import os
import shutil
import re

# Update these paths
source_root = "./raw_gpt/direct_api"  # <- where your current files are
target_root = "./results/ICL/ChatGPT"         # <- where you want the new folder structure

# Patterns
prompt_pattern = re.compile(r'_(io|zsr)_prompt')
encoding_pattern = re.compile(r'(many2one|one2many|one2one)')
k_pattern = re.compile(r'_k=(\d+)')
lang_pattern = re.compile(r'^([a-zA-Z0-9\-]+)_gpt')

# Encoding name mapping (excluding one2many which is special)
encoding_map = {
    "many2one": "many_to_one",
    "one2one": "one_to_one"
}

for folder in os.listdir(source_root):
    folder_path = os.path.join(source_root, folder)
    if not os.path.isdir(folder_path):
        continue

    # Extract components from folder name
    prompt_match = prompt_pattern.search(folder)
    encoding_match = encoding_pattern.search(folder)
    k_match = k_pattern.search(folder)

    if not (prompt_match and encoding_match):
        continue

    prompt_type = prompt_match.group(1) + "_prompt"
    encoding = encoding_match.group(1)

    # Determine proper encoding folder name
    if encoding == "one2many" and k_match:
        encoding_folder = f"one_to_many_{k_match.group(1)}"
    else:
        encoding_folder = encoding_map.get(encoding, encoding)

    for filename in os.listdir(folder_path):
        if not filename.endswith('.json'):
            continue

        lang_match = lang_pattern.match(filename)
        if not lang_match:
            continue

        language = lang_match.group(1)
        src_path = os.path.join(folder_path, filename)

        # Construct destination directory
        dst_dir = os.path.join(target_root, prompt_type, encoding_folder, language)
        os.makedirs(dst_dir, exist_ok=True)

        dst_path = os.path.join(dst_dir, filename)
        shutil.move(src_path, dst_path)
        print(f"Moved: {src_path} → {dst_path}")
