"""Trace objects for verified learning and future controller training."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from uuid import uuid4

from .bytecode import Instruction


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


@dataclass(slots=True)
class ToolCallTrace:
    operation: str
    success: bool
    summary: str
    args: dict[str, Any] = field(default_factory=dict)
    stdout_tail: str = ""
    stderr_tail: str = ""
    returncode: int | None = None
    duration_ms: float = 0.0
    created_at: str = field(default_factory=utc_now)

    def to_dict(self) -> dict[str, Any]:
        return {
            "operation": self.operation,
            "success": self.success,
            "summary": self.summary,
            "args": self.args,
            "stdout_tail": self.stdout_tail,
            "stderr_tail": self.stderr_tail,
            "returncode": self.returncode,
            "duration_ms": self.duration_ms,
            "created_at": self.created_at,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ToolCallTrace":
        return cls(
            operation=str(data["operation"]),
            success=bool(data["success"]),
            summary=str(data["summary"]),
            args=dict(data.get("args") or {}),
            stdout_tail=str(data.get("stdout_tail", "")),
            stderr_tail=str(data.get("stderr_tail", "")),
            returncode=int(data["returncode"]) if data.get("returncode") is not None else None,
            duration_ms=float(data.get("duration_ms", 0.0)),
            created_at=str(data.get("created_at") or utc_now()),
        )


@dataclass(slots=True)
class StepTrace:
    step_index: int
    instruction: Instruction
    result_summary: str
    success: bool
    cost: float = 0.0
    data: dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=utc_now)

    def to_dict(self) -> dict[str, Any]:
        return {
            "step_index": self.step_index,
            "instruction": self.instruction.to_dict(),
            "result_summary": self.result_summary,
            "success": self.success,
            "cost": self.cost,
            "data": self.data,
            "created_at": self.created_at,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "StepTrace":
        return cls(
            step_index=int(data["step_index"]),
            instruction=Instruction.from_dict(data["instruction"]),
            result_summary=str(data["result_summary"]),
            success=bool(data["success"]),
            cost=float(data.get("cost", 0.0)),
            data=dict(data.get("data") or {}),
            created_at=str(data.get("created_at") or utc_now()),
        )


@dataclass(slots=True)
class Episode:
    task: str
    steps: list[StepTrace]
    outcome: str
    score: float
    tags: list[str] = field(default_factory=list)
    episode_id: str = field(default_factory=lambda: uuid4().hex)
    created_at: str = field(default_factory=utc_now)
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "episode_id": self.episode_id,
            "task": self.task,
            "steps": [step.to_dict() for step in self.steps],
            "outcome": self.outcome,
            "score": self.score,
            "tags": self.tags,
            "created_at": self.created_at,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Episode":
        return cls(
            episode_id=str(data.get("episode_id") or uuid4().hex),
            task=str(data["task"]),
            steps=[StepTrace.from_dict(item) for item in data.get("steps", [])],
            outcome=str(data.get("outcome", "unknown")),
            score=float(data.get("score", 0.0)),
            tags=list(data.get("tags") or []),
            created_at=str(data.get("created_at") or utc_now()),
            metadata=dict(data.get("metadata") or {}),
        )


class TraceRecorder:
    def __init__(self, task: str) -> None:
        self.task = task
        self.steps: list[StepTrace] = []

    def record(
        self,
        instruction: Instruction,
        result_summary: str,
        success: bool,
        cost: float = 0.0,
        data: dict[str, Any] | None = None,
    ) -> StepTrace:
        trace = StepTrace(
            step_index=len(self.steps),
            instruction=instruction,
            result_summary=result_summary,
            success=success,
            cost=cost,
            data=data or {},
        )
        self.steps.append(trace)
        return trace

    def to_episode(
        self,
        outcome: str,
        score: float,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> Episode:
        return Episode(
            task=self.task,
            steps=list(self.steps),
            outcome=outcome,
            score=score,
            tags=tags or [],
            metadata=metadata or {},
        )
