"""Pydantic schemas for trace data."""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator


SpanType = Literal["LLM", "CHAIN", "TOOL", "FUNCTION"]
SpanStatus = Literal["SUCCESS", "ERROR", "RUNNING"]


def _ms_to_datetime(value: Any) -> Any:
    """Convert millisecond timestamp to datetime."""
    if isinstance(value, (int, float)):
        return datetime.fromtimestamp(value / 1000.0, tz=timezone.utc)
    return value


class TokenUsage(BaseModel):
    """Token usage for an LLM span."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    prompt_tokens: int = Field(alias="promptTokens")
    completion_tokens: int = Field(alias="completionTokens")
    total_tokens: int = Field(alias="totalTokens")


class TokenUsageSummary(BaseModel):
    """Aggregated token usage."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    total_prompt: int = Field(default=0, alias="totalPromptTokens")
    total_completion: int = Field(default=0, alias="totalCompletionTokens")
    total: int = Field(default=0, alias="totalTokens")
    estimated_cost: float | None = Field(default=0, alias="estimatedCost")


class SpanSummary(BaseModel):
    """Lightweight span summary without input/output."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    id: str
    trace_id: str = Field(alias="traceId")
    parent_id: str | None = Field(default=None, alias="parentId")
    name: str
    type: SpanType
    status: SpanStatus
    start_time: datetime = Field(alias="startTime")
    end_time: datetime | None = Field(default=None, alias="endTime")
    duration_ms: int | None = Field(default=None, alias="duration")
    model: str | None = None
    token_usage: TokenUsage | None = Field(default=None, alias="tokenUsage")
    has_error: bool = Field(default=False, alias="hasError")
    tags: list[str] = Field(default_factory=list)

    @model_validator(mode="before")
    @classmethod
    def convert_timestamps(cls, data: dict) -> dict:
        if not isinstance(data, dict):
            return data
        if "startTime" in data:
            data["startTime"] = _ms_to_datetime(data["startTime"])
        if "endTime" in data:
            data["endTime"] = _ms_to_datetime(data["endTime"])
        return data


class SpanDetail(SpanSummary):
    """Full span details with input/output."""

    input: Any | None = None
    output: Any | None = None
    error: dict | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class TraceSummary(BaseModel):
    """Trace summary for list views."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    id: str
    project_id: str | None = Field(default=None, alias="projectId")
    name: str
    status: SpanStatus
    start_time: datetime = Field(alias="startTime")
    end_time: datetime | None = Field(default=None, alias="endTime")
    duration_ms: int | None = Field(default=None, alias="durationMs")
    span_count: int = Field(alias="spanCount")
    error_count: int = Field(alias="errorCount")
    llm_count: int = Field(default=0, alias="llmCount")
    tool_call_count: int = Field(default=0, alias="toolCallCount")
    tags: list[str] = Field(default_factory=list)
    user_id: str | None = Field(default=None, alias="userId")
    session_id: str | None = Field(default=None, alias="sessionId")
    total_tokens: int | None = Field(default=None, alias="totalTokens")


class TraceDetail(TraceSummary):
    """Full trace with spans."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    spans: list[SpanDetail] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="before")
    @classmethod
    def unwrap_trace(cls, data: dict) -> dict:
        if isinstance(data, dict) and "trace" in data:
            return data["trace"]
        return data


class Pagination(BaseModel):
    """Pagination info."""

    model_config = ConfigDict(populate_by_name=True)
    total: int
    limit: int
    offset: int
    has_more: bool = Field(alias="hasMore")


class TracesResponse(BaseModel):
    """Response for listing traces."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    traces: list[TraceSummary]
    pagination: Pagination


class TimeSeriesPoint(BaseModel):
    """Single point in analytics time series."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    timestamp: datetime
    traces: int = Field(alias="traceCount")
    spans: int = Field(alias="spanCount")
    errors: int = Field(alias="errorCount")


class ModelMetrics(BaseModel):
    """Per-model analytics."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    model: str
    spans: int = Field(alias="spanCount")
    avg_duration_ms: float | None = Field(default=None, alias="avgDurationMs")
    total_tokens: int | None = Field(default=None, alias="totalTokens")
    estimated_cost: float | None = None

    @model_validator(mode="before")
    @classmethod
    def from_collector_format(cls, data: dict) -> dict:
        if not isinstance(data, dict):
            return data
        if "name" in data and "metrics" in data:
            metrics = data.get("metrics") or {}
            return {
                "model": data.get("name", ""),
                "spans": metrics.get("totalCalls", 0),
                "avg_duration_ms": metrics.get("avgLatencyMs"),
                "total_tokens": metrics.get("totalTokens"),
                "estimated_cost": metrics.get("estimatedCost"),
            }
        return data


class AnalyticsSummary(BaseModel):
    """Analytics summary with time series."""

    model_config = ConfigDict(populate_by_name=True, extra="ignore")
    total_traces: int
    total_spans: int
    error_rate: float
    avg_duration_ms: float
    p50_duration_ms: float
    p95_duration_ms: float
    p99_duration_ms: float
    token_usage: TokenUsageSummary = Field(default_factory=TokenUsageSummary)
    time_series: list[TimeSeriesPoint] = Field(default_factory=list)
    by_model: list[ModelMetrics] = Field(default_factory=list)

    @model_validator(mode="before")
    @classmethod
    def from_collector_format(cls, data: dict) -> dict:
        if not isinstance(data, dict):
            return data
        if "summary" in data:
            summary = data.get("summary") or {}
            return {
                "total_traces": summary.get("totalTraces", 0),
                "total_spans": summary.get("totalSpans", 0),
                "error_rate": summary.get("errorRate", 0.0),
                "avg_duration_ms": summary.get("avgDurationMs", 0.0),
                "p50_duration_ms": summary.get("p50DurationMs", 0.0),
                "p95_duration_ms": summary.get("p95DurationMs", 0.0),
                "p99_duration_ms": summary.get("p99DurationMs", 0.0),
                "token_usage": data.get("tokenUsage") or {},
                "time_series": data.get("timeSeries") or [],
                "by_model": data.get("byModel") or [],
            }
        return data
