"""Trajectory collection — the seed of opt-in model-level evolution.

Records (prompt, response, outcome) trajectories from agent runs and exports them as
SFT or DPO datasets. Actual fine-tuning (LoRA/SFT/DPO) is intentionally **opt-in and
external** — Chimera collects the data and gets out of the way; nothing here trains a
model or changes weights automatically.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, Field

Outcome = Literal["success", "failure", "unknown"]


class Trajectory(BaseModel):
    seq: int
    prompt: str
    response: str
    outcome: Outcome = "unknown"
    reward: float = 0.0
    steps: int = 0  # tool-calling steps — a long-horizon signal for data curation
    events: list[dict[str, Any]] = Field(default_factory=list)  # per-tool-step {tool, ok}
    # Diff-gate (nanobot "Dream"): did the run actually change the workspace? None = untracked
    # (no guard/workspace), True/False = certified by the real snapshot diff. ``diff_summary`` is
    # machine-derived (files +/-), so it — not the model's narrative — is the honest audit record.
    diff_productive: bool | None = None
    diff_summary: str | None = None

    def process_score(self) -> float:
        """Step-following process signal in [0, 1] — the SkillCoach data-quality filter."""
        from chimera.ecosystem.events import step_following_score

        return step_following_score(self.events)


class TrajectoryCollector:
    """A JSONL-backed append-only log of trajectories, exportable for training."""

    def __init__(self, path: Path) -> None:
        self.path = Path(path)
        self._items: list[Trajectory] = []
        self.load()

    def load(self) -> None:
        self._items = []
        if not self.path.exists():
            return
        for line in self.path.read_text(encoding="utf-8").splitlines():
            if line.strip():
                self._items.append(Trajectory.model_validate_json(line))

    def record(
        self,
        prompt: str,
        response: str,
        *,
        outcome: Outcome = "unknown",
        reward: float = 0.0,
        steps: int = 0,
        events: list[dict[str, Any]] | None = None,
        diff_productive: bool | None = None,
        diff_summary: str | None = None,
    ) -> Trajectory:
        item = Trajectory(
            seq=len(self._items),
            prompt=prompt,
            response=response,
            outcome=outcome,
            reward=reward,
            steps=steps,
            events=list(events or []),
            diff_productive=diff_productive,
            diff_summary=diff_summary,
        )
        self._items.append(item)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with self.path.open("a", encoding="utf-8") as handle:
            handle.write(item.model_dump_json() + "\n")
        return item

    def all(self) -> list[Trajectory]:
        return list(self._items)

    def __len__(self) -> int:
        return len(self._items)

    def export_sft(self, out_path: Path) -> int:
        """Write successful trajectories as chat-format SFT examples. Returns count."""
        rows = [
            {
                "messages": [
                    {"role": "user", "content": item.prompt},
                    {"role": "assistant", "content": item.response},
                ]
            }
            for item in self._items
            if item.outcome == "success"
        ]
        _write_jsonl(out_path, rows)
        return len(rows)

    def export_dpo(self, out_path: Path) -> int:
        """Write preference pairs (chosen=success, rejected=failure) per prompt."""
        by_prompt: dict[str, dict[str, str]] = {}
        for item in self._items:
            slot = by_prompt.setdefault(item.prompt, {})
            if item.outcome == "success":
                slot.setdefault("chosen", item.response)
            elif item.outcome == "failure":
                slot.setdefault("rejected", item.response)
        rows = [
            {"prompt": prompt, "chosen": slot["chosen"], "rejected": slot["rejected"]}
            for prompt, slot in by_prompt.items()
            if "chosen" in slot and "rejected" in slot
        ]
        _write_jsonl(out_path, rows)
        return len(rows)


def _write_jsonl(out_path: Path, rows: list[dict[str, Any]]) -> None:
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text("\n".join(json.dumps(row) for row in rows) + ("\n" if rows else ""), encoding="utf-8")
