"""Fluent span data builder."""
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from decimal import Decimal
from typing import Any


# Pricing table (USD per 1M tokens) — used for auto-cost calculation
_PRICING: dict[str, tuple[float, float]] = {
    # model_name: (input_per_1m, output_per_1m)
    # OpenAI
    "gpt-4o": (2.50, 10.00),
    "gpt-4o-mini": (0.15, 0.60),
    "gpt-4-turbo": (10.00, 30.00),
    "gpt-3.5-turbo": (0.50, 1.50),
    # Anthropic
    "claude-sonnet-4-20250514": (3.00, 15.00),
    "claude-3-opus-20240229": (15.00, 75.00),
    "claude-3-haiku-20240307": (0.25, 1.25),
    "claude-3-5-sonnet-20241022": (3.00, 15.00),
    # Cohere
    "command": (1.00, 2.00),
    "command-r": (0.50, 1.50),
    "command-r-plus": (3.00, 15.00),
    # Google
    "gemini-pro": (0.50, 1.50),
    "gemini-1.5-pro": (3.50, 10.50),
    "gemini-1.5-flash": (0.075, 0.30),
    # Mistral
    "mistral-small-latest": (0.20, 0.60),
    "mistral-medium-latest": (2.70, 8.10),
    "mistral-large-latest": (4.00, 12.00),
    # Groq (hosted models)
    "llama-3.1-70b-versatile": (0.59, 0.79),
    "mixtral-8x7b-32768": (0.24, 0.24),
    # Bedrock
    "anthropic.claude-3-sonnet-20240229-v1:0": (3.00, 15.00),
    "anthropic.claude-3-haiku-20240307-v1:0": (0.25, 1.25),
    "amazon.titan-text-express-v1": (0.20, 0.60),
    "amazon.titan-text-lite-v1": (0.15, 0.20),
}


def estimate_cost(model_name: str | None, tokens_input: int | None, tokens_output: int | None) -> Decimal | None:
    if not model_name or (not tokens_input and not tokens_output):
        return None
    pricing = _PRICING.get(model_name)
    if pricing is None:
        return None
    inp = (tokens_input or 0) / 1_000_000 * pricing[0]
    out = (tokens_output or 0) / 1_000_000 * pricing[1]
    return Decimal(str(round(inp + out, 6)))


@dataclass
class SpanData:
    """Mutable span data populated during execution."""

    trace_id: str = ""
    span_id: str = ""
    parent_span_id: str | None = None
    span_kind: str = "custom"
    provider: str | None = None
    model_type: str | None = None
    started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    ended_at: datetime | None = None
    duration_ms: float | None = None
    status: str | None = None
    error_type: str | None = None
    error_message: str | None = None

    # Token counts
    tokens_input: int | None = None
    tokens_output: int | None = None
    tokens_total: int | None = None
    cost_usd: Decimal | None = None

    # Context
    session_id: str | None = None
    user_id: str | None = None
    workflow_name: str | None = None
    environment: str = "production"

    # Extension data
    model_call: dict | None = None
    tool_call: dict | None = None
    retrieval: dict | None = None
    agent_planning: dict | None = None

    # Arbitrary extra
    custom: dict | None = None

    def set_model_call(self, **kwargs: Any) -> SpanData:
        self.model_call = kwargs
        # Auto-calculate cost if we have model info
        model_name = kwargs.get("model_name")
        if model_name and self.cost_usd is None:
            self.cost_usd = estimate_cost(model_name, self.tokens_input, self.tokens_output)
        return self

    def set_tool_call(self, **kwargs: Any) -> SpanData:
        self.tool_call = kwargs
        return self

    def set_tokens(self, *, input: int = 0, output: int = 0) -> SpanData:
        self.tokens_input = input
        self.tokens_output = output
        self.tokens_total = input + output
        # Recalculate cost if model_call has model_name
        if self.model_call and self.cost_usd is None:
            self.cost_usd = estimate_cost(
                self.model_call.get("model_name"), input, output
            )
        return self

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {
            "trace_id": self.trace_id,
            "span_id": self.span_id,
            "parent_span_id": self.parent_span_id,
            "span_kind": self.span_kind,
            "started_at": self.started_at.isoformat(),
            "status": self.status or "ok",
            "environment": self.environment,
        }
        if self.ended_at:
            d["ended_at"] = self.ended_at.isoformat()
        if self.duration_ms is not None:
            d["duration_ms"] = self.duration_ms
        if self.provider:
            d["provider"] = self.provider
        if self.model_type:
            d["model_type"] = self.model_type
        if self.error_type:
            d["error_type"] = self.error_type
        if self.error_message:
            d["error_message"] = self.error_message
        if self.tokens_input is not None:
            d["tokens_input"] = self.tokens_input
        if self.tokens_output is not None:
            d["tokens_output"] = self.tokens_output
        if self.tokens_total is not None:
            d["tokens_total"] = self.tokens_total
        if self.cost_usd is not None:
            d["cost_usd"] = str(self.cost_usd)
        if self.session_id:
            d["session_id"] = self.session_id
        if self.user_id:
            d["user_id"] = self.user_id
        if self.workflow_name:
            d["workflow_name"] = self.workflow_name
        if self.model_call:
            d["model_call"] = self.model_call
        if self.tool_call:
            d["tool_call"] = self.tool_call
        if self.retrieval:
            d["retrieval"] = self.retrieval
        if self.agent_planning:
            d["agent_planning"] = self.agent_planning
        if self.custom:
            d["custom"] = self.custom
        return d
