"""Precompute code coverage for each task's test suite over the original function.

Runs coverage.py inside Docker at the pre-optimization commit for each task.
Stores results in coverage_map.json.

Usage:
    python rl_env/precompute_coverage.py --limit 5
"""

from __future__ import annotations

import argparse
import json
import logging
import re
import subprocess
import tempfile
from pathlib import Path

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

DOCKER_IMAGE = "codeflash-inference-base:latest"
TASKS_DIR = Path(__file__).parent / "codeflash-perf-tasks"
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 compute_coverage_for_task(
    task_dir: Path, entry: dict,
) -> float | None:
    """Run coverage.py inside Docker for one task. Returns coverage percentage or None."""
    commit = entry["pre_optimization_commit"]
    file_path = entry["file_path"]
    func_name = entry["function_name"].rsplit(".", maxsplit=1)[-1]

    test_dir = task_dir / "tests" / "codeflash_eval"
    behavioral_tests = sorted(test_dir.glob("*behaviorinstrumented*.py"))
    if not behavioral_tests:
        return None

    # Build a script that runs coverage inside Docker
    script = f"""
import subprocess, json, sys
from pathlib import Path

project = Path("/workspace/inference")
subprocess.run(["git", "checkout", "{commit}"], cwd=project, capture_output=True)

# Run behavioral tests with coverage
result = subprocess.run(
    [sys.executable, "-m", "coverage", "run",
     "--source={file_path.rsplit('/', 1)[0]}",
     "--include={file_path}",
     "-m", "pytest", "/tmp/tests/", "-q", "--no-header",
     "-p", "no:benchmark", "-p", "no:codspeed", "-p", "no:xdist",
     "--override-ini=addopts="],
    cwd=project, capture_output=True, text=True, timeout=120,
    env={{**__import__("os").environ, "PYTHONPATH": str(project),
         "CODEFLASH_TEST_ITERATION": "0", "CODEFLASH_LOOP_INDEX": "1",
         "CODEFLASH_TRACER_DISABLE": "1"}},
)

# Get coverage report
result2 = subprocess.run(
    [sys.executable, "-m", "coverage", "json", "-o", "/dev/stdout"],
    cwd=project, capture_output=True, text=True,
)

try:
    cov_data = json.loads(result2.stdout)
    total = cov_data.get("totals", {{}})
    pct = total.get("percent_covered", 0)
    print(json.dumps({{"coverage_pct": pct}}))
except Exception:
    print(json.dumps({{"coverage_pct": None, "error": result2.stderr[:200]}}))
"""

    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8") as f:
        f.write(script)
        script_path = Path(f.name)

    # Mount test files into Docker
    result = subprocess.run(
        [
            "docker", "run", "--rm",
            "-v", f"{script_path}:/tmp/coverage_script.py:ro",
            "-v", f"{test_dir}:/tmp/tests:ro",
            DOCKER_IMAGE,
            "python", "/tmp/coverage_script.py",
        ],
        capture_output=True, text=True, encoding="utf-8", timeout=180,
    )

    script_path.unlink(missing_ok=True)

    if result.returncode != 0:
        return None

    for line in reversed(result.stdout.strip().splitlines()):
        if line.strip().startswith("{"):
            try:
                data = json.loads(line)
                return data.get("coverage_pct")
            except json.JSONDecodeError:
                continue

    return None


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--output", type=Path, default=Path("rl_env/coverage_map.json"))
    args = parser.parse_args()

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

    coverage_map: dict[str, float | None] = {}
    processed = 0

    for entry in catalog:
        task_name = sanitize_task_name(entry["function_name"], entry["file_path"])
        task_dir = TASKS_DIR / task_name
        if not task_dir.exists():
            continue

        if args.limit > 0 and processed >= args.limit:
            break

        log.info(f"  [{processed + 1}] {entry['function_name']}...")
        coverage = compute_coverage_for_task(task_dir, entry)
        coverage_map[entry["trace_id"]] = coverage
        if coverage is not None:
            log.info(f"    Coverage: {coverage:.1f}%")
        processed += 1

    args.output.write_text(json.dumps(coverage_map, indent=2), encoding="utf-8")
    log.info(f"\nWrote {args.output}")

    with_coverage = sum(1 for v in coverage_map.values() if v is not None)
    log.info(f"Coverage computed for {with_coverage}/{len(coverage_map)} tasks")


if __name__ == "__main__":
    main()
