"""TRL GRPO `environment_factory` for deepagents SWE-bench sandboxes.

Kaggle has no Docker for instance images — `get_reward` POSTs the model patch
to `SWE_GRADER_URL` (local `scripts/swe_grader_server.py` + tunnel).

Each rollout opens a Langfuse agent trace; every tool call is a nested `tool`
span; grading is an `evaluator` span with reward scores.
"""
from __future__ import annotations

import json
import os
import shutil
import tempfile
import traceback
import uuid
from functools import wraps
from pathlib import Path
from typing import Any, Callable


def _vlog(msg: str, **fields: Any) -> None:
    """Always-on verbose progress (Kaggle logs + local stderr)."""
    extra = ""
    if fields:
        try:
            extra = " | " + json.dumps(fields, default=str)[:1200]
        except Exception:
            extra = f" | {fields!r}"[:1200]
    line = f"[darl] {msg}{extra}"
    print(line, flush=True)
    try:
        from deepagents_rl.observability.logging import step

        step("swe_env", msg, **fields)
    except Exception:
        pass


def _result_to_str(obj: Any) -> str:
    if obj is None:
        return ""
    if isinstance(obj, str):
        return obj[:8000]
    if hasattr(obj, "output") and isinstance(obj.output, str):
        return obj.output[:8000]
    # Prefer listing entries before error (LsResult / GlobResult)
    entries = getattr(obj, "entries", None)
    if entries is not None:
        lines: list[str] = []
        for e in entries:
            if isinstance(e, dict):
                lines.append(str(e.get("path") or e))
            else:
                lines.append(str(e))
        return ("\n".join(lines) if lines else "(empty directory)")[:8000]
    matches = getattr(obj, "matches", None)
    if matches is not None and not getattr(obj, "error", None):
        try:
            return json.dumps(matches, default=str)[:8000]
        except Exception:
            return str(matches)[:8000]
    if hasattr(obj, "error") and obj.error:
        return f"ERROR: {obj.error}"[:8000]
    try:
        return json.dumps(obj, default=str)[:8000]
    except Exception:
        return str(obj)[:8000]


# Common dirs Qwen invents that most SWE-bench repos do not have.
_PHANTOM_DIRS = frozenset(
    {
        "src",
        "lib",
        "app",
        "source",
        "sources",
        "code",
        "project",
        "workspace",
    }
)


def _as_str_list(val: Any) -> list[str]:
    if val is None:
        return []
    if isinstance(val, list):
        return [str(x) for x in val]
    if isinstance(val, str):
        s = val.strip()
        if not s:
            return []
        try:
            parsed = json.loads(s)
            if isinstance(parsed, list):
                return [str(x) for x in parsed]
        except json.JSONDecodeError:
            pass
        return [s]
    return [str(val)]


def _truncate(obj: Any, n: int = 4000) -> Any:
    if isinstance(obj, str):
        return obj[:n]
    if isinstance(obj, dict):
        return {k: _truncate(v, n) for k, v in list(obj.items())[:40]}
    if isinstance(obj, list):
        return [_truncate(v, n) for v in obj[:40]]
    return obj


def grade_patch_remote(
    *,
    instance_id: str,
    docker_image: str,
    eval_script: str,
    fail_to_pass: list[str],
    pass_to_pass: list[str],
    model_patch: str,
    timeout_s: int = 1800,
) -> dict[str, Any]:
    """POST patch to the local Docker grader service."""
    import requests

    base = (os.environ.get("SWE_GRADER_URL") or "").rstrip("/")
    if not base:
        raise RuntimeError("SWE_GRADER_URL is not set (point at the local grader tunnel)")

    headers = {
        "Content-Type": "application/json",
        "ngrok-skip-browser-warning": "1",
    }
    token = os.environ.get("SWE_GRADER_TOKEN") or ""
    if token:
        headers["Authorization"] = f"Bearer {token}"

    payload = {
        "instance_id": instance_id,
        "image": docker_image,
        "eval_script": eval_script,
        "fail_to_pass": fail_to_pass,
        "pass_to_pass": pass_to_pass,
        "model_patch": model_patch or "",
        "timeout_s": timeout_s,
    }
    _vlog(
        "grader.POST /grade",
        url=base,
        instance_id=instance_id,
        patch_bytes=len((model_patch or "").encode("utf-8")),
        timeout_s=timeout_s,
    )
    resp = requests.post(f"{base}/grade", headers=headers, json=payload, timeout=timeout_s + 60)
    _vlog("grader.response", status=resp.status_code, body_preview=(resp.text or "")[:400])
    resp.raise_for_status()
    data = resp.json()
    _vlog(
        "grader.result",
        instance_id=instance_id,
        reward_total=data.get("reward_total"),
        resolution=data.get("resolution"),
        resolved=data.get("resolved"),
        error=data.get("error"),
    )
    return data


def _traced_tool(method: Callable[..., str]) -> Callable[..., str]:
    """Wrap a public tool method so each call becomes a Langfuse tool span."""

    @wraps(method)
    def wrapper(self: "DeepAgentSWEEnv", *args: Any, **kwargs: Any) -> str:
        return self._call_tool(method.__name__, method, args, kwargs)

    return wrapper


class DeepAgentSWEEnv:
    """Per-rollout TRL environment: LocalShellBackend tools + remote Docker reward."""

    def __init__(self) -> None:
        self._sandbox_parent: Path | None = None
        self._sandbox: Path | None = None
        self._backend: Any = None
        self._task: dict[str, Any] = {}
        self._last_reward: float = 0.0
        self._last_report: dict[str, Any] = {}
        self._rollout_id: str = ""
        self._tool_calls: list[dict[str, Any]] = []
        self._lf: Any = None
        self._root_cm: Any = None
        self._root: Any = None
        self._trace_id: str = ""
        self._root_span_id: str = ""

    def _langfuse_enabled(self) -> bool:
        if os.getenv("DARL_LANGFUSE", "1").strip() in ("0", "false", "no"):
            return False
        return bool(os.getenv("LANGFUSE_PUBLIC_KEY") and os.getenv("LANGFUSE_SECRET_KEY"))

    def _ensure_langfuse(self) -> Any | None:
        if not self._langfuse_enabled():
            return None
        if self._lf is not None:
            return self._lf
        try:
            from langfuse import Langfuse, get_client

            Langfuse(
                public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
                secret_key=os.environ["LANGFUSE_SECRET_KEY"],
                host=os.getenv("LANGFUSE_HOST") or os.getenv("LANGFUSE_BASE_URL"),
            )
            self._lf = get_client()
            return self._lf
        except Exception as e:  # noqa: BLE001
            print(f"[langfuse] init failed: {e}")
            return None

    def _child_trace_context(self) -> dict[str, str] | None:
        """Attach children to THIS rollout's root — never ambient OTEL context.

        Passing only ``trace_id`` (no parent) makes the Langfuse SDK invent a
        NonRecordingSpan parent id → orphan parents + mega-trace pollution when
        parallel GRPO rollouts share get_client().
        """
        if not self._trace_id or not self._root_span_id:
            return None
        return {"trace_id": self._trace_id, "parent_span_id": self._root_span_id}

    def _end_trace(self) -> None:
        if self._root_cm is not None:
            try:
                self._root_cm.__exit__(None, None, None)
            except Exception:
                pass
        self._root_cm = None
        self._root = None
        self._root_span_id = ""
        # Keep _trace_id until next reset for logging; clear after flush.
        if self._lf is not None:
            try:
                self._lf.flush()
            except Exception:
                pass

    def _call_tool(
        self,
        name: str,
        method: Callable[..., str],
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> str:
        payload = _truncate({"args": list(args), "kwargs": kwargs})
        n = len(self._tool_calls) + 1
        _vlog(
            f"tool[{n}].begin {name}",
            instance_id=self._task.get("instance_id"),
            rollout_id=self._rollout_id,
            input=payload,
        )
        lf = self._lf
        ctx = self._child_trace_context()
        if lf is None or ctx is None:
            try:
                out = method(self, *args, **kwargs)
            except Exception as e:  # noqa: BLE001
                out = f"TOOL_ERROR: {type(e).__name__}: {e}"
            self._tool_calls.append({"name": name, "input": payload, "output": _truncate(out)})
            _vlog(
                f"tool[{n}].end {name}",
                output_preview=str(out)[:500],
                output_chars=len(str(out)),
            )
            return out

        with lf.start_as_current_observation(
            name=name,
            as_type="tool",
            trace_context=ctx,
            input=payload,
            metadata={
                "rollout_id": self._rollout_id,
                "instance_id": self._task.get("instance_id"),
                "harness": "deepagents-swe-grpo",
                "tool_index": n,
                "trace_id": self._trace_id,
            },
        ) as span:
            try:
                out = method(self, *args, **kwargs)
                span.update(output=_truncate(out))
            except Exception as e:  # noqa: BLE001
                out = f"TOOL_ERROR: {type(e).__name__}: {e}"
                span.update(output=out, level="ERROR")
            self._tool_calls.append({"name": name, "input": payload, "output": _truncate(out)})
            _vlog(
                f"tool[{n}].end {name}",
                output_preview=str(out)[:500],
                output_chars=len(str(out)),
            )
            return out

    def reset(self, **kwargs: Any) -> str:
        """Materialize repo sandbox; return tool-forcing text TRL appends to the user message."""
        self._cleanup()
        self._end_trace()

        instance_id = str(kwargs.get("instance_id") or kwargs.get("task_id") or "")
        repo = str(kwargs.get("repo") or "")
        base_commit = str(kwargs.get("base_commit") or "")
        docker_image = str(kwargs.get("docker_image") or kwargs.get("image") or "")
        eval_script = str(kwargs.get("eval_script") or "")
        fail_to_pass = _as_str_list(
            kwargs.get("fail_to_pass")
            if kwargs.get("fail_to_pass") is not None
            else kwargs.get("fail_to_pass_json")
        )
        pass_to_pass = _as_str_list(
            kwargs.get("pass_to_pass")
            if kwargs.get("pass_to_pass") is not None
            else kwargs.get("pass_to_pass_json")
        )

        _vlog(
            "env.reset.begin",
            instance_id=instance_id,
            repo=repo,
            base_commit=(base_commit[:12] if base_commit else ""),
            keys=list(kwargs.keys())[:30],
        )

        if not instance_id or not repo or not base_commit:
            raise ValueError(
                f"DeepAgentSWEEnv.reset requires instance_id/repo/base_commit; got keys={list(kwargs)}"
            )
        if not docker_image or not eval_script.strip():
            raise ValueError(f"{instance_id}: missing docker_image/eval_script")

        from deepagents.backends import LocalShellBackend

        from deepagents_rl.harness.repo import materialize_repo_at_commit

        self._rollout_id = str(uuid.uuid4())
        self._tool_calls = []
        self._task = {
            "instance_id": instance_id,
            "task_id": instance_id,
            "repo": repo,
            "base_commit": base_commit,
            "image": docker_image,
            "docker_image": docker_image,
            "eval_script": eval_script,
            "fail_to_pass": fail_to_pass,
            "pass_to_pass": pass_to_pass,
        }
        self._last_reward = 0.0
        self._last_report = {}
        self._trace_id = ""

        lf = self._ensure_langfuse()
        _vlog("env.reset.langfuse", enabled=lf is not None)
        if lf is not None:
            try:
                # IMPORTANT: do NOT pass trace_context={"trace_id": ...} alone.
                # Langfuse SDK then fabricates a NonRecordingSpan parent → orphan
                # parent ids in the UI and children from other rollouts can leak in.
                # A plain start_as_current_observation creates a clean root trace.
                self._root_cm = lf.start_as_current_observation(
                    name=f"grpo-swe-rollout:{instance_id}",
                    as_type="agent",
                    input={
                        "instance_id": instance_id,
                        "repo": repo,
                        "base_commit": base_commit,
                        "rollout_id": self._rollout_id,
                        "prompt_preview": _truncate(
                            kwargs.get("prompt")
                            if isinstance(kwargs.get("prompt"), str)
                            else (
                                kwargs.get("problem_statement")
                                or (kwargs.get("prompt") or "")
                            ),
                            2000,
                        ),
                    },
                    metadata={
                        "rollout_id": self._rollout_id,
                        "instance_id": instance_id,
                        "docker_image": docker_image,
                        "phase": "grpo_train_or_eval",
                        "harness": "deepagents-swe-grpo",
                        "session_id": self._rollout_id,
                        "tags": ["deepagents-rl", "swe-bench", "grpo", "tools"],
                    },
                    end_on_exit=True,
                )
                self._root = self._root_cm.__enter__()
                self._trace_id = str(
                    getattr(self._root, "trace_id", None)
                    or getattr(self._root, "traceId", None)
                    or ""
                )
                self._root_span_id = str(
                    getattr(self._root, "id", None)
                    or getattr(self._root, "observation_id", None)
                    or ""
                )
                _vlog(
                    "env.reset.langfuse_trace",
                    trace_id=self._trace_id,
                    root_span_id=self._root_span_id,
                )
            except Exception as e:  # noqa: BLE001
                _vlog("env.reset.langfuse_failed", error=str(e))
                self._root_cm = None
                self._root = None
                self._trace_id = ""
                self._root_span_id = ""

        def _materialize() -> None:
            _vlog("env.reset.materialize", repo=repo, base_commit=base_commit[:12])
            self._sandbox_parent = Path(
                tempfile.mkdtemp(prefix=f"darl_grpo_{instance_id[:40]}_")
            )
            self._sandbox = self._sandbox_parent / "repo"
            materialize_repo_at_commit(repo=repo, base_commit=base_commit, dest=self._sandbox)
            self._backend = LocalShellBackend(
                root_dir=str(self._sandbox),
                virtual_mode=True,
                timeout=60,
            )
            _vlog("env.reset.sandbox_ready", sandbox=str(self._sandbox))

        ctx = self._child_trace_context()
        if lf is not None and ctx is not None:
            with lf.start_as_current_observation(
                name="repo_materialize",
                as_type="span",
                trace_context=ctx,
                input={"repo": repo, "base_commit": base_commit},
            ) as span:
                _materialize()
                span.update(output={"sandbox": str(self._sandbox)})
        else:
            _materialize()
        # Pre-list root so the model sees real paths even before its first tool call.
        root_listing = ""
        try:
            root_listing = self._root_listing()
        except Exception as e:  # noqa: BLE001
            root_listing = f"(root listing failed: {e})"
        _vlog(
            "env.reset.done",
            instance_id=instance_id,
            rollout_id=self._rollout_id,
            n_f2p=len(fail_to_pass),
            n_p2p=len(pass_to_pass),
            langfuse_trace_id=self._trace_id,
            root_listing_chars=len(root_listing),
        )
        # TRL appends this string to the last user message — force tool use.
        return (
            "\n\nREPO ROOT LISTING (already gathered for you):\n"
            f"{root_listing[:3500]}\n\n"
            "CRITICAL REMINDER:\n"
            "- You MUST use structured tool/function calls (not markdown ```json```).\n"
            "- FIRST tool call: ls path='.' OR grep/glob/read_file on a REAL path from the listing above.\n"
            "- Never invent src/lib/app. Prefer edit_file with a unique old_string.\n"
            "- Text-only DONE with zero tools scores NEGATIVE reward. Call tools.\n"
            "- After a minimal fix, stop tools and reply: DONE <summary>."
        )

    def _cleanup(self) -> None:
        self._backend = None
        if self._sandbox_parent is not None and self._sandbox_parent.exists():
            shutil.rmtree(self._sandbox_parent, ignore_errors=True)
        self._sandbox_parent = None
        self._sandbox = None

    def _require_backend(self) -> Any:
        if self._backend is None:
            raise RuntimeError("DeepAgentSWEEnv tool called before reset()")
        return self._backend

    @staticmethod
    def _norm_path(path: str | None) -> str:
        p = (path or ".").strip().strip("'\"")
        if p in ("", "/", "./", ".\\", "~"):
            return "."
        # virtual_mode often wants leading slash for absolute-in-sandbox paths
        if p.startswith("./"):
            p = p[2:]
        if p.startswith("/"):
            p = p.lstrip("/") or "."
        # Drop trailing slash for alias matching
        p = p.rstrip("/") or "."
        return p or "."

    @classmethod
    def _resolve_ls_path(cls, path: str | None) -> tuple[str, str | None]:
        """Map invented dirs (src/lib/…) to repo root before hitting the backend."""
        p = cls._norm_path(path)
        head = p.split("/", 1)[0].lower()
        if head in _PHANTOM_DIRS:
            return ".", p
        return p, None

    def _is_missing(self, result_obj: Any, text: str) -> bool:
        err = getattr(result_obj, "error", None)
        if err and "path_not_found" in str(err).lower():
            return True
        return "path_not_found" in (text or "").lower()

    def _root_listing(self) -> str:
        backend = self._require_backend()
        # Prefer "/" then "." for virtual_mode backends
        for candidate in ("/", "."):
            result = backend.ls(candidate)
            out = _result_to_str(result)
            if not self._is_missing(result, out):
                return out
        return "(unable to list repo root)"

    @_traced_tool
    def ls(self, path: str = ".") -> str:
        """List files in a directory inside the sandbox repo root.

        Args:
            path: Directory relative to the repo root. Use "." for the root.
                Do NOT assume a "src/" folder exists — many repos have none.

        Returns:
            Directory listing as text (falls back to repo root if path missing).
        """
        backend = self._require_backend()
        path, aliased_from = self._resolve_ls_path(path)
        result = backend.ls(path)
        out = _result_to_str(result)
        if self._is_missing(result, out):
            root = self._root_listing()
            _vlog("tool.ls.fallback_root", requested=aliased_from or path)
            return (
                f"Path '{aliased_from or path}' does not exist here.\n"
                f"Listing repo root instead (use these real paths):\n{root}"
            )
        if aliased_from and aliased_from != path:
            _vlog("tool.ls.alias", requested=aliased_from, used=path)
            return (
                f"Note: '{aliased_from}' is not used in this repo; showing repo root:\n{out}"
            )
        return out

    @_traced_tool
    def glob(self, pattern: str, path: str = ".") -> str:
        """Glob for files matching pattern under path.

        Args:
            pattern: Glob pattern (e.g. '**/*.py').
            path: Root path for the search (use '.' for repo root).

        Returns:
            Matching paths as text.
        """
        backend = self._require_backend()
        path, aliased_from = self._resolve_ls_path(path)
        result = backend.glob(pattern, path)
        out = _result_to_str(result)
        if self._is_missing(result, out):
            return (
                f"Path '{aliased_from or path}' not found for glob; retrying at repo root.\n"
                f"{_result_to_str(backend.glob(pattern, '.'))}"
            )
        return out

    @_traced_tool
    def grep(self, pattern: str, path: str = ".") -> str:
        """Search file contents for a regex pattern.

        Args:
            pattern: Regex pattern to search for.
            path: Root path for the search (use '.' for repo root).

        Returns:
            Matching lines as text.
        """
        backend = self._require_backend()
        path, aliased_from = self._resolve_ls_path(path)
        result = backend.grep(pattern, path)
        out = _result_to_str(result)
        if self._is_missing(result, out):
            return (
                f"Path '{aliased_from or path}' not found for grep; retrying at repo root.\n"
                f"{_result_to_str(backend.grep(pattern, '.'))}"
            )
        return out

    @_traced_tool
    def read_file(self, file_path: str, offset: int = 0, limit: int = 200) -> str:
        """Read a text file from the sandbox.

        Args:
            file_path: Path relative to the repo root (discover via ls/glob first).
            offset: Starting line offset.
            limit: Max number of lines to return.

        Returns:
            File contents as text.
        """
        backend = self._require_backend()
        file_path = self._norm_path(file_path)
        # If model invents src/foo.py, try stripping the phantom prefix once.
        head = file_path.split("/", 1)[0].lower()
        candidates = [file_path]
        if head in _PHANTOM_DIRS and "/" in file_path:
            candidates.append(file_path.split("/", 1)[1])
        last_out = ""
        for cand in candidates:
            result = backend.read(cand, offset=offset, limit=limit)
            out = _result_to_str(result)
            last_out = out
            if not self._is_missing(result, out):
                if cand != file_path:
                    return f"(read via '{cand}' — '{file_path}' missing)\n{out}"
                return out
        return (
            f"File '{file_path}' not found. Use ls path='.' or glob to discover real paths.\n"
            f"Repo root listing:\n{self._root_listing()}\nLast error: {last_out}"
        )

    @_traced_tool
    def write_file(self, file_path: str, content: str) -> str:
        """Write or create a text file in the sandbox.

        Args:
            file_path: Path relative to the repo root.
            content: Full file contents to write.

        Returns:
            Write result as text.
        """
        return _result_to_str(self._require_backend().write(file_path, content))

    @_traced_tool
    def edit_file(self, file_path: str, old_string: str, new_string: str) -> str:
        """Replace old_string with new_string in an existing file.

        Args:
            file_path: Path relative to the repo root.
            old_string: Exact text to find.
            new_string: Replacement text.

        Returns:
            Edit result as text.
        """
        return _result_to_str(
            self._require_backend().edit(file_path, old_string, new_string)
        )

    @_traced_tool
    def execute(self, command: str) -> str:
        """Run a shell command in the sandbox.

        Args:
            command: Shell command (e.g. python -c '...').

        Returns:
            Command stdout/stderr as text.
        """
        return _result_to_str(self._require_backend().execute(command))

    def _shaped_reward(
        self,
        *,
        docker_reward: float,
        model_patch: str,
        report: dict[str, Any],
    ) -> tuple[float, dict[str, float]]:
        """Dense process reward so GRPO sees non-zero variance before resolve.

        Components (additive, then clamped):
          - no tool calls: -0.15
          - tool calls: +0.05 each (cap +0.25)
          - used edit_file/write_file: +0.15
          - non-empty git patch: +0.30
          - docker/pass reward (0..1 typically): +1.0 * docker_reward
        """
        names = [str(c.get("name") or "") for c in self._tool_calls]
        n = len(names)
        parts = {
            "tool_penalty": -0.15 if n == 0 else 0.0,
            "tool_bonus": 0.0 if n == 0 else float(min(0.25, 0.05 * n)),
            "edit_bonus": 0.15
            if any(x in names for x in ("edit_file", "write_file"))
            else 0.0,
            "patch_bonus": 0.30 if (model_patch or "").strip() else 0.0,
            "docker": float(docker_reward or 0.0),
        }
        total = sum(parts.values())
        total = max(-0.2, min(1.6, total))
        parts["total"] = total
        return total, parts

    def get_reward(self) -> float:
        """Shaped reward = process bonuses + remote Docker SWE score (if patch)."""
        lf = self._lf
        reward = 0.0
        report: dict[str, Any] = {}
        shape_parts: dict[str, float] = {}
        model_patch = ""
        _vlog(
            "env.get_reward.begin",
            instance_id=self._task.get("instance_id"),
            n_tool_calls=len(self._tool_calls),
            tools_used=[c.get("name") for c in self._tool_calls],
        )

        def _grade() -> float:
            nonlocal report, model_patch, shape_parts
            docker_reward = 0.0
            if self._sandbox is None:
                report = {"error": "no_sandbox", "reward_total": 0.0}
                _vlog("env.get_reward.no_sandbox")
                shaped, shape_parts = self._shaped_reward(
                    docker_reward=0.0, model_patch="", report=report
                )
                report["shaped"] = shape_parts
                report["reward_total"] = shaped
                return shaped

            from deepagents_rl.harness.repo import git_diff_patch

            try:
                _vlog("env.get_reward.git_diff")
                model_patch = git_diff_patch(self._sandbox)
                _vlog(
                    "env.get_reward.patch",
                    patch_bytes=len(model_patch.encode("utf-8")),
                    empty=not bool(model_patch.strip()),
                    preview=model_patch[:300],
                )
            except Exception as e:  # noqa: BLE001
                report = {"error": f"git_diff_failed: {e}"}
                _vlog("env.get_reward.git_diff_failed", error=str(e))
                shaped, shape_parts = self._shaped_reward(
                    docker_reward=0.0, model_patch="", report=report
                )
                report["shaped"] = shape_parts
                report["reward_total"] = shaped
                return shaped

            if not model_patch.strip():
                report = {"error": "empty_model_patch", "reward_total": 0.0}
                _vlog("env.get_reward.empty_patch")
            else:
                try:
                    report = grade_patch_remote(
                        instance_id=self._task["instance_id"],
                        docker_image=self._task["docker_image"],
                        eval_script=self._task["eval_script"],
                        fail_to_pass=list(self._task["fail_to_pass"]),
                        pass_to_pass=list(self._task["pass_to_pass"]),
                        model_patch=model_patch,
                    )
                    docker_reward = float(report.get("reward_total") or 0.0)
                except Exception as e:  # noqa: BLE001
                    report = {
                        "error": f"grader_failed: {e}",
                        "traceback": traceback.format_exc()[-1500:],
                    }
                    _vlog("env.get_reward.grader_failed", error=str(e))
                    docker_reward = 0.0

            shaped, shape_parts = self._shaped_reward(
                docker_reward=docker_reward,
                model_patch=model_patch,
                report=report,
            )
            report["docker_reward"] = docker_reward
            report["shaped"] = shape_parts
            report["reward_total"] = shaped
            _vlog("env.get_reward.shaped", **shape_parts)
            return shaped

        ctx = self._child_trace_context()
        if lf is not None and ctx is not None:
            with lf.start_as_current_observation(
                name="reward_grader",
                as_type="evaluator",
                trace_context=ctx,
                input={
                    "instance_id": self._task.get("instance_id"),
                    "n_tool_calls": len(self._tool_calls),
                    "tools_used": [c["name"] for c in self._tool_calls],
                },
            ) as span:
                reward = _grade()
                span.update(
                    output=_truncate(
                        {
                            "reward_total": reward,
                            "docker_reward": report.get("docker_reward"),
                            "shaped": report.get("shaped"),
                            "pass_frac": report.get("pass_frac"),
                            "resolved": report.get("resolved"),
                            "resolution": report.get("resolution"),
                            "error": report.get("error"),
                            "n_tool_calls": len(self._tool_calls),
                        }
                    )
                )
                try:
                    lf.score_current_trace(name="reward_total", value=float(reward))
                    lf.score_current_trace(
                        name="pass_frac", value=float(report.get("pass_frac") or 0.0)
                    )
                    lf.score_current_trace(
                        name="resolved", value=1.0 if report.get("resolved") else 0.0
                    )
                    lf.score_current_trace(
                        name="n_tool_calls", value=float(len(self._tool_calls))
                    )
                    if isinstance(report.get("shaped"), dict):
                        for sk, sv in report["shaped"].items():
                            lf.score_current_trace(name=f"shaped_{sk}", value=float(sv))
                except Exception:
                    pass
                if self._root is not None:
                    try:
                        self._root.update(
                            output={
                                "reward_total": reward,
                                "n_tool_calls": len(self._tool_calls),
                                "tools_used": [c["name"] for c in self._tool_calls],
                                "resolution": report.get("resolution"),
                                "shaped": report.get("shaped"),
                            }
                        )
                    except Exception:
                        pass
        else:
            reward = _grade()

        self._last_report = report
        self._last_reward = reward
        _vlog(
            "env.get_reward.done",
            instance_id=self._task.get("instance_id"),
            reward_total=reward,
            pass_frac=report.get("pass_frac"),
            resolution=report.get("resolution"),
            resolved=report.get("resolved"),
            error=report.get("error"),
            n_tool_calls=len(self._tool_calls),
            shaped=report.get("shaped"),
        )

        # W&B: one point per completed rollout (train + ckpt eval)
        try:
            import wandb

            if wandb.run is not None:
                tools_used = [c.get("name") for c in self._tool_calls]
                shaped = report.get("shaped") if isinstance(report.get("shaped"), dict) else {}
                wandb.log(
                    {
                        "rollout/reward_total": float(reward),
                        "rollout/docker_reward": float(report.get("docker_reward") or 0.0),
                        "rollout/pass_frac": float(report.get("pass_frac") or 0.0),
                        "rollout/resolved": 1.0 if report.get("resolved") else 0.0,
                        "rollout/n_tool_calls": float(len(self._tool_calls)),
                        "rollout/instance_id": self._task.get("instance_id") or "",
                        "rollout/tools_used": ",".join(tools_used),
                        "rollout/resolution": str(report.get("resolution") or ""),
                        "rollout/grader_error": str(report.get("error") or ""),
                        "rollout/shaped_tool_bonus": float(shaped.get("tool_bonus") or 0.0),
                        "rollout/shaped_edit_bonus": float(shaped.get("edit_bonus") or 0.0),
                        "rollout/shaped_patch_bonus": float(shaped.get("patch_bonus") or 0.0),
                    }
                )
                _vlog("env.get_reward.wandb_logged")
        except Exception as e:  # noqa: BLE001
            _vlog("env.get_reward.wandb_failed", error=str(e))

        self._end_trace()
        _vlog("env.get_reward.trace_closed", trace_id=self._trace_id)
        return reward


def prepare_swe_grpo_rows(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Normalize HF SWE Lite rows for TRL GRPO + DeepAgentSWEEnv.

    Renames `image` -> `docker_image` so TRL does not treat it as multimodal.
    Converts `prompt` to conversational chat messages (required for tool calling).
    """
    from deepagents_rl.data.swe_bench_lite import build_agent_prompt, from_hf_record

    rows: list[dict[str, Any]] = []
    for rec in records:
        if "fail_to_pass_json" in rec or (
            "prompt" in rec and isinstance(rec.get("prompt"), str) and "image" in rec
        ):
            task = from_hf_record(rec)
        else:
            task = dict(rec)

        # Rebuild prompt locally so tool-forcing text stays current (HF rows may be stale).
        messages = [
            {
                "role": "user",
                "content": build_agent_prompt(
                    str(task.get("problem_statement") or ""),
                    instance_id=str(task.get("instance_id") or task.get("task_id") or ""),
                ),
            }
        ]

        docker_image = str(task.get("docker_image") or task.get("image") or "")
        rows.append(
            {
                "prompt": messages,
                "task_id": task["task_id"],
                "instance_id": task["instance_id"],
                "repo": task["repo"],
                "base_commit": task["base_commit"],
                "docker_image": docker_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 []),
                "problem_statement": task.get("problem_statement") or "",
            }
        )
    return rows