#!/usr/bin/env python3
"""End-to-end test: run solve.sh logic + evaluate.py for selected tasks.

Uses git worktrees to simulate Docker environment locally.
Requires: codeflash installed, inference repo available.

Usage:
    python rl_env/e2e_test.py --tasks 3
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

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

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


def get_python() -> str:
    """Get the codeflash venv python."""
    venv_python = CODEFLASH_ROOT / ".venv" / "bin" / "python"
    if venv_python.exists():
        return str(venv_python)
    return sys.executable


def create_worktree(commit: str, dest: Path) -> bool:
    """Create a git worktree at the given commit."""
    # Clean up if exists
    subprocess.run(
        ["git", "worktree", "remove", "--force", str(dest)],
        cwd=INFERENCE_REPO, capture_output=True,
    )
    result = subprocess.run(
        ["git", "worktree", "add", str(dest), commit],
        cwd=INFERENCE_REPO, capture_output=True, text=True,
    )
    return result.returncode == 0


def cleanup_worktree(dest: Path) -> None:
    subprocess.run(
        ["git", "worktree", "remove", "--force", str(dest)],
        cwd=INFERENCE_REPO, capture_output=True,
    )


def run_solve(task_dir: Path, workspace: Path) -> tuple[bool, str]:
    """Run solve.sh logic (apply optimization) against the workspace."""
    solve_sh = task_dir / "solution" / "solve.sh"
    if not solve_sh.exists():
        return False, "solve.sh not found"

    # Check if git_parent (has patch) or db method (has apply_optimization.py)
    patch_file = task_dir / "solution" / "optimization.patch"
    apply_script = task_dir / "solution" / "apply_optimization.py"
    config_file = task_dir / "solution" / "config.json"

    if patch_file.exists():
        # Git apply
        result = subprocess.run(
            ["git", "apply", str(patch_file)],
            cwd=workspace, capture_output=True, text=True,
        )
        if result.returncode != 0:
            return False, f"git apply failed: {result.stderr[:300]}"
        return True, "patch applied"

    elif apply_script.exists() and config_file.exists():
        # Use apply_optimization.py with adapted paths
        script_content = apply_script.read_text(encoding="utf-8")
        script_content = script_content.replace("/workspace/inference", str(workspace))
        script_content = script_content.replace("/solution/config.json", str(config_file))
        script_content = script_content.replace(
            "/solution/optimized_code.py",
            str(task_dir / "solution" / "optimized_code.py"),
        )

        adapted_script = Path(tempfile.mktemp(suffix=".py"))
        adapted_script.write_text(script_content, encoding="utf-8")

        result = subprocess.run(
            [get_python(), str(adapted_script)],
            capture_output=True, text=True, encoding="utf-8", timeout=30,
            env={**os.environ, "PYTHONPATH": str(workspace)},
        )
        adapted_script.unlink(missing_ok=True)

        if result.returncode != 0:
            return False, f"apply failed: {(result.stdout + result.stderr)[-300:]}"
        return True, result.stdout.strip()

    return False, "no solve method found"


def run_evaluate(task_dir: Path, workspace: Path, entry: dict) -> tuple[bool, dict | None, str]:
    """Run evaluate.py against the workspace."""
    test_dir = task_dir / "tests" / "codeflash_eval"
    evaluate_script = test_dir / "evaluate.py"
    short_name = entry["function_name"].rsplit(".", maxsplit=1)[-1]

    if not evaluate_script.exists():
        return False, None, "evaluate.py not found"

    with tempfile.TemporaryDirectory(prefix="eval_output_") as output_dir:
        # Adapt evaluate.py paths (it expects /workspace/inference and /tests/codeflash_eval)
        script_content = evaluate_script.read_text(encoding="utf-8")

        result = subprocess.run(
            [
                get_python(), str(evaluate_script),
                "--function", short_name,
                "--file-path", entry["file_path"],
                "--repo-path", str(workspace),
                "--test-dir", str(test_dir),
                "--output-dir", output_dir,
            ],
            capture_output=True, text=True, encoding="utf-8", timeout=600,
            env={**os.environ, "PYTHONPATH": str(workspace)},
        )

        reward_path = Path(output_dir) / "reward.json"
        if reward_path.exists():
            reward = json.loads(reward_path.read_text(encoding="utf-8"))
        else:
            reward = None

        output = result.stdout[-1000:] if result.stdout else ""
        if result.stderr:
            output += "\nSTDERR: " + result.stderr[-500:]

        return result.returncode == 0, reward, output


def select_test_tasks(catalog: list[dict], n: int) -> list[dict]:
    """Select diverse tasks for testing: mix of git_parent and db_code_match."""
    git_parent = [e for e in catalog if e["resolution_method"] == "git_parent"]
    db_match = [e for e in catalog if e["resolution_method"] == "db_code_match"]

    selected = []
    # Pick 1 git_parent if available
    if git_parent and n > 0:
        selected.append(git_parent[0])
        n -= 1
    # Fill with db_code_match (prefer those with high speedup)
    db_match.sort(key=lambda e: e.get("speedup_x") or 0, reverse=True)
    for e in db_match:
        if n <= 0:
            break
        # Skip if task dir doesn't exist or has no behavioral tests
        task_name = find_task_dir(e)
        if task_name:
            test_dir = TASKS_DIR / task_name / "tests" / "codeflash_eval"
            behavioral = list(test_dir.glob("*behaviorinstrumented*.py"))
            if behavioral:
                selected.append(e)
                n -= 1

    return selected


def find_task_dir(entry: dict) -> str | None:
    """Find the task directory name for an entry."""
    import re
    func_name = entry["function_name"]
    file_path = entry["file_path"]
    grandparent = Path(file_path).parent.name
    name = f"{grandparent}-{func_name}".replace(".", "-")
    name = re.sub(r"[^a-zA-Z0-9_-]", "-", name)
    name = re.sub(r"-+", "-", name).strip("-").lower()[:80]
    if (TASKS_DIR / name).exists():
        return name
    return None


def main() -> None:
    parser = argparse.ArgumentParser(description="End-to-end test of Harbor tasks")
    parser.add_argument("--tasks", type=int, default=3, help="Number of tasks to test")
    parser.add_argument("--task-name", type=str, help="Test a specific task by directory name")
    args = parser.parse_args()

    catalog = json.loads(CATALOG_PATH.read_text(encoding="utf-8"))

    if args.task_name:
        entries = [e for e in catalog if find_task_dir(e) == args.task_name]
        if not entries:
            log.error(f"Task '{args.task_name}' not found in catalog")
            return
        test_entries = entries
    else:
        test_entries = select_test_tasks(catalog, args.tasks)

    log.info(f"Testing {len(test_entries)} tasks end-to-end")
    log.info(f"Python: {get_python()}")
    log.info("")

    results = {"solve_pass": 0, "solve_fail": 0, "eval_pass": 0, "eval_fail": 0, "details": []}

    for i, entry in enumerate(test_entries):
        func = entry["function_name"]
        commit = entry["pre_optimization_commit"]
        method = entry["resolution_method"]
        task_name = find_task_dir(entry)

        if not task_name:
            log.warning(f"  [{i+1}] No task dir for {func}")
            continue

        task_dir = TASKS_DIR / task_name
        log.info(f"{'='*60}")
        log.info(f"[{i+1}/{len(test_entries)}] {func}")
        log.info(f"  File: {entry['file_path']}")
        log.info(f"  Method: {method}, Commit: {commit[:12]}")
        log.info(f"  Task: {task_name}")

        # Create worktree
        workspace = Path(f"/tmp/e2e_test_{i}")
        if not create_worktree(commit, workspace):
            log.error(f"  Failed to create worktree at {commit[:12]}")
            results["solve_fail"] += 1
            continue

        try:
            # Step 1: Run solve (apply optimization)
            log.info("  SOLVE...")
            solve_ok, solve_msg = run_solve(task_dir, workspace)
            if solve_ok:
                results["solve_pass"] += 1
                log.info(f"  SOLVE: PASS ({solve_msg})")
            else:
                results["solve_fail"] += 1
                log.error(f"  SOLVE: FAIL ({solve_msg})")
                results["details"].append({"task": task_name, "stage": "solve", "error": solve_msg})
                continue

            # Step 2: Run evaluate
            log.info("  EVALUATE...")
            eval_ok, reward, eval_output = run_evaluate(task_dir, workspace, entry)

            if reward:
                log.info(f"  EVALUATE: correct={reward.get('correct')} speedup={reward.get('speedup')} "
                        f"passed={reward.get('tests_passed')}/{reward.get('tests_total')}")
                if reward.get("correct") == 1.0:
                    results["eval_pass"] += 1
                else:
                    results["eval_fail"] += 1
                    results["details"].append({
                        "task": task_name, "stage": "eval", "reward": reward,
                        "output": eval_output[-300:]
                    })
            else:
                results["eval_fail"] += 1
                log.error(f"  EVALUATE: FAIL (no reward.json)")
                log.error(f"  Output: {eval_output[-500:]}")
                results["details"].append({
                    "task": task_name, "stage": "eval", "error": "no reward.json",
                    "output": eval_output[-500:]
                })

        finally:
            cleanup_worktree(workspace)

    log.info("")
    log.info("=" * 60)
    log.info("E2E TEST RESULTS")
    log.info("=" * 60)
    log.info(f"  Solve: {results['solve_pass']} pass, {results['solve_fail']} fail")
    log.info(f"  Evaluate: {results['eval_pass']} pass, {results['eval_fail']} fail")

    if results["details"]:
        log.info("\nFailure details:")
        for d in results["details"]:
            log.info(f"  {d['task']} [{d['stage']}]: {d.get('error', d.get('reward', ''))}")


if __name__ == "__main__":
    main()
