from __future__ import annotations

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


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _iso(dt: datetime | None) -> str | None:
    if dt is None:
        return None
    s = dt.isoformat()
    # Normalize +00:00 to Z for consistency
    if s.endswith("+00:00"):
        s = s[:-6] + "Z"
    return s


@dataclass
class Span:
    name: str
    kind: str
    span_id: str = field(default_factory=lambda: str(uuid4()))
    parent_id: str | None = None
    input: Any = None
    output: Any = None
    started_at: datetime = field(default_factory=_now)
    ended_at: datetime | None = None
    status: str = "ok"
    error: str | None = None
    metadata: dict[str, Any] | None = None

    # LLM-specific fields (only serialized when set)
    model: str | None = None
    provider: str | None = None
    response_model: str | None = None
    response_id: str | None = None
    model_params: dict[str, Any] | None = None
    input_tokens: int | None = None
    output_tokens: int | None = None
    token_details: dict[str, Any] | None = None
    cost: float | None = None
    finish_reason: str | None = None
    time_to_first_token: float | None = None
    provider_data: dict[str, Any] | None = None

    _LLM_FIELDS: tuple[str, ...] = field(
        default=(
            "model", "provider", "response_model", "response_id",
            "model_params", "input_tokens", "output_tokens", "token_details",
            "cost", "finish_reason", "time_to_first_token", "provider_data",
        ),
        init=False,
        repr=False,
    )

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {
            "span_id": self.span_id,
            "parent_id": self.parent_id,
            "name": self.name,
            "kind": self.kind,
            "started_at": _iso(self.started_at),
            "status": self.status,
        }
        if self.input is not None:
            d["input"] = self.input
        if self.output is not None:
            d["output"] = self.output
        if self.ended_at is not None:
            d["ended_at"] = _iso(self.ended_at)
        if self.error is not None:
            d["error"] = self.error
        if self.metadata is not None:
            d["metadata"] = self.metadata

        for f in self._LLM_FIELDS:
            v = getattr(self, f)
            if v is not None:
                d[f] = v

        return d


@dataclass
class Trace:
    project: str
    trace_id: str = field(default_factory=lambda: str(uuid4()))
    input: Any = None
    output: Any = None
    session_id: str | None = None
    user_id: str | None = None
    metadata: dict[str, Any] | None = None
    tags: list[str] | None = None
    started_at: datetime = field(default_factory=_now)
    ended_at: datetime | None = None
    status: str = "ok"
    error: str | None = None
    spans: list[Span] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {
            "trace_id": self.trace_id,
            "project": self.project,
            "started_at": _iso(self.started_at),
            "status": self.status,
            "spans": [s.to_dict() for s in self.spans],
        }
        if self.input is not None:
            d["input"] = self.input
        if self.output is not None:
            d["output"] = self.output
        if self.session_id is not None:
            d["session_id"] = self.session_id
        if self.user_id is not None:
            d["user_id"] = self.user_id
        if self.metadata is not None:
            d["metadata"] = self.metadata
        if self.tags is not None:
            d["tags"] = self.tags
        if self.ended_at is not None:
            d["ended_at"] = _iso(self.ended_at)
        if self.error is not None:
            d["error"] = self.error
        return d


@dataclass
class Evaluation:
    trace_id: str
    human_score: int | None = None
    human_comment: str | None = None
    human_span_ids: list[str] | None = None
    human_graded_at: datetime | None = None
    human_graded_by: str | None = None
    judge_score: int | None = None
    judge_critique: str | None = None
    judge_graded_at: datetime | None = None
    created_at: datetime = field(default_factory=_now)

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {
            "trace_id": self.trace_id,
            "created_at": _iso(self.created_at),
        }
        if self.human_score is not None:
            d["human_score"] = self.human_score
        if self.human_comment is not None:
            d["human_comment"] = self.human_comment
        if self.human_span_ids is not None:
            d["human_span_ids"] = self.human_span_ids
        if self.human_graded_at is not None:
            d["human_graded_at"] = _iso(self.human_graded_at)
        if self.human_graded_by is not None:
            d["human_graded_by"] = self.human_graded_by
        if self.judge_score is not None:
            d["judge_score"] = self.judge_score
        if self.judge_critique is not None:
            d["judge_critique"] = self.judge_critique
        if self.judge_graded_at is not None:
            d["judge_graded_at"] = _iso(self.judge_graded_at)
        return d


def get_eval_score(eval_data: dict) -> int | None:
    """Get primary score from evaluation dict, preferring judge_score over human_score."""
    if eval_data.get("judge_score") is not None:
        return eval_data["judge_score"]
    return eval_data.get("human_score")
