from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from datetime import datetime
from typing import Any, Literal
from uuid import UUID

SpanType = Literal["llm", "tool", "retrieval", "chain", "agent"]


def _ensure_json_serializable(value: Any, field_name: str) -> None:
    try:
        json.dumps(value)
    except TypeError as exc:
        raise ValueError(f"{field_name} must be JSON serializable") from exc


@dataclass(slots=True)
class Run:
    run_id: UUID
    name: str
    created_at: datetime
    metadata_json: dict[str, Any] = field(default_factory=dict)

    def validate(self) -> None:
        if not self.name:
            raise ValueError("run.name must be non-empty")
        _ensure_json_serializable(self.metadata_json, "run.metadata_json")

    def to_dict(self) -> dict[str, Any]:
        self.validate()
        payload = asdict(self)
        payload["run_id"] = str(self.run_id)
        payload["created_at"] = self.created_at.isoformat()
        return payload


@dataclass(slots=True)
class Span:
    span_id: UUID
    run_id: UUID
    parent_span_id: UUID | None
    span_type: SpanType
    name: str
    start_time: datetime
    end_time: datetime
    duration_ms: int
    input_json: Any
    output_json: Any
    error: str | None
    metadata_json: dict[str, Any] = field(default_factory=dict)

    def validate(self) -> None:
        if not self.name:
            raise ValueError("span.name must be non-empty")
        if self.duration_ms < 0:
            raise ValueError("span.duration_ms must be >= 0")
        _ensure_json_serializable(self.input_json, "span.input_json")
        _ensure_json_serializable(self.output_json, "span.output_json")
        _ensure_json_serializable(self.metadata_json, "span.metadata_json")
        if self.start_time > self.end_time:
            raise ValueError("span.start_time must be <= span.end_time")

    def to_dict(self) -> dict[str, Any]:
        self.validate()
        payload = asdict(self)
        payload["span_id"] = str(self.span_id)
        payload["run_id"] = str(self.run_id)
        payload["parent_span_id"] = str(self.parent_span_id) if self.parent_span_id else None
        payload["start_time"] = self.start_time.isoformat()
        payload["end_time"] = self.end_time.isoformat()
        return payload


def build_ingest_payload(runs: list[Run], spans: list[Span]) -> dict[str, Any]:
    run_ids = {run.run_id for run in runs}
    span_ids_by_run: dict[UUID, set[UUID]] = {}

    for span in spans:
        span.validate()
        span_ids_by_run.setdefault(span.run_id, set()).add(span.span_id)

    for span in spans:
        if span.run_id not in run_ids:
            raise ValueError(f"span {span.span_id} references unknown run_id {span.run_id}")
        if span.parent_span_id and span.parent_span_id not in span_ids_by_run.get(span.run_id, set()):
            raise ValueError(
                f"span {span.span_id} has parent_span_id {span.parent_span_id} "
                f"not found in run {span.run_id}"
            )

    payload = {
        "runs": [run.to_dict() for run in runs],
        "spans": [span.to_dict() for span in spans],
    }
    _ensure_json_serializable(payload, "ingest payload")
    return payload
