"""Validate all Harbor tasks end-to-end in Docker and generate a markdown report.

Runs solve.sh + evaluate.py for each task, captures reward.json and debug info.

Usage:
    python rl_env/batch_validate.py --limit 5   # test first 5
    python rl_env/batch_validate.py              # all tasks
"""

from __future__ import annotations

import argparse
import json
import logging
import re
import subprocess
import time
from dataclasses import dataclass, field
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").resolve()
DOCKER_IMAGE = "codeflash-inference-base:latest"


@dataclass
class TaskResult:
    task_name: str
    function_name: str
    file_path: str
    commit: str
    resolution_method: str
    speedup_db: float | None  # from catalog
    solve_ok: bool = False
    solve_output: str = ""
    eval_ok: bool = False
    reward: dict = field(default_factory=dict)
    eval_output: str = ""
    duration_secs: float = 0
    error: str = ""


def get_task_commit(task_dir: Path) -> str:
    dockerfile = task_dir / "environment" / "Dockerfile"
    if dockerfile.exists():
        content = dockerfile.read_text(encoding="utf-8")
        match = re.search(r"git checkout (\S+)", content)
        if match:
            return match.group(1)[:12]
    return "HEAD"


def run_task(task_dir: Path, entry: dict) -> TaskResult:
    task_name = task_dir.name
    func_name = entry.get("function_name", "")
    file_path = entry.get("file_path", "")
    commit = entry.get("pre_optimization_commit", "HEAD")

    # Read function/file from test.sh if available
    test_sh = task_dir / "tests" / "test.sh"
    if test_sh.exists():
        content = test_sh.read_text(encoding="utf-8")
        m = re.search(r'--function "([^"]*)"', content)
        if m:
            func_name = m.group(1)
        m = re.search(r'--file-path "([^"]*)"', content)
        if m:
            file_path = m.group(1)

    result = TaskResult(
        task_name=task_name,
        function_name=func_name,
        file_path=file_path,
        commit=commit,
        resolution_method=entry.get("resolution_method", ""),
        speedup_db=entry.get("speedup_x"),
    )

    start = time.time()
    try:
        # Create temp dir for output (mounted as volume)
        import tempfile

        with tempfile.TemporaryDirectory(prefix="cf_eval_") as output_dir:
            output_path = Path(output_dir)

            docker_result = subprocess.run(
                [
                    "docker",
                    "run",
                    "--rm",
                    "-v",
                    f"{task_dir}/solution:/solution:ro",
                    "-v",
                    f"{task_dir}/tests:/tests:ro",
                    "-v",
                    f"{output_path}:/logs/verifier",
                    DOCKER_IMAGE,
                    "bash",
                    "-c",
                    f"""
cd /workspace/inference

# --- SETUP: Get to pre-optimization state ---
if git cat-file -t {commit} >/dev/null 2>&1; then
    git checkout -f {commit} 2>&1
else
    echo "Commit {commit} not found, reconstructing from original.py"
    cp /tests/codeflash_eval/original.py /workspace/inference/{file_path}
    git add -A && git commit -m "Pre-optimization state (reconstructed)" 2>&1
fi

# --- SOLVE ---
if [ -f /solution/optimization.patch ]; then
    git apply /solution/optimization.patch 2>&1
else
    python /solution/apply_optimization.py 2>&1
fi
git add -A && git commit --allow-empty -m "Apply optimization" 2>&1

# --- EVALUATE ---
python /tests/codeflash_eval/evaluate.py \
    --function "{func_name}" \
    --file-path "{file_path}" \
    --repo-path /workspace/inference \
    --test-dir /tests/codeflash_eval \
    --output-dir /logs/verifier \
    --debug 2>&1
""",
                ],
                capture_output=True,
                text=True,
                encoding="utf-8",
                timeout=600,
            )

            result.eval_output = docker_result.stdout[-3000:] if docker_result.stdout else ""
            result.solve_ok = (
                "Syntax check passed" in result.eval_output
                or "Apply optimization" in result.eval_output
                or "patch applied" in result.eval_output.lower()
            )

            # Read reward.json directly from mounted volume
            reward_file = output_path / "reward.json"
            if reward_file.exists():
                result.reward = json.loads(reward_file.read_text(encoding="utf-8"))
                result.eval_ok = result.reward.get("correct") == 1.0

            if docker_result.returncode != 0:
                result.error = f"Docker exit code {docker_result.returncode}"
                if docker_result.stderr:
                    result.error += f": {docker_result.stderr[-200:]}"

    except subprocess.TimeoutExpired:
        result.error = "Timeout (600s)"
    except Exception as e:
        result.error = str(e)

    result.duration_secs = time.time() - start
    return result


def generate_report(results: list[TaskResult], output_path: Path) -> None:
    lines = ["# Codeflash RL Environment — Batch Validation Report\n"]

    # Summary
    total = len(results)
    solve_pass = sum(1 for r in results if r.solve_ok)
    eval_pass = sum(1 for r in results if r.eval_ok)
    faster = sum(1 for r in results if r.reward.get("speedup", 0) > 1.0)
    all_tests_pass = sum(
        1
        for r in results
        if r.reward.get("tests_passed", 0) == r.reward.get("tests_total", -1) and r.reward.get("tests_total", 0) > 0
    )

    lines.append("## Summary\n")
    lines.append("| Metric | Count | % |")
    lines.append("|--------|-------|---|")
    lines.append(f"| Total tasks | {total} | 100% |")
    lines.append(f"| Solve passes | {solve_pass} | {solve_pass * 100 // total}% |")
    lines.append(f"| Eval correct (all behavioral tests pass) | {eval_pass} | {eval_pass * 100 // total}% |")
    lines.append(f"| Faster than original (speedup > 1.0) | {faster} | {faster * 100 // total}% |")
    lines.append(f"| All test cases pass | {all_tests_pass} | {all_tests_pass * 100 // total}% |")
    lines.append("")

    # Speedup distribution
    speedups = [r.reward.get("speedup", 0) for r in results if r.eval_ok]
    if speedups:
        lines.append("## Speedup Distribution (correct tasks only)\n")
        ranges = [
            (0, 1, "Slower (< 1x)"),
            (1, 1.5, "1-1.5x"),
            (1.5, 2, "1.5-2x"),
            (2, 5, "2-5x"),
            (5, 100, "5-100x"),
            (100, float("inf"), ">100x"),
        ]
        for lo, hi, label in ranges:
            count = sum(1 for s in speedups if lo <= s < hi)
            if count:
                lines.append(f"- **{label}**: {count} tasks")
        lines.append("")

    # Results table — successful tasks
    lines.append("## Successful Tasks (correct=1.0)\n")
    lines.append("| Task | Function | Speedup | Tests | Coverage | Quality | DB Speedup |")
    lines.append("|------|----------|---------|-------|----------|---------|------------|")

    for r in sorted(results, key=lambda x: x.reward.get("speedup", 0), reverse=True):
        if not r.eval_ok:
            continue
        rw = r.reward
        speedup = rw.get("speedup", 0)
        tests = f"{rw.get('tests_passed', 0)}/{rw.get('tests_total', 0)}"
        coverage = f"{rw.get('coverage_pct', 'N/A')}%" if rw.get("coverage_pct") else "N/A"
        quality = rw.get("optimization_quality", "")
        # DB stores speedup as fractional improvement (0.3 = 30% faster = 1.3x ratio)
        db_spd = f"{r.speedup_db + 1:.2f}x" if r.speedup_db else "N/A"
        lines.append(
            f"| {r.task_name} | `{r.function_name}` | {speedup:.4f}x | {tests} | {coverage} | {quality} | {db_spd} |"
        )
    lines.append("")

    # Failed tasks — detailed
    failed = [r for r in results if not r.eval_ok]
    if failed:
        lines.append(f"## Failed Tasks ({len(failed)})\n")
        for r in failed:
            lines.append(f"### {r.task_name}\n")
            lines.append(f"- **Function**: `{r.function_name}`")
            lines.append(f"- **File**: `{r.file_path}`")
            lines.append(f"- **Commit**: `{r.commit}`")
            lines.append(f"- **Method**: {r.resolution_method}")
            lines.append(f"- **DB Speedup**: {r.speedup_db + 1:.2f}x" if r.speedup_db else "- **DB Speedup**: N/A")
            lines.append(f"- **Solve OK**: {r.solve_ok}")
            lines.append(f"- **Error**: {r.error}" if r.error else "")
            lines.append(f"- **Duration**: {r.duration_secs:.1f}s")

            rw = r.reward
            if rw:
                lines.append(
                    f"- **Reward**: correct={rw.get('correct')}, speedup={rw.get('speedup')}, "
                    f"tests={rw.get('tests_passed')}/{rw.get('tests_total')}"
                )

            # Solve output
            if r.solve_output:
                lines.append("\n<details><summary>Solve output</summary>\n")
                lines.append(f"```\n{r.solve_output[-500:]}\n```\n</details>\n")

            # Eval output (last 1000 chars for failures)
            if r.eval_output:
                # Extract key error lines
                error_lines = [
                    l
                    for l in r.eval_output.splitlines()
                    if "ERROR" in l or "FAIL" in l or "INCORRECT" in l or "error" in l.lower()
                ]
                if error_lines:
                    lines.append("\n<details><summary>Key errors</summary>\n")
                    lines.append("```")
                    for el in error_lines[:20]:
                        lines.append(el)
                    lines.append("```\n</details>\n")

            # Reproduction
            lines.append(f"\n**Reproduce**: `bash docker_e2e_test.sh {r.task_name} --debug`\n")
            lines.append("---\n")

    output_path.write_text("\n".join(lines), encoding="utf-8")
    log.info(f"Report written to {output_path}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--output", type=Path, default=Path("rl_env/validation_report.md"))
    parser.add_argument("--catalog", type=Path, default=Path("rl_env/resolved_catalog_full.json"))
    parser.add_argument("--tasks-file", type=Path, default=None, help="File with task names to run (one per line)")
    args = parser.parse_args()

    catalog = json.loads(args.catalog.read_text(encoding="utf-8"))
    catalog_by_func = {}
    for e in catalog:
        gp = Path(e["file_path"]).parent.name
        name = f"{gp}-{e['function_name']}".replace(".", "-")
        name = re.sub(r"[^a-zA-Z0-9_-]", "-", name)
        name = re.sub(r"-+", "-", name).strip("-").lower()[:80]
        catalog_by_func[name] = e

    if args.tasks_file:
        task_names = [l.strip() for l in args.tasks_file.read_text(encoding="utf-8").splitlines() if l.strip()]
        task_dirs = [TASKS_DIR / name for name in task_names if (TASKS_DIR / name).is_dir()]
    else:
        task_dirs = sorted([d for d in TASKS_DIR.iterdir() if d.is_dir() and d.name != "base-image"])
    if args.limit > 0:
        task_dirs = task_dirs[: args.limit]

    log.info(f"Validating {len(task_dirs)} tasks...")

    results = []
    for i, task_dir in enumerate(task_dirs):
        entry = catalog_by_func.get(task_dir.name, {})
        log.info(f"  [{i + 1}/{len(task_dirs)}] {task_dir.name}...")

        result = run_task(task_dir, entry)
        results.append(result)

        status = "OK" if result.eval_ok else "FAIL"
        speedup = result.reward.get("speedup", 0)
        tests = f"{result.reward.get('tests_passed', 0)}/{result.reward.get('tests_total', 0)}"
        log.info(f"    {status} speedup={speedup:.2f}x tests={tests} ({result.duration_secs:.0f}s)")

    generate_report(results, args.output)

    # Quick summary
    total = len(results)
    ok = sum(1 for r in results if r.eval_ok)
    log.info(f"\nDone: {ok}/{total} correct")


if __name__ == "__main__":
    main()
