"""Official-style SWE-bench evaluation via the instance Docker image.

Protocol (matches SWE-bench harness):
  1. Start the prebuilt instance image (`task['image']`)
  2. Apply the model's git patch into `/testbed`
  3. Run the instance `eval_script` (applies test_patch, runs pytest)
  4. Parse status between `>>>>> Start/End Test Output`
  5. Resolve = all FAIL_TO_PASS pass AND all PASS_TO_PASS pass

Uses the Docker CLI (not the Python SDK) for reliability on Windows.
"""
from __future__ import annotations

import json
import shutil
import subprocess
import tempfile
import time
import uuid
from pathlib import Path
from typing import Any

from deepagents_rl.observability.logging import call, error, result, span, step, truncate

START_TEST_OUTPUT = ">>>>> Start Test Output"
END_TEST_OUTPUT = ">>>>> End Test Output"

GIT_APPLY_CMDS = [
    "git apply --verbose",
    "git apply --verbose --3way",
    "git apply --verbose --reject",
]


def fix_eval_script_for_docker_desktop(script: str) -> str:
    """
    Upstream eval_script uses `conda activate testbed`. On Docker Desktop the
    image's *base* conda entrypoint is often broken (`conda` -> TypeError), while
    `/opt/miniconda3/envs/testbed/bin/*` works. Prefetch that env onto PATH and
    stub `conda activate` so the official test_patch + pytest body still runs.
    """
    preamble = (
        "#!/bin/bash\n"
        "set -uxo pipefail\n"
        'export PATH="/opt/miniconda3/envs/testbed/bin:/opt/miniconda3/bin:$PATH"\n'
        "conda() {\n"
        '  if [ "$1" = "activate" ]; then\n'
        '    export PATH="/opt/miniconda3/envs/${2:-testbed}/bin:$PATH"\n'
        "    hash -r\n"
        "    return 0\n"
        "  fi\n"
        "  return 0\n"
        "}\n"
    )
    body: list[str] = []
    for line in script.replace("\r\n", "\n").splitlines():
        if line.startswith("#!"):
            continue
        if line.strip().startswith("source /opt/miniconda3/bin/activate"):
            continue
        if "conda activate" in line:
            continue
        body.append(line)
    return preamble + "\n".join(body) + "\n"


def parse_log_pytest(log: str) -> dict[str, str]:
    """SWE-bench pytest log parser (status token first on the line)."""
    status_map: dict[str, str] = {}
    statuses = ("PASSED", "FAILED", "ERROR", "SKIPPED", "XFAIL", "XPASS")
    for line in log.splitlines():
        if not any(line.startswith(s) for s in statuses):
            continue
        if line.startswith("FAILED"):
            line = line.replace(" - ", " ")
        parts = line.split()
        if len(parts) <= 1:
            continue
        status_map[parts[1]] = parts[0]
    return status_map


def _extract_test_region(log: str) -> str | None:
    if START_TEST_OUTPUT not in log or END_TEST_OUTPUT not in log:
        return None
    return log.split(START_TEST_OUTPUT, 1)[1].split(END_TEST_OUTPUT, 1)[0]


def _test_passed(case: str, sm: dict[str, str]) -> bool:
    return case in sm and sm[case] in ("PASSED", "XFAIL")


def _test_failed(case: str, sm: dict[str, str]) -> bool:
    return case not in sm or sm[case] in ("FAILED", "ERROR", "SKIPPED")


def _test_maintained(case: str, sm: dict[str, str]) -> bool:
    return _test_passed(case, sm) or (case in sm and sm[case] == "SKIPPED")


def grade_status_map(
    status_map: dict[str, str],
    fail_to_pass: list[str],
    pass_to_pass: list[str],
) -> dict[str, Any]:
    """Official resolution: FULL iff all F2P pass and all P2P maintained."""
    f2p_ok = [t for t in fail_to_pass if _test_passed(t, status_map)]
    f2p_fail = [t for t in fail_to_pass if _test_failed(t, status_map)]
    p2p_ok = [t for t in pass_to_pass if _test_maintained(t, status_map)]
    p2p_fail = [t for t in pass_to_pass if not _test_maintained(t, status_map)]

    f2p_frac = (len(f2p_ok) / len(fail_to_pass)) if fail_to_pass else 1.0
    p2p_frac = (len(p2p_ok) / len(pass_to_pass)) if pass_to_pass else 1.0

    if f2p_frac == 1.0 and p2p_frac == 1.0:
        resolved = "FULL"
    elif 0.0 < f2p_frac < 1.0 and p2p_frac == 1.0:
        resolved = "PARTIAL"
    else:
        resolved = "NO"

    return {
        "resolved": resolved == "FULL",
        "resolution": resolved,
        "fail_to_pass": {"success": f2p_ok, "failure": f2p_fail, "frac": f2p_frac},
        "pass_to_pass": {"success": p2p_ok, "failure": p2p_fail, "frac": p2p_frac},
        "pass_frac": float(f2p_frac),
        # P2P regressions discount reward (cannot claim FULL/PARTIAL)
        "reward_total": float(f2p_frac if p2p_frac == 1.0 else f2p_frac * 0.5),
        "status_map": status_map,
    }


def _docker(
    *args: str,
    timeout: int | None = None,
    merge_stderr: bool = False,
) -> subprocess.CompletedProcess[str]:
    exe = shutil.which("docker")
    if not exe:
        raise RuntimeError("docker CLI not found on PATH")
    cmdline = "docker " + " ".join(args)
    call("docker", cmdline, timeout_s=timeout, merge_stderr=merge_stderr)
    t0 = time.perf_counter()
    try:
        if merge_stderr:
            proc = subprocess.run(
                [exe, *args],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                timeout=timeout,
                encoding="utf-8",
                errors="replace",
            )
        else:
            proc = subprocess.run(
                [exe, *args],
                capture_output=True,
                text=True,
                timeout=timeout,
                encoding="utf-8",
                errors="replace",
            )
    except subprocess.TimeoutExpired:
        error("docker", f"timeout after {timeout}s", cmd=cmdline)
        raise
    dt = int((time.perf_counter() - t0) * 1000)
    out = (proc.stdout or "") + ("" if merge_stderr else (proc.stderr or ""))
    result(
        "docker",
        f"exit={proc.returncode} ({dt}ms)",
        exit_code=proc.returncode,
        duration_ms=dt,
        output=truncate(out, 800),
    )
    return proc


def evaluate_model_patch(
    task: dict[str, Any],
    model_patch: str,
    *,
    timeout_s: int = 1800,
    log_dir: Path | None = None,
    keep_container: bool = False,
) -> dict[str, Any]:
    """
    Grade one prediction with the official instance Docker image + eval_script.

    `task` must include: instance_id, image, eval_script, fail_to_pass, pass_to_pass
    """
    instance_id = task["instance_id"]
    image = task.get("image")
    eval_script = task.get("eval_script") or ""
    fail_to_pass = list(task.get("fail_to_pass") or [])
    pass_to_pass = list(task.get("pass_to_pass") or [])

    if not image:
        raise ValueError(f"{instance_id}: missing Docker image field from SWE-bench Lite")
    if not eval_script.strip():
        raise ValueError(f"{instance_id}: missing eval_script")

    if log_dir is None:
        log_dir = Path(tempfile.mkdtemp(prefix="swe_eval_"))
    log_dir.mkdir(parents=True, exist_ok=True)

    report: dict[str, Any] = {
        "instance_id": instance_id,
        "image": image,
        "patch_exists": bool(model_patch and model_patch.strip()),
        "patch_successfully_applied": False,
        "resolved": False,
        "resolution": "NO",
        "pass_frac": 0.0,
        "reward_total": 0.0,
        "fail_to_pass": {"success": [], "failure": fail_to_pass, "frac": 0.0},
        "pass_to_pass": {"success": [], "failure": [], "frac": 0.0},
        "status_map": {},
        "error": None,
        "log_dir": str(log_dir),
    }

    step(
        "eval.begin",
        f"instance={instance_id}",
        image=image,
        patch_bytes=len((model_patch or "").encode("utf-8")),
        n_f2p=len(fail_to_pass),
        n_p2p=len(pass_to_pass),
        log_dir=str(log_dir),
    )

    if not model_patch or not model_patch.strip():
        step("eval.empty_patch", "no model_patch — unresolved")
        report["error"] = "empty_model_patch"
        (log_dir / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
        return report

    patch_path = log_dir / "model.patch"
    eval_path = log_dir / "eval.sh"
    eval_raw_path = log_dir / "eval.upstream.sh"
    patch_path.write_text(model_patch.replace("\r\n", "\n"), encoding="utf-8", newline="\n")
    eval_raw_path.write_text(eval_script.replace("\r\n", "\n"), encoding="utf-8", newline="\n")
    eval_path.write_text(
        fix_eval_script_for_docker_desktop(eval_script),
        encoding="utf-8",
        newline="\n",
    )
    step(
        "eval.write_files",
        "wrote model.patch + eval.sh",
        patch_path=str(patch_path),
        eval_path=str(eval_path),
        patch_preview=truncate(model_patch, 400),
    )

    run_tag = uuid.uuid4().hex[:10]
    container = f"darl-eval-{run_tag}"

    try:
        with span("eval.docker_lifecycle", container=container, image=image):
            inspect = _docker("image", "inspect", image, timeout=60)
            if inspect.returncode != 0:
                step("eval.docker_pull", f"image missing locally, pulling {image}")
                pull = _docker("pull", image, timeout=max(timeout_s, 3600))
                if pull.returncode != 0:
                    report["error"] = f"docker_pull_failed: {(pull.stderr or pull.stdout)[-2000:]}"
                    error("eval.docker_pull", report["error"])
                    (log_dir / "report.json").write_text(
                        json.dumps(report, indent=2), encoding="utf-8"
                    )
                    return report
            else:
                step("eval.docker_image", "image already present")

            _docker("rm", "-f", container, timeout=60)

            create = _docker(
                "create",
                "--name",
                container,
                "--workdir",
                "/testbed",
                image,
                "tail",
                "-f",
                "/dev/null",
                timeout=120,
            )
            if create.returncode != 0:
                report["error"] = (
                    f"docker_create_failed: {(create.stderr or create.stdout)[-2000:]}"
                )
                error("eval.docker_create", report["error"])
                (log_dir / "report.json").write_text(
                    json.dumps(report, indent=2), encoding="utf-8"
                )
                return report

            start = _docker("start", container, timeout=120)
            if start.returncode != 0:
                report["error"] = (
                    f"docker_start_failed: {(start.stderr or start.stdout)[-2000:]}"
                )
                error("eval.docker_start", report["error"])
                (log_dir / "report.json").write_text(
                    json.dumps(report, indent=2), encoding="utf-8"
                )
                return report

            step("eval.docker_cp", "copy patch + eval.sh into container")
            cp1 = _docker("cp", str(patch_path), f"{container}:/tmp/model.patch", timeout=120)
            cp2 = _docker("cp", str(eval_path), f"{container}:/eval.sh", timeout=120)
            if cp1.returncode != 0 or cp2.returncode != 0:
                report["error"] = (
                    f"docker_cp_failed: patch={(cp1.stderr or '')[-500:]} "
                    f"eval={(cp2.stderr or '')[-500:]}"
                )
                error("eval.docker_cp", report["error"])
                (log_dir / "report.json").write_text(
                    json.dumps(report, indent=2), encoding="utf-8"
                )
                return report

            _docker(
                "exec",
                "-w",
                "/testbed",
                container,
                "git",
                "config",
                "--global",
                "--add",
                "safe.directory",
                "/testbed",
                timeout=60,
            )

            apply_log = []
            applied = False
            step("eval.apply_patch", "trying git apply cascade")
            for cmd in GIT_APPLY_CMDS:
                parts = cmd.split()
                resp = _docker(
                    "exec",
                    "-w",
                    "/testbed",
                    container,
                    *parts,
                    "/tmp/model.patch",
                    timeout=120,
                )
                apply_log.append(
                    {
                        "cmd": f"{cmd} /tmp/model.patch",
                        "exit": resp.returncode,
                        "out": ((resp.stdout or "") + (resp.stderr or ""))[-2000:],
                    }
                )
                if resp.returncode == 0:
                    applied = True
                    step("eval.apply_patch_ok", cmd)
                    break
                step("eval.apply_patch_retry", f"failed: {cmd}", exit_code=resp.returncode)
            (log_dir / "apply_patch.json").write_text(
                json.dumps(apply_log, indent=2), encoding="utf-8"
            )

            if not applied:
                report["error"] = "patch_apply_failed"
                report["apply_log"] = apply_log
                error("eval.apply_patch", "all git apply strategies failed")
                (log_dir / "test_output.txt").write_text(
                    json.dumps(apply_log, indent=2), encoding="utf-8"
                )
                (log_dir / "report.json").write_text(
                    json.dumps(report, indent=2), encoding="utf-8"
                )
                return report

            report["patch_successfully_applied"] = True

            step("eval.run_eval_script", "bash /eval.sh (official test_patch + pytest)")
            t0 = time.time()
            eval_proc = _docker(
                "exec",
                "-w",
                "/testbed",
                container,
                "/bin/bash",
                "-lc",
                "/bin/bash /eval.sh 2>&1",
                timeout=timeout_s,
                merge_stderr=True,
            )
            elapsed = time.time() - t0
            test_output = eval_proc.stdout or ""
            (log_dir / "test_output.txt").write_text(test_output, encoding="utf-8")
            report["eval_exit_code"] = eval_proc.returncode
            report["eval_seconds"] = elapsed
            step(
                "eval.eval_script_done",
                f"exit={eval_proc.returncode} seconds={elapsed:.1f}",
                output_chars=len(test_output),
                output_tail=truncate(test_output, 600),
            )

            region = _extract_test_region(test_output)
            status_map = parse_log_pytest(region or "")
            if not status_map:
                step("eval.parse_fallback", "marked region empty — parsing full log")
                status_map = parse_log_pytest(test_output)
            if not status_map:
                report["error"] = (
                    "missing_test_output_markers"
                    if region is None
                    else "no_parsed_test_statuses"
                )
                error("eval.parse", report["error"])
                (log_dir / "report.json").write_text(
                    json.dumps(report, indent=2), encoding="utf-8"
                )
                return report

            step(
                "eval.parse_ok",
                f"parsed {len(status_map)} test statuses",
                sample=list(status_map.items())[:8],
            )
            graded = grade_status_map(status_map, fail_to_pass, pass_to_pass)
            report.update(graded)
            report["error"] = None
            step(
                "eval.grade",
                f"resolution={graded['resolution']} f2p={graded['pass_frac']}",
                resolved=graded["resolved"],
                resolution=graded["resolution"],
                pass_frac=graded["pass_frac"],
                reward_total=graded["reward_total"],
                f2p_success=graded["fail_to_pass"]["success"],
                f2p_failure=graded["fail_to_pass"]["failure"],
                p2p_frac=graded["pass_to_pass"]["frac"],
            )
            (log_dir / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
            return report

    except subprocess.TimeoutExpired as e:
        report["error"] = f"timeout: {e}"
        error("eval.timeout", report["error"])
        (log_dir / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
        return report
    except Exception as e:  # noqa: BLE001
        report["error"] = f"{type(e).__name__}: {e}"
        error("eval.exception", report["error"])
        (log_dir / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
        return report
    finally:
        if not keep_container:
            try:
                step("eval.cleanup", f"rm container {container}")
                _docker("rm", "-f", container, timeout=120)
            except Exception as e:  # noqa: BLE001
                error("eval.cleanup", str(e))