import os
from pathlib import Path
import sys

from experiments_utils.scripts.consolidate_experiments import get_score_from_csv

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
from utils.file_utils import add_annotation_to_path, copy_move, find_files, get_task_id_from_file_path

base_dir = "/home/mashalimay/webarena/experiments/gpt-4o-2024-08-06/base_stateaware"
domains = ["reddit", "shopping", "classifieds"]

for domain in domains:
    htmls_dir = f"{base_dir}/{domain}/lm_data/critique_agent/conversation"
    csv_with_scores = f"{base_dir}/{domain}/consolidated_summary_data.csv"

    trace_dirs = [
        f"{base_dir}/{domain}/lm_data/critique_agent/conversation",
        f"{base_dir}/{domain}/lm_data/critique_agent/usage",
        f"{base_dir}/{domain}/lm_data/executor_agent/conversation",
        f"{base_dir}/{domain}/lm_data/executor_agent/usage",
        f"{base_dir}/{domain}/htmls",
    ]

    for trace_dir in trace_dirs:
        trace_dir = Path(trace_dir)
        if not trace_dir.exists():
            continue

        target_dir = trace_dir

        for file in trace_dir.iterdir():
            if not file.is_file():
                continue

            task_id = get_task_id_from_file_path(file)
            if not task_id:
                continue

            dest_path = Path(target_dir) / file.name
            score = get_score_from_csv(csv_with_scores, task_id)
            dest_path = Path(
                add_annotation_to_path(dest_path, f"s={score:.0f}/" if score is not None else "s=NA/", add_to_end=False)
            )
            os.makedirs(dest_path.parent, exist_ok=True)

            if file != dest_path:
                copy_move(file, dest_path, mode="move")
