from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
import glob
import json
import os
from pathlib import Path
import shutil
import sys
from pprint import pprint

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
from experiments_utils.constants import ALL_EXPER_DIR_STRUCT
from utils.file_utils import (
    find_files,
    is_finished_execution,
    remove_empty_dirs,
    get_common_paths,
    copy_move,
)


source_folder = input("Enter a dir with experiment runs to move.")
# source_folder = "results/google/gemini-2.0-flash-001/____EXEC____mem=text-a_t=1____CRIT____m=2p_eval=tri"


# find all args.json files
args_files = find_files(source_folder, "args.json", upwards=False, downwards=True)

args_files = [f for f in args_files if is_finished_execution(Path(f).parent)]
longest_p, shortest_p, common_p = get_common_paths(args_files, relative_to=source_folder)

config_directories = []

if len(common_p) == 0:
    config_directories = sorted([os.path.commonpath(args_files)])
elif len(longest_p) == len(shortest_p):
    config_directories = sorted([os.path.commonpath(shortest_p)])
else:
    config_directories = sorted(shortest_p)

if not (len(config_directories) == 1 and str(config_directories[0]) == source_folder):
    while True:
        try:
            print("Below is the list of experiment config dirs to move\n")
            for i, config_dir in enumerate(config_directories):
                print(f"{i}: {config_dir}")
            ans = input(
                "Enter the index of the config dir to remove, OR press enter to continue, OR press 'q' to quit.\n"
            )
            if ans == "q":
                exit()
            elif ans == "":
                break
            elif ans.isdigit() and 0 <= int(ans) < len(config_directories):
                config_directories.pop(int(ans))
            else:
                print("Invalid input. Please enter a valid index.")

        except ValueError:
            print("Invalid input. Please enter a valid index.")

copy_move_mode = input("\nCopy or move directories? ('c' | 'm'): ").strip().lower()

copy_move_mode = "copy" if copy_move_mode == "c" else "move"

for config_dir in config_directories:
    print(f"\nMoving {config_dir}...")
    dir_name = os.path.basename(config_dir)
    config_name = input(
        f"\nEnter the experiment config name.\nOr press enter to use the name below:\n {dir_name}"
    ).strip()
    if config_name == "":
        config_name = dir_name

    args_files = find_files(config_dir, "args.json", upwards=False, downwards=True)
    result_dirs = [os.path.dirname(args_file) for args_file in args_files]

    src_to_dst_map = {}
    for args_file, result_dir in zip(args_files, result_dirs):
        # Load args.json
        args = json.load(open(args_file))

        # Get domain
        domain = args["test_config_base_dir"].split("/")[-1].split("_")[-1]

        # Get model
        model = args["agents_configs"]["executor_agent"]["lm_config"]["model"]

        # Create folder structure in target folder
        target_folder = ALL_EXPER_DIR_STRUCT.format(model=model, config_name=config_name, domain=domain)

        src_to_dst_map[result_dir] = os.path.join(target_folder, os.path.basename(result_dir))

    print("\nDirectories will be move from -> to:")
    for result_dir in src_to_dst_map:
        print(f"{result_dir} -> {src_to_dst_map[result_dir]}.")

    if input("\nPress enter to continue or q to skip this config.").strip().lower() == "q":
        continue

    for target_folder in src_to_dst_map.values():
        if not os.path.exists(target_folder):
            abs_path = os.path.abspath(target_folder)
            print(f"Creating folder {abs_path}")
            os.makedirs(target_folder, exist_ok=True)

    # Move or Copy result_dir as a subdir of target_folder.
    # E.g.: "results/openai/gpt-4o-2024-08-06/exper_001/execution1" -> "experiments/gpt-4o/exper_001/shopping/results/execution1"
    with ThreadPoolExecutor(max_workers=8) as executor:
        for src, dst in src_to_dst_map.items():
            executor.submit(
                copy_move,
                src,
                dst,
                mode=copy_move_mode,
                merge_dirs=True,
            )

    # Sequential copy/move
    # for src, dst in src_to_dst_map.items():
    #     copy_move(src, dst, mode=copy_move_mode)

    # Delete empty directories
# remove_empty_dirs(os.path.dirname(source_folder))
