"""Run one deepagents episode on a long-horizon coding task and record everything."""
from __future__ import annotations

import json
import os
import tempfile
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from deepagents_rl.harness.grade import grade_episode
from deepagents_rl.harness.repo import materialize_repo_at_commit
from deepagents_rl.observability.logging import (
    call,
    error as log_error,
    result,
    span,
    step,
    truncate,
)


SYSTEM_PROMPT = (
    "You are a software engineering agent with tools.\n"
    "You MUST use tools to inspect and edit files. Never answer with only prose.\n"
    "First action: call ls or glob to see the sandbox.\n"
    "Then read_file the relevant source, edit_file/write_file to fix the bug, "
    "optionally execute tests.\n"
    "When finished, stop calling tools and reply DONE plus a short summary.\n"
    "Do not delete unrelated files."
)


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _msg_to_dict(msg: Any) -> dict[str, Any]:
    role = getattr(msg, "type", None) or getattr(msg, "role", None) or msg.__class__.__name__
    content = getattr(msg, "content", None)
    if not isinstance(content, str):
        content = str(content) if content is not None else ""
    out: dict[str, Any] = {"role": role, "content": content[:8000]}
    tool_calls = getattr(msg, "tool_calls", None)
    if tool_calls:
        out["tool_calls"] = tool_calls
    name = getattr(msg, "name", None)
    if name:
        out["name"] = name
    tool_call_id = getattr(msg, "tool_call_id", None)
    if tool_call_id:
        out["tool_call_id"] = tool_call_id
    return out


def _extract_tool_calls(messages: list[Any]) -> tuple[list[str], list[dict[str, Any]]]:
    tools_used: list[str] = []
    tool_calls_log: list[dict[str, Any]] = []
    pending: dict[str, dict[str, Any]] = {}
    turn = 0
    for msg in messages:
        tname = msg.__class__.__name__
        if tname == "AIMessage" and getattr(msg, "tool_calls", None):
            turn += 1
            for tc in msg.tool_calls:
                name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", None)
                args = tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {})
                tc_id = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
                if name:
                    tools_used.append(name)
                entry = {
                    "turn": turn,
                    "name": name,
                    "args": args or {},
                    "tool_call_id": tc_id,
                    "result": None,
                    "ok": None,
                }
                if tc_id:
                    pending[tc_id] = entry
                tool_calls_log.append(entry)
        elif tname == "ToolMessage":
            tc_id = getattr(msg, "tool_call_id", None)
            content = getattr(msg, "content", "")
            content_s = content if isinstance(content, str) else str(content)
            entry = pending.get(tc_id) if tc_id else None
            if entry is None:
                entry = {
                    "turn": turn,
                    "name": getattr(msg, "name", None),
                    "args": {},
                    "tool_call_id": tc_id,
                }
                tool_calls_log.append(entry)
            entry["result"] = content_s[:4000]
            entry["ok"] = not content_s.startswith("Error") and "ERROR" not in content_s[:40]
    return tools_used, tool_calls_log


def _short_prompt(task: dict[str, Any]) -> str:
    """Compact prompt — long issue bodies break weak GGUF tool-calling templates."""
    body = (task.get("problem_statement") or "").strip()
    if len(body) > 1800:
        body = body[:1800] + "\n...(truncated)..."
    paths = task.get("workspace_paths") or list((task.get("workspace_files") or {}).keys())
    path_hint = ", ".join(paths[:8]) if paths else "(explore with ls/glob)"
    return (
        f"Fix this GitHub issue in the sandbox using tools.\n"
        f"Instance: {task['task_id']}\n"
        f"Likely files: {path_hint}\n\n"
        f"Issue:\n{body}\n\n"
        f"Required: call tools (start with ls or read_file). "
        f"When done, reply DONE."
    )


def _result_to_str(obj: Any) -> str:
    if obj is None:
        return ""
    if hasattr(obj, "model_dump"):
        try:
            return json.dumps(obj.model_dump(), default=str)[:4000]
        except Exception:
            pass
    if hasattr(obj, "__dict__"):
        try:
            return json.dumps(obj.__dict__, default=str)[:4000]
        except Exception:
            pass
    return str(obj)[:4000]


def _build_backend_tools(backend: Any):
    """LangChain tools wrapping deepagents LocalShellBackend (same FS + execute)."""
    from langchain_core.tools import tool

    @tool
    def ls(path: str = ".") -> str:
        """List files in a directory inside the sandbox."""
        return _result_to_str(backend.ls(path))

    @tool
    def glob(pattern: str, path: str = ".") -> str:
        """Glob for files matching pattern under path."""
        return _result_to_str(backend.glob(pattern, path))

    @tool
    def grep(pattern: str, path: str = ".") -> str:
        """Search file contents for a regex pattern."""
        return _result_to_str(backend.grep(pattern, path))

    @tool
    def read_file(file_path: str, offset: int = 0, limit: int = 200) -> str:
        """Read a text file from the sandbox (line offset/limit)."""
        return _result_to_str(backend.read(file_path, offset=offset, limit=limit))

    @tool
    def write_file(file_path: str, content: str) -> str:
        """Write/create a text file in the sandbox."""
        return _result_to_str(backend.write(file_path, content))

    @tool
    def edit_file(file_path: str, old_string: str, new_string: str) -> str:
        """Replace old_string with new_string in an existing file."""
        return _result_to_str(backend.edit(file_path, old_string, new_string))

    @tool
    def execute(command: str) -> str:
        """Run a shell command in the sandbox (e.g. python -c '...')."""
        return _result_to_str(backend.execute(command))

    return [ls, glob, grep, read_file, write_file, edit_file, execute]


def _run_tool_loop(
    *,
    model: Any,
    backend: Any,
    prompt: str,
    max_turns: int,
    langfuse_client: Any | None = None,
    callback_handler: Any | None = None,
    episode_id: str | None = None,
) -> tuple[list[Any], str, str]:
    """Reliable multi-turn tool loop on deepagents LocalShellBackend tools.

    When `langfuse_client` is set, each LLM turn is traced via CallbackHandler
    (generation) and each tool execution is a nested `as_type="tool"` span.
    """
    from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage

    from deepagents_rl.observability.langfuse_trace import llm_invoke_config

    tools = _build_backend_tools(backend)
    tool_map = {t.name: t for t in tools}
    tool_names = list(tool_map.keys())
    step("agent.bind_tools", f"tools={tool_names}", tools=tool_names)
    bound = model.bind_tools(tools)
    messages: list[Any] = [
        SystemMessage(content=SYSTEM_PROMPT),
        HumanMessage(content=prompt),
    ]
    step("agent.prompt", truncate(prompt, 500), prompt_chars=len(prompt))
    final_text = ""
    termination = "max_steps"

    def _llm_call(turn: int) -> AIMessage:
        cfg = None
        if callback_handler is not None:
            cfg = llm_invoke_config(
                handler=callback_handler,
                turn=turn,
                episode_id=episode_id,
            )
        if cfg is not None:
            return bound.invoke(messages, config=cfg)
        return bound.invoke(messages)

    for turn in range(max_turns):
        def _one_turn(turn: int = turn) -> str | None:
            """Run one LLM+tools turn. Returns termination if loop should stop."""
            nonlocal final_text, termination
            with span("agent.llm_turn", turn=turn):
                call("llm.invoke", f"turn={turn} n_messages={len(messages)}")
                t0 = time.time()
                ai: AIMessage = _llm_call(turn)
                dt = int((time.time() - t0) * 1000)
                messages.append(ai)
                content_preview = (
                    ai.content if isinstance(ai.content, str) else str(ai.content)
                )
                tcs = getattr(ai, "tool_calls", None) or []
                result(
                    "llm.invoke",
                    f"turn={turn} ({dt}ms)",
                    duration_ms=dt,
                    content=truncate(content_preview, 400),
                    n_tool_calls=len(tcs),
                    tool_calls=[
                        {
                            "name": tc.get("name"),
                            "args": truncate(tc.get("args") or {}, 300),
                        }
                        for tc in tcs
                    ],
                )

            if not tcs:
                final_text = content_preview
                termination = (
                    "agent_done"
                    if "DONE" in (final_text or "").upper()
                    else "no_tool_final"
                )
                step("agent.final", termination, final_text=truncate(final_text, 400))
                return termination

            tool_results_for_turn: list[dict[str, Any]] = []
            for tc in tcs:
                name = tc["name"]
                args = tc.get("args") or {}
                tool_out: Any
                with span("agent.tool", turn=turn, tool=name):
                    call("tool", name, args=truncate(args, 500))

                    def _exec_tool() -> Any:
                        try:
                            return tool_map[name].invoke(args)
                        except Exception as e:  # noqa: BLE001
                            err = f"ERROR: {type(e).__name__}: {e}"
                            log_error("tool", err, tool=name)
                            return err

                    if langfuse_client is not None:
                        with langfuse_client.start_as_current_observation(
                            name=name,
                            as_type="tool",
                            input=args if isinstance(args, dict) else {"value": args},
                            metadata={
                                "turn": turn,
                                "tool_call_id": tc.get("id"),
                                "episode_id": episode_id,
                            },
                        ) as tool_span:
                            tool_out = _exec_tool()
                            tool_span.update(
                                output={"result": truncate(str(tool_out), 4000)}
                            )
                    else:
                        tool_out = _exec_tool()

                    result("tool", name, output=truncate(str(tool_out), 600))
                    tool_results_for_turn.append(
                        {
                            "name": name,
                            "args": truncate(args, 300),
                            "result": truncate(str(tool_out), 300),
                        }
                    )
                    messages.append(
                        ToolMessage(
                            content=str(tool_out)[:4000],
                            tool_call_id=tc["id"],
                            name=name,
                        )
                    )
            return None  # continue

        if langfuse_client is not None:
            with langfuse_client.start_as_current_observation(
                name=f"turn-{turn}",
                as_type="span",
                metadata={"turn": turn, "episode_id": episode_id},
                input={"n_messages": len(messages)},
            ) as turn_span:
                stop = _one_turn()
                if stop:
                    turn_span.update(
                        output={
                            "termination": stop,
                            "final_text": truncate(final_text, 1000),
                            "n_tool_calls": 0,
                        }
                    )
                    break
                n_tc = 0
                for m in reversed(messages):
                    if m.__class__.__name__ == "AIMessage" and getattr(
                        m, "tool_calls", None
                    ):
                        n_tc = len(m.tool_calls)
                        break
                turn_span.update(output={"n_tool_calls": n_tc, "continued": True})
        else:
            if _one_turn():
                break
    else:
        termination = "max_steps"
        final_text = "(hit max turns)"
        step("agent.max_turns", f"hit max_turns={max_turns}")

    return messages, final_text, termination


def _run_graph(
    *,
    model: Any,
    backend: Any,
    prompt: str,
    recursion_limit: int,
    langfuse_client: Any | None = None,
    callback_handler: Any | None = None,
    episode_id: str | None = None,
) -> tuple[list[Any], str, str]:
    from deepagents import create_deep_agent

    step("agent.create_deep_agent", f"recursion_limit={recursion_limit}")
    agent = create_deep_agent(
        model=model,
        backend=backend,
        system_prompt=SYSTEM_PROMPT,
        name="deepagents-rl-swe",
    )
    step("agent.prompt", truncate(prompt, 500), prompt_chars=len(prompt))

    invoke_cfg: dict[str, Any] = {"recursion_limit": recursion_limit}
    if callback_handler is not None:
        invoke_cfg["callbacks"] = [callback_handler]
        invoke_cfg["run_name"] = "deepagent-graph"
        invoke_cfg["tags"] = ["deepagents-rl", "graph"]
        invoke_cfg["metadata"] = {"episode_id": episode_id, "harness": "graph"}

    def _invoke_graph() -> Any:
        with span("agent.graph_invoke", recursion_limit=recursion_limit):
            call("graph.invoke", "create_deep_agent.invoke")
            return agent.invoke(
                {"messages": [{"role": "user", "content": prompt}]},
                config=invoke_cfg,
            )

    if langfuse_client is not None:
        with langfuse_client.start_as_current_observation(
            name="agent-graph",
            as_type="span",
            metadata={"episode_id": episode_id, "recursion_limit": recursion_limit},
            input={"prompt": prompt[:4000]},
        ) as graph_span:
            graph_out = _invoke_graph()
            msgs = list(graph_out.get("messages") or [])
            final_text = ""
            for m in reversed(msgs):
                if m.__class__.__name__ == "AIMessage" and not getattr(
                    m, "tool_calls", None
                ):
                    content = getattr(m, "content", "")
                    final_text = content if isinstance(content, str) else str(content)
                    break
            tools_used, tool_calls_log = _extract_tool_calls(msgs)
            # Replay tool spans under this graph span for UI granularity
            for tc in tool_calls_log:
                with langfuse_client.start_as_current_observation(
                    name=tc.get("name") or "tool",
                    as_type="tool",
                    input=tc.get("args") or {},
                    metadata={
                        "turn": tc.get("turn"),
                        "tool_call_id": tc.get("tool_call_id"),
                        "episode_id": episode_id,
                        "source": "graph_replay",
                    },
                ) as tool_span:
                    tool_span.update(
                        output={"result": truncate(tc.get("result") or "", 4000)}
                    )
            if tools_used:
                termination = (
                    "agent_done"
                    if "DONE" in (final_text or "").upper()
                    else "max_steps_or_halt"
                )
            else:
                termination = "no_tools"
            graph_span.update(
                output={
                    "termination": termination,
                    "n_messages": len(msgs),
                    "tools_used": tools_used,
                    "final_text": truncate(final_text, 1000),
                }
            )
    else:
        graph_out = _invoke_graph()
        msgs = list(graph_out.get("messages") or [])
        final_text = ""
        for m in reversed(msgs):
            if m.__class__.__name__ == "AIMessage" and not getattr(m, "tool_calls", None):
                content = getattr(m, "content", "")
                final_text = content if isinstance(content, str) else str(content)
                break
        tools_used, tool_calls_log = _extract_tool_calls(msgs)
        if tools_used:
            termination = (
                "agent_done" if "DONE" in (final_text or "").upper() else "max_steps_or_halt"
            )
        else:
            termination = "no_tools"

    for tc in tool_calls_log:
        result(
            "tool",
            tc.get("name") or "?",
            turn=tc.get("turn"),
            args=truncate(tc.get("args") or {}, 400),
            output=truncate(tc.get("result") or "", 400),
        )
    step(
        "agent.graph_done",
        termination,
        n_messages=len(msgs),
        tools_used=tools_used,
        final_text=truncate(final_text, 400),
    )
    return msgs, final_text, termination


def run_deepagent_episode(
    task: dict[str, Any],
    *,
    model: Any | None = None,
    recursion_limit: int = 80,
    max_turns: int = 16,
    mode: str | None = None,
    langfuse_client: Any | None = None,
    artifacts_dir: Path | None = None,
    keep_sandbox: bool = False,
) -> dict[str, Any]:
    """
    Full-repo git checkout -> deepagents LocalShellBackend + agent loop ->
    git diff model_patch -> official SWE-bench Docker eval -> trajectory.
    """
    from deepagents.backends import LocalShellBackend
    from langchain_openai import ChatOpenAI

    episode_id = str(uuid.uuid4())
    started = time.time()
    if not task.get("image") or not task.get("eval_script"):
        raise ValueError(
            f"Task {task.get('task_id')} missing image/eval_script — "
            "reload from SWE-bench/SWE-bench_Lite (not the old sparse dataset)."
        )

    step(
        "episode.begin",
        task["task_id"],
        episode_id=episode_id,
        repo=task.get("repo"),
        base_commit=task.get("base_commit"),
        image=task.get("image"),
        mode=mode or os.getenv("HARNESS_MODE") or "tool_loop",
        n_f2p=len(task.get("fail_to_pass") or []),
        n_p2p=len(task.get("pass_to_pass") or []),
    )

    sandbox_parent = Path(tempfile.mkdtemp(prefix=f"darl_swe_{task['task_id'][:40]}_"))
    sandbox = sandbox_parent / "repo"
    materialize_repo_at_commit(
        repo=task["repo"],
        base_commit=task["base_commit"],
        dest=sandbox,
    )

    mode = (mode or os.getenv("HARNESS_MODE") or "tool_loop").strip().lower()
    # tool_loop: langchain bind_tools over LocalShellBackend (reliable on local GGUF)
    # graph: create_deep_agent CompiledStateGraph

    if model is None:
        step(
            "episode.model",
            "ChatOpenAI",
            base_url=os.environ.get("MODEL_BASE_URL"),
            model=os.environ.get("MODEL_NAME"),
        )
        model = ChatOpenAI(
            base_url=os.environ["MODEL_BASE_URL"],
            api_key=os.getenv("MODEL_API_KEY", "local"),
            model=os.environ["MODEL_NAME"],
            temperature=0,
            max_tokens=1024,
        )

    step("episode.backend", f"LocalShellBackend root={sandbox}")
    backend = LocalShellBackend(root_dir=sandbox, virtual_mode=True, timeout=60)
    prompt = _short_prompt(task)
    eval_log_dir = (artifacts_dir / "evals" / episode_id) if artifacts_dir else None

    messages_dump: list[dict[str, Any]] = [{"role": "user", "content": prompt}]
    final_text = ""
    termination = "unknown"
    tools_used: list[str] = []
    tool_calls_log: list[dict[str, Any]] = []
    trace_id = ""
    error: str | None = None
    msgs: list[Any] = []

    callback_handler = None
    if langfuse_client is not None:
        from deepagents_rl.observability.langfuse_trace import make_callback_handler

        callback_handler = make_callback_handler()
        step("episode.langfuse_handler", "CallbackHandler ready for LLM generations")

    def _invoke() -> None:
        nonlocal final_text, termination, tools_used, tool_calls_log, messages_dump, msgs
        step("episode.agent_loop", f"mode={mode}")
        if mode == "graph":
            msgs, final_text, termination = _run_graph(
                model=model,
                backend=backend,
                prompt=prompt,
                recursion_limit=recursion_limit,
                langfuse_client=langfuse_client,
                callback_handler=callback_handler,
                episode_id=episode_id,
            )
        else:
            msgs, final_text, termination = _run_tool_loop(
                model=model,
                backend=backend,
                prompt=prompt,
                max_turns=max_turns,
                langfuse_client=langfuse_client,
                callback_handler=callback_handler,
                episode_id=episode_id,
            )
        messages_dump = [_msg_to_dict(m) for m in msgs]
        tools_used, tool_calls_log = _extract_tool_calls(msgs)
        step(
            "episode.agent_done",
            termination,
            n_messages=len(messages_dump),
            n_tool_calls=len(tool_calls_log),
            tools_used=tools_used,
        )

    def _grade() -> dict[str, Any]:
        from deepagents_rl.harness.repo import git_diff_patch

        step("episode.grade", "starting Docker SWE-bench eval")
        if langfuse_client is not None:
            with langfuse_client.start_as_current_observation(
                name="reward_grader",
                as_type="evaluator",
                input={
                    "fail_to_pass": task.get("fail_to_pass"),
                    "pass_to_pass": task.get("pass_to_pass"),
                    "grader": "swebench_docker_official_protocol",
                    "image": task.get("image"),
                    "n_tool_calls": len(tool_calls_log),
                    "tools_used": tools_used,
                },
                metadata={
                    "episode_id": episode_id,
                    "task_id": task["task_id"],
                },
            ) as reward_span:
                with langfuse_client.start_as_current_observation(
                    name="git_diff_model_patch",
                    as_type="span",
                    metadata={"episode_id": episode_id},
                ) as diff_span:
                    model_patch = git_diff_patch(sandbox)
                    diff_span.update(
                        output={
                            "patch_bytes": len(model_patch.encode("utf-8")),
                            "patch_preview": truncate(model_patch, 1500),
                            "empty": not bool(model_patch.strip()),
                        }
                    )
                with langfuse_client.start_as_current_observation(
                    name="docker_swe_eval",
                    as_type="span",
                    input={
                        "image": task.get("image"),
                        "patch_bytes": len(model_patch.encode("utf-8")),
                    },
                    metadata={"episode_id": episode_id},
                ) as docker_span:
                    g = grade_episode(
                        sandbox=sandbox,
                        task=task,
                        tools_used=tools_used,
                        tool_calls=tool_calls_log,
                        model_patch=model_patch,
                        eval_log_dir=eval_log_dir,
                    )
                    docker_span.update(
                        output={
                            "resolved": g.get("resolved"),
                            "resolution": g.get("resolution"),
                            "pass_frac": g.get("pass_frac"),
                            "patch_successfully_applied": g.get(
                                "patch_successfully_applied"
                            ),
                            "fail_to_pass": g.get("fail_to_pass"),
                            "pass_to_pass": {
                                "frac": (g.get("pass_to_pass") or {}).get("frac"),
                                "n_failure": len(
                                    (g.get("pass_to_pass") or {}).get("failure") or []
                                ),
                            },
                            "status_map_sample": list(
                                (g.get("status_map") or {}).items()
                            )[:15],
                            "eval_error": g.get("eval_error"),
                            "eval_exit_code": g.get("eval_exit_code"),
                            "eval_log_dir": g.get("eval_log_dir"),
                        }
                    )
                reward_span.update(
                    output={
                        k: g.get(k)
                        for k in (
                            "pass_frac",
                            "reward_total",
                            "resolved",
                            "resolution",
                            "patch_successfully_applied",
                            "fail_to_pass",
                            "pass_to_pass",
                            "eval_error",
                        )
                    }
                )
        else:
            g = grade_episode(
                sandbox=sandbox,
                task=task,
                tools_used=tools_used,
                tool_calls=tool_calls_log,
                eval_log_dir=eval_log_dir,
            )
        step(
            "episode.grade_done",
            f"resolved={g.get('resolved')} f2p={g.get('pass_frac')}",
            resolution=g.get("resolution"),
            eval_error=g.get("eval_error"),
        )
        return g

    try:
        if langfuse_client is not None:
            with langfuse_client.start_as_current_observation(
                name="deepagent-swe-episode",
                as_type="agent",
                metadata={
                    "episode_id": episode_id,
                    "task_id": task["task_id"],
                    "repo": task.get("repo"),
                    "split": task.get("split"),
                    "source_dataset": task.get("source_dataset"),
                    "harness": f"deepagents:{mode}",
                    "model": os.getenv("MODEL_NAME"),
                    "tags": ["deepagents-rl", "swe-bench", mode],
                },
                input={
                    "task_id": task["task_id"],
                    "prompt": prompt[:4000],
                    "workspace_paths": task.get("workspace_paths"),
                    "mode": mode,
                    "fail_to_pass": task.get("fail_to_pass"),
                },
            ) as root:
                # Nested agent-loop span so LLM/tool children group cleanly
                with langfuse_client.start_as_current_observation(
                    name="agent-loop",
                    as_type="span",
                    metadata={"mode": mode, "episode_id": episode_id},
                    input={"max_turns": max_turns, "recursion_limit": recursion_limit},
                ) as loop_span:
                    _invoke()
                    loop_span.update(
                        output={
                            "termination": termination,
                            "n_tool_calls": len(tool_calls_log),
                            "tools_used": tools_used,
                            "final_text": truncate(final_text, 1000),
                        }
                    )

                grade = _grade()

                langfuse_client.score_current_trace(
                    name="pass_frac", value=float(grade["pass_frac"])
                )
                langfuse_client.score_current_trace(
                    name="reward_total", value=float(grade["reward_total"])
                )
                langfuse_client.score_current_trace(
                    name="resolved", value=1.0 if grade.get("resolved") else 0.0
                )
                langfuse_client.score_current_trace(
                    name="n_tool_calls", value=float(len(tool_calls_log))
                )
                root.update(
                    output={
                        "termination": termination,
                        "final_text": final_text[:2000],
                        "tools_used": tools_used,
                        "n_tool_calls": len(tool_calls_log),
                        "tool_calls": [
                            {
                                "turn": t.get("turn"),
                                "name": t.get("name"),
                                "args": t.get("args"),
                                "result": truncate(t.get("result") or "", 500),
                            }
                            for t in tool_calls_log[:40]
                        ],
                        "grade": {
                            k: grade.get(k)
                            for k in (
                                "pass_frac",
                                "reward_total",
                                "resolved",
                                "resolution",
                                "eval_error",
                            )
                        },
                    }
                )
                trace_id = langfuse_client.get_current_trace_id() or ""
            langfuse_client.flush()
            step("episode.langfuse_flush", trace_id or "(no trace id)")
        else:
            _invoke()
            grade = _grade()
    except Exception as e:  # noqa: BLE001
        error = f"{type(e).__name__}: {e}"
        log_error("episode", error)
        termination = "error"
        grade = {
            "pass_frac": 0.0,
            "reward_total": 0.0,
            "resolved": False,
            "resolution": "NO",
            "eval_error": error,
            "grader": "swebench_docker_official_protocol",
        }

    duration_s = time.time() - started
    trajectory = {
        "episode_id": episode_id,
        "task_id": task["task_id"],
        "instance_id": task.get("instance_id"),
        "split": task.get("split"),
        "repo": task.get("repo"),
        "base_commit": task.get("base_commit"),
        "source_dataset": task.get("source_dataset") or "SWE-bench/SWE-bench_Lite",
        "harness": f"deepagents:{mode}",
        "model_name": os.getenv("MODEL_NAME"),
        "model_base_url": os.getenv("MODEL_BASE_URL"),
        "created_at": _utc_now(),
        "duration_s": duration_s,
        "recursion_limit": recursion_limit,
        "max_turns": max_turns,
        "sandbox": str(sandbox),
        "prompt": prompt,
        "messages": messages_dump,
        "tool_calls": tool_calls_log,
        "tools_used": tools_used,
        "final_text": final_text,
        "termination": termination,
        "grade": grade,
        "reward": {
            "pass_frac": grade.get("pass_frac", 0.0),
            "reward_total": grade.get("reward_total", 0.0),
            "resolved": grade.get("resolved", False),
            "resolution": grade.get("resolution"),
        },
        "model_patch": grade.get("model_patch"),
        "langfuse_trace_id": trace_id,
        "error": error,
        "gold_paths": task.get("workspace_paths"),
        "image": task.get("image"),
    }

    if artifacts_dir is not None:
        artifacts_dir.mkdir(parents=True, exist_ok=True)
        traj_path = artifacts_dir / f"episode_{episode_id}.json"
        traj_path.write_text(json.dumps(trajectory, indent=2), encoding="utf-8")
        jsonl_path = artifacts_dir / "trajectories.jsonl"
        with jsonl_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(trajectory) + "\n")
        if grade.get("model_patch"):
            (artifacts_dir / f"patch_{episode_id}.diff").write_text(
                grade["model_patch"], encoding="utf-8"
            )
        trajectory["artifact_json"] = str(traj_path)
        trajectory["artifact_jsonl"] = str(jsonl_path)

    if not keep_sandbox:
        try:
            import shutil

            step("episode.cleanup_sandbox", str(sandbox_parent))
            shutil.rmtree(sandbox_parent, ignore_errors=True)
            trajectory["sandbox"] = None
        except Exception:
            pass

    step(
        "episode.end",
        termination,
        episode_id=episode_id,
        duration_s=round(duration_s, 2),
        n_tool_calls=len(tool_calls_log),
        resolved=grade.get("resolved"),
        pass_frac=grade.get("pass_frac"),
        artifact=trajectory.get("artifact_json"),
    )
    return trajectory