"""Grade long-horizon coding episodes with official SWE-bench Docker evaluation.

Reward = FAIL_TO_PASS fraction after applying the agent's git patch inside the
instance image and running the official eval_script (test_patch + pytest).

Resolution follows SWE-bench: FULL only if all F2P pass and all P2P maintained.
"""
from __future__ import annotations

from pathlib import Path
from typing import Any

from deepagents_rl.harness.repo import git_diff_patch
from deepagents_rl.harness.swe_docker_eval import evaluate_model_patch
from deepagents_rl.observability.logging import span, step, truncate


def grade_episode(
    *,
    sandbox: Path,
    task: dict[str, Any],
    tools_used: list[str] | None = None,
    tool_calls: list[dict[str, Any]] | None = None,
    model_patch: str | None = None,
    eval_log_dir: Path | None = None,
    timeout_s: int = 1800,
) -> dict[str, Any]:
    """
    Produce model_patch from the agent sandbox (if not provided), then run
    Docker SWE-bench evaluation.
    """
    tools_used = tools_used or []
    tool_calls = tool_calls or []

    with span(
        "grade.episode",
        task.get("instance_id") or task.get("task_id"),
        sandbox=str(sandbox),
        n_tool_calls=len(tool_calls),
    ):
        if model_patch is None:
            step("grade.diff", "collecting git diff from agent sandbox")
            model_patch = git_diff_patch(sandbox)
        else:
            step("grade.diff", "using provided model_patch")

        step(
            "grade.patch_ready",
            f"bytes={len(model_patch.encode('utf-8'))}",
            empty=not bool(model_patch.strip()),
            preview=truncate(model_patch, 400),
        )
        step(
            "grade.docker_eval",
            f"image={task.get('image')}",
            n_f2p=len(task.get("fail_to_pass") or []),
            n_p2p=len(task.get("pass_to_pass") or []),
        )
        report = evaluate_model_patch(
            task,
            model_patch,
            timeout_s=timeout_s,
            log_dir=eval_log_dir,
        )

        out = {
            "pass_frac": float(report.get("pass_frac") or 0.0),
            "reward_total": float(report.get("reward_total") or 0.0),
            "resolved": bool(report.get("resolved")),
            "resolution": report.get("resolution") or "NO",
            "patch_successfully_applied": bool(report.get("patch_successfully_applied")),
            "model_patch": model_patch,
            "fail_to_pass": report.get("fail_to_pass"),
            "pass_to_pass": report.get("pass_to_pass"),
            "status_map": report.get("status_map") or {},
            "eval_error": report.get("error"),
            "eval_log_dir": report.get("log_dir"),
            "eval_exit_code": report.get("eval_exit_code"),
            "n_tool_calls": len(tool_calls),
            "tools_used": tools_used,
            "image": report.get("image") or task.get("image"),
            "grader": "swebench_docker_official_protocol",
        }
        step(
            "grade.done",
            f"resolution={out['resolution']}",
            resolved=out["resolved"],
            pass_frac=out["pass_frac"],
            eval_error=out["eval_error"],
        )
        return out