#!/usr/bin/env python3
"""
Script to map browser_gym task_ids to VWA task_ids by comparing specific fields.
Maps tasks based on: intent, eval, reasoning_difficulty, visual_difficulty, overall_difficulty, comments
"""

import csv
import hashlib
import json
from typing import Any, Dict, List, Optional, Tuple


def normalize_field(field: Any) -> str:
    """Normalize a field for comparison by converting to string and removing whitespace"""
    if field is None:
        return ""
    return str(field).strip()


def create_task_signature(task: Dict[str, Any]) -> str:
    """Create a unique signature for a task based on the comparison fields"""
    fields_to_compare = [
        "intent",
        "eval",
        "reasoning_difficulty",
        "visual_difficulty",
        "overall_difficulty",
        "comments",
    ]

    signature_parts = []
    for field in fields_to_compare:
        value = normalize_field(task.get(field, ""))
        signature_parts.append(f"{field}:{value}")

    # Create hash of the signature for easier comparison
    signature_str = "|".join(signature_parts)
    return hashlib.md5(signature_str.encode()).hexdigest()


def load_tasks(file_path: str) -> List[Dict[str, Any]]:
    """Load tasks from a JSON file"""
    with open(file_path, "r", encoding="utf-8") as f:
        return json.load(f)


def find_matching_tasks(
    browser_gym_tasks: List[Dict[str, Any]], vwa_tasks_by_domain: Dict[str, List[Dict[str, Any]]]
) -> List[Tuple[int, int, str, str]]:
    """Find matching tasks between browser_gym and VWA based on signature comparison"""

    # Create signatures for VWA tasks with domain info
    vwa_signatures = {}
    for domain, tasks in vwa_tasks_by_domain.items():
        for task in tasks:
            signature = create_task_signature(task)
            task_id = task.get("task_id")
            sites = task.get("sites", [])
            if signature not in vwa_signatures:
                vwa_signatures[signature] = []
            vwa_signatures[signature].append((task_id, sites, domain))

    matches = []
    unmatched_browser_gym = []

    # Find matches for browser_gym tasks
    for bg_task in browser_gym_tasks:
        bg_signature = create_task_signature(bg_task)
        bg_task_id = bg_task.get("task_id")
        bg_sites = bg_task.get("sites", [])

        if bg_signature in vwa_signatures:
            # Found matching signature(s)
            for vwa_task_id, vwa_sites, domain in vwa_signatures[bg_signature]:
                # Also check if sites match
                if set(bg_sites) == set(vwa_sites):
                    site_str = ",".join(sorted(bg_sites))
                    matches.append((bg_task_id, vwa_task_id, site_str, domain))
                    break
        else:
            unmatched_browser_gym.append(bg_task_id)

    print(f"Found {len(matches)} matches")
    print(f"Unmatched browser_gym tasks: {len(unmatched_browser_gym)}")
    if unmatched_browser_gym:
        print(f"First 10 unmatched browser_gym task_ids: {unmatched_browser_gym[:10]}")

    return matches


def create_detailed_comparison(
    browser_gym_tasks: List[Dict[str, Any]], vwa_tasks_by_domain: Dict[str, List[Dict[str, Any]]]
) -> List[Dict[str, Any]]:
    """Create detailed comparison for manual verification"""

    comparison_results = []
    fields_to_compare = [
        "intent",
        "eval",
        "reasoning_difficulty",
        "visual_difficulty",
        "overall_difficulty",
        "comments",
    ]

    for bg_task in browser_gym_tasks:
        bg_task_id = bg_task.get("task_id")
        bg_sites = bg_task.get("sites", [])

        # Find potential matches based on individual field comparisons
        potential_matches = []

        for domain, vwa_tasks in vwa_tasks_by_domain.items():
            for vwa_task in vwa_tasks:
                vwa_task_id = vwa_task.get("task_id")
                vwa_sites = vwa_task.get("sites", [])

            # Check if sites match first
            if set(bg_sites) != set(vwa_sites):
                continue

            # Count matching fields
            matching_fields = 0
            field_details = {}

            for field in fields_to_compare:
                bg_value = normalize_field(bg_task.get(field))
                vwa_value = normalize_field(vwa_task.get(field))

                is_match = bg_value == vwa_value
                if is_match:
                    matching_fields += 1

                field_details[field] = {
                    "match": is_match,
                    "bg_value": bg_value[:100] + "..." if len(bg_value) > 100 else bg_value,
                    "vwa_value": vwa_value[:100] + "..." if len(vwa_value) > 100 else vwa_value,
                }

                if matching_fields >= 4:  # Require at least 4 out of 6 fields to match
                    potential_matches.append(
                        {
                            "vwa_task_id": vwa_task_id,
                            "domain": domain,
                            "matching_fields": matching_fields,
                            "field_details": field_details,
                        }
                    )

        # Sort by number of matching fields
        potential_matches.sort(key=lambda x: x["matching_fields"], reverse=True)

        comparison_results.append(
            {
                "browser_gym_task_id": bg_task_id,
                "sites": ",".join(sorted(bg_sites)),
                "potential_matches": potential_matches[:3],  # Top 3 matches
            }
        )

    return comparison_results


def save_mapping_csv(matches: List[Tuple[int, int, str, str]], output_file: str):
    """Save the mapping to a CSV file"""
    with open(output_file, "w", newline="", encoding="utf-8") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["browser_gym_task_id", "vwa_task_id", "sites", "domain"])

        for bg_task_id, vwa_task_id, sites, domain in sorted(matches):
            writer.writerow([bg_task_id, vwa_task_id, sites, domain])

    print(f"Mapping saved to {output_file}")


def save_detailed_comparison(comparison: List[Dict[str, Any]], output_file: str):
    """Save detailed comparison for manual review"""
    with open(output_file, "w", newline="", encoding="utf-8") as csvfile:
        fieldnames = [
            "browser_gym_task_id",
            "sites",
            "best_vwa_match_id",
            "domain",
            "matching_fields_count",
            "intent_match",
            "eval_match",
            "reasoning_difficulty_match",
            "visual_difficulty_match",
            "overall_difficulty_match",
            "comments_match",
        ]
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

        for item in comparison:
            bg_task_id = item["browser_gym_task_id"]
            sites = item["sites"]

            if item["potential_matches"]:
                best_match = item["potential_matches"][0]
                row = {
                    "browser_gym_task_id": bg_task_id,
                    "sites": sites,
                    "best_vwa_match_id": best_match["vwa_task_id"],
                    "matching_fields_count": best_match["matching_fields"],
                    "domain": best_match.get("domain", ""),
                }

                # Add individual field match status
                for field in [
                    "intent",
                    "eval",
                    "reasoning_difficulty",
                    "visual_difficulty",
                    "overall_difficulty",
                    "comments",
                ]:
                    row[f"{field}_match"] = best_match["field_details"][field]["match"]

                writer.writerow(row)
            else:
                writer.writerow(
                    {
                        "browser_gym_task_id": bg_task_id,
                        "sites": sites,
                        "best_vwa_match_id": "NO_MATCH",
                        "domain": "",
                        "matching_fields_count": 0,
                    }
                )

    print(f"Detailed comparison saved to {output_file}")


def main():
    # File paths
    browser_gym_file = "config_files/vwa_browser_gym/test_raw.json"
    vwa_files = {
        "classifieds": "config_files/vwa/test_classifieds.raw.json",
        "shopping": "config_files/vwa/test_shopping.raw.json",
        "reddit": "config_files/vwa/test_reddit.raw.json",
    }

    # Load browser_gym tasks
    print("Loading browser_gym tasks...")
    browser_gym_tasks = load_tasks(browser_gym_file)
    print(f"Loaded {len(browser_gym_tasks)} browser_gym tasks")

    # Load all VWA tasks by domain
    print("Loading VWA tasks...")
    vwa_tasks_by_domain = {}
    total_vwa_tasks = 0
    for site, file_path in vwa_files.items():
        tasks = load_tasks(file_path)
        print(f"Loaded {len(tasks)} {site} tasks")
        vwa_tasks_by_domain[site] = tasks
        total_vwa_tasks += len(tasks)

    print(f"Total VWA tasks: {total_vwa_tasks}")

    # Find exact matches
    print("\nFinding exact matches...")
    matches = find_matching_tasks(browser_gym_tasks, vwa_tasks_by_domain)

    # Create flat list for detailed comparison
    all_vwa_tasks = []
    for tasks in vwa_tasks_by_domain.values():
        all_vwa_tasks.extend(tasks)

    # Create detailed comparison for all tasks
    print("\nCreating detailed comparison...")
    detailed_comparison = create_detailed_comparison(browser_gym_tasks, vwa_tasks_by_domain)

    # Save results
    save_mapping_csv(matches, "browser_gym_to_vwa_mapping.csv")
    save_detailed_comparison(detailed_comparison, "browser_gym_vwa_detailed_comparison.csv")

    print("\nSummary:")
    print(f"- Total browser_gym tasks: {len(browser_gym_tasks)}")
    print(f"- Total VWA tasks: {len(all_vwa_tasks)}")
    print(f"- Exact matches found: {len(matches)}")
    print(f"- Match rate: {len(matches) / len(browser_gym_tasks) * 100:.1f}%")


if __name__ == "__main__":
    main()
