"""Validate all Harbor tasks and exclude broken ones.

Checks:
1. solve.sh can apply the optimization without errors
2. Task has at least one behavioral test file
3. Task has at least one perf test file
4. original.py is valid Python

Moves invalid tasks to a _broken/ directory.

Usage:
    python rl_env/validate_tasks.py
"""

from __future__ import annotations

import ast
import json
import logging
import re
import shutil
import subprocess
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
log = logging.getLogger(__name__)

TASKS_DIR = Path(__file__).parent / "codeflash-perf-tasks"
INFERENCE_REPO = Path("/Users/saurabh/Dropbox/hacks/inference")
CATALOG_PATH = Path(__file__).parent / "resolved_catalog_full.json"


def sanitize_task_name(function_name: str, file_path: str) -> str:
    grandparent = Path(file_path).parent.name
    name = f"{grandparent}-{function_name}".replace(".", "-")
    name = re.sub(r"[^a-zA-Z0-9_-]", "-", name)
    name = re.sub(r"-+", "-", name).strip("-").lower()
    return name[:80]


def validate_task(task_dir: Path, entry: dict, repo: Path) -> tuple[bool, str]:
    """Validate a task. Returns (valid, reason)."""
    test_dir = task_dir / "tests" / "codeflash_eval"

    # Check 1: Has behavioral tests
    behavioral = list(test_dir.glob("*behaviorinstrumented*.py"))
    if not behavioral:
        return False, "no behavioral tests"

    # Check 2: Has perf tests
    perf = list(test_dir.glob("*perfinstrumented*.py"))
    if not perf:
        return False, "no perf tests"

    # Check 3: original.py is valid Python
    original_py = test_dir / "original.py"
    if not original_py.exists():
        return False, "no original.py"
    try:
        ast.parse(original_py.read_text(encoding="utf-8"))
    except SyntaxError:
        return False, "original.py has syntax error"

    # Check 4: solve.sh can apply optimization
    commit = entry.get("pre_optimization_commit", "HEAD")
    if commit == "HEAD":
        return True, "ok (HEAD commit, skip solve check)"

    patch_file = task_dir / "solution" / "optimization.patch"
    apply_script = task_dir / "solution" / "apply_optimization.py"

    if patch_file.exists():
        # Test git apply in a temp worktree
        wt = Path(f"/tmp/validate_{task_dir.name}")
        subprocess.run(["git", "worktree", "remove", "--force", str(wt)], cwd=repo, capture_output=True)
        result = subprocess.run(
            ["git", "worktree", "add", str(wt), commit],
            cwd=repo, capture_output=True, text=True,
        )
        if result.returncode != 0:
            return False, f"cannot create worktree at {commit[:12]}"

        result = subprocess.run(
            ["git", "apply", "--check", str(patch_file)],
            cwd=wt, capture_output=True, text=True,
        )
        subprocess.run(["git", "worktree", "remove", "--force", str(wt)], cwd=repo, capture_output=True)

        if result.returncode != 0:
            return False, "git patch doesn't apply"

    elif apply_script.exists():
        config_file = task_dir / "solution" / "config.json"
        if not config_file.exists():
            return False, "no config.json for apply script"

        config = json.loads(config_file.read_text(encoding="utf-8"))
        func_name = config.get("function_name", "")
        target_file = config.get("target_file", "")

        # Check function exists at the commit
        result = subprocess.run(
            ["git", "show", f"{commit}:{target_file}"],
            cwd=repo, capture_output=True, text=True, encoding="utf-8",
        )
        if result.returncode != 0:
            return False, f"target file not in git at {commit[:12]}"

        # Check function exists in the file at that commit
        try:
            tree = ast.parse(result.stdout)
            found = any(
                isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name
                for node in ast.walk(tree)
            )
            if not found:
                return False, f"function '{func_name}' not found at {commit[:12]}"
        except SyntaxError:
            return False, f"target file has syntax error at {commit[:12]}"

    return True, "ok"


def main() -> None:
    catalog = json.loads(CATALOG_PATH.read_text(encoding="utf-8"))
    catalog_by_trace = {e["trace_id"]: e for e in catalog}

    broken_dir = TASKS_DIR.parent / "codeflash-perf-tasks-broken"
    broken_dir.mkdir(exist_ok=True)

    valid_count = 0
    broken_count = 0
    reasons: dict[str, int] = {}

    task_dirs = sorted([d for d in TASKS_DIR.iterdir() if d.is_dir() and d.name != "base-image"])
    log.info(f"Validating {len(task_dirs)} tasks...")

    for i, task_dir in enumerate(task_dirs):
        if (i + 1) % 30 == 0:
            log.info(f"  [{i + 1}/{len(task_dirs)}]...")

        # Find catalog entry
        entry = None
        for e in catalog:
            if sanitize_task_name(e["function_name"], e["file_path"]) == task_dir.name:
                entry = e
                break

        if not entry:
            # Try to find by partial match
            for e in catalog:
                tn = sanitize_task_name(e["function_name"], e["file_path"])
                if tn == task_dir.name:
                    entry = e
                    break

        if not entry:
            broken_count += 1
            reason = "no catalog entry"
            reasons[reason] = reasons.get(reason, 0) + 1
            shutil.move(str(task_dir), str(broken_dir / task_dir.name))
            continue

        valid, reason = validate_task(task_dir, entry, INFERENCE_REPO)

        if valid:
            valid_count += 1
        else:
            broken_count += 1
            reasons[reason] = reasons.get(reason, 0) + 1
            log.debug(f"  BROKEN: {task_dir.name} — {reason}")
            shutil.move(str(task_dir), str(broken_dir / task_dir.name))

    log.info("=" * 60)
    log.info(f"VALIDATION RESULTS:")
    log.info(f"  Valid: {valid_count}")
    log.info(f"  Broken: {broken_count}")
    if reasons:
        log.info(f"  Reasons:")
        for reason, count in sorted(reasons.items(), key=lambda x: -x[1]):
            log.info(f"    {reason}: {count}")


if __name__ == "__main__":
    main()
