"""
Structured per-agent I/O tracing for FormScout.
Records every agent's input/output as JSON-serializable dicts.
Used for the Sharing-is-Caring badge (publish full trace to Hub).
"""
from __future__ import annotations

import json
import time
from dataclasses import asdict, is_dataclass
from pathlib import Path
from typing import Any


class TraceRecord:
    """A single agent execution record."""

    def __init__(self, agent_name: str, input_data: Any, output_data: Any, duration_ms: float):
        self.agent_name = agent_name
        self.input_summary = self._summarize(input_data)
        self.output_summary = self._summarize(output_data)
        self.duration_ms = duration_ms
        self.timestamp = time.time()

    def _summarize(self, data: Any) -> dict:
        """Convert dataclass or dict to JSON-safe summary."""
        if is_dataclass(data) and not isinstance(data, type):
            d = asdict(data)
            # Don't serialize raw frames (numpy arrays)
            if "frames" in d:
                d["frames"] = f"[{len(d['frames'])} frames]"
            return d
        if isinstance(data, dict):
            return data
        return {"value": str(data)}

    def to_dict(self) -> dict:
        return {
            "agent": self.agent_name,
            "timestamp": self.timestamp,
            "duration_ms": self.duration_ms,
            "input": self.input_summary,
            "output": self.output_summary,
        }


class PipelineTrace:
    """Collects trace records for a full pipeline run."""

    def __init__(self):
        self.records: list[TraceRecord] = []
        self.start_time = time.time()

    def add(self, record: TraceRecord):
        self.records.append(record)

    def to_dict(self) -> dict:
        return {
            "total_duration_ms": (time.time() - self.start_time) * 1000,
            "n_agents": len(self.records),
            "agents": [r.to_dict() for r in self.records],
        }

    def save(self, path: str | Path):
        """Save trace as JSON."""
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        with open(p, "w") as f:
            json.dump(self.to_dict(), f, indent=2, default=str)