#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import shutil

# === CONFIG ===
OUTPUTS_DIR = "/home/hyrenko/Diploma/outputs"   # откуда брать
DEST_DIR = "/home/hyrenko/Diploma/response"     # базовая папка назначения

# создаём папки, если их нет
os.makedirs(DEST_DIR, exist_ok=True)
os.makedirs(os.path.join(DEST_DIR, "gemma"), exist_ok=True)
os.makedirs(os.path.join(DEST_DIR, "llama"), exist_ok=True)
os.makedirs(os.path.join(DEST_DIR, "qwen"), exist_ok=True)

# === PICK MODEL INTERACTIVELY ===
def pick_model():
    print("\n=== Select model ===")
    print("[1] gemma")
    print("[2] llama")
    print("[3] qwen")
    while True:
        ch = input("Choose model [1-3]: ").strip()
        if ch == "1":
            return "gemma"
        elif ch == "2":
            return "llama"
        elif ch == "3":
            return "qwen"
        else:
            print("Invalid selection, try again.")

# === MAIN ===
def main():
    model_name = pick_model()
    print(f"\n[INFO] Selected model: {model_name}")

    # Папка назначения — /response/<model_name>/
    dest_path = os.path.join(DEST_DIR, model_name)
    os.makedirs(dest_path, exist_ok=True)

    # Сканируем все подпапки в outputs
    subdirs = sorted([os.path.join(OUTPUTS_DIR, d) for d in os.listdir(OUTPUTS_DIR)
                      if os.path.isdir(os.path.join(OUTPUTS_DIR, d))])

    counter = 1
    copied = 0

    for subdir in subdirs:
        # фильтруем только те, где в названии есть имя модели
        if model_name not in subdir:
            continue

        src_file = os.path.join(subdir, "responses.json")
        if os.path.isfile(src_file):
            dst_file = os.path.join(dest_path, f"{counter}.json")
            shutil.copy2(src_file, dst_file)
            print(f"[{counter}] Copied from {subdir} -> {dst_file}")
            counter += 1
            copied += 1
        else:
            print(f"[SKIP] No response.json in {subdir}")

    print("\n=== Summary ===")
    print(f"Model: {model_name}")
    print(f"Copied files: {copied}")
    print(f"Destination: {dest_path}")

    if copied == 0:
        print("[WARN] No matching response.json files were found.")

if __name__ == "__main__":
    main()