"""Data types and models for nanomon."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any


class RunStatus(str, Enum):
    """Status of a run."""

    SUCCESS = "success"
    ERROR = "error"
    RUNNING = "running"


@dataclass
class ModelStats:
    """Accumulated statistics for a single model within a run."""

    model: str
    calls: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    cost_usd: float = 0.0
    latency_ms: int = 0

    def add_call(
        self,
        input_tokens: int,
        output_tokens: int,
        latency_ms: int,
    ) -> None:
        """Add a call's statistics to this model's accumulator."""
        self.calls += 1
        self.input_tokens += input_tokens
        self.output_tokens += output_tokens
        self.latency_ms += latency_ms

    def to_dict(self) -> dict[str, Any]:
        """Convert to API-compatible dictionary."""
        return {
            "model": self.model,
            "calls": self.calls,
            "inputTokens": self.input_tokens,
            "outputTokens": self.output_tokens,
            "costUsd": self.cost_usd,
            "latencyMs": self.latency_ms,
        }


@dataclass
class RunState:
    """Internal state for an active run being tracked."""

    id: str
    started_at: datetime
    run_context_id: str | None = None
    provider: str | None = None
    primary_model: str | None = None
    tags: list[str] = field(default_factory=list)
    metadata: dict[str, str | int | float | bool | None] = field(default_factory=dict)
    by_model: dict[str, ModelStats] = field(default_factory=dict)
    ended_at: datetime | None = None
    status: RunStatus = RunStatus.RUNNING
    error: str | None = None
    input: dict[str, Any] | None = None
    output: dict[str, Any] | None = None

    def get_or_create_model_stats(self, model: str) -> ModelStats:
        """Get or create a ModelStats instance for the given model."""
        if model not in self.by_model:
            self.by_model[model] = ModelStats(model=model)
        return self.by_model[model]

    def add_call(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        latency_ms: int,
    ) -> None:
        """Add a call's statistics to the run."""
        stats = self.get_or_create_model_stats(model)
        stats.add_call(input_tokens, output_tokens, latency_ms)
        if self.primary_model is None:
            self.primary_model = model

    def finalize(self, ended_at: datetime, error: Exception | None = None) -> None:
        """Finalize the run state."""
        self.ended_at = ended_at
        if error is not None:
            self.status = RunStatus.ERROR
            self.error = str(error)
        else:
            self.status = RunStatus.SUCCESS

    @property
    def input_tokens(self) -> int:
        """Total input tokens across all models."""
        return sum(s.input_tokens for s in self.by_model.values())

    @property
    def output_tokens(self) -> int:
        """Total output tokens across all models."""
        return sum(s.output_tokens for s in self.by_model.values())

    @property
    def total_tokens(self) -> int:
        """Total tokens (input + output) across all models."""
        return self.input_tokens + self.output_tokens

    @property
    def cost_usd(self) -> float:
        """Total cost in USD across all models."""
        return sum(s.cost_usd for s in self.by_model.values())

    @property
    def latency_ms(self) -> int:
        """Total latency in milliseconds across all models."""
        return sum(s.latency_ms for s in self.by_model.values())

    @property
    def duration_ms(self) -> int:
        """Duration of the run in milliseconds."""
        if self.ended_at is None:
            return 0
        delta = self.ended_at - self.started_at
        return int(delta.total_seconds() * 1000)

    def to_payload(self) -> dict[str, Any]:
        """Convert to API-compatible payload dictionary."""
        return {
            "id": self.id,
            "runContextId": self.run_context_id,
            "startedAt": self.started_at.isoformat().replace("+00:00", "Z"),
            "endedAt": (
                self.ended_at.isoformat().replace("+00:00", "Z")
                if self.ended_at
                else None
            ),
            "durationMs": self.duration_ms,
            "status": self.status.value,
            "error": self.error,
            "provider": self.provider,
            "primaryModel": self.primary_model,
            "inputTokens": self.input_tokens,
            "outputTokens": self.output_tokens,
            "totalTokens": self.total_tokens,
            "costUsd": self.cost_usd,
            "latencyMs": self.latency_ms,
            "tags": self.tags,
            "metadata": self.metadata,
            "byModel": [s.to_dict() for s in self.by_model.values()],
            "input": self.input,
            "output": self.output,
        }
