from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Iterator, Literal

Role = Literal["system", "user", "assistant"]


@dataclass(frozen=True)
class Message:
    role: Role
    content: str


@dataclass
class Usage:
    prompt_tokens: int = 0
    completion_tokens: int = 0

    def to_dict(self) -> dict[str, int]:
        return {
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
        }


@dataclass
class CompletionResult:
    text: str
    model: str
    provider: str
    slot: str = "primary"
    usage: Usage | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "text": self.text,
            "model": self.model,
            "provider": self.provider,
            "slot": self.slot,
            "usage": self.usage.to_dict() if self.usage else None,
        }


@dataclass
class EmbeddingResult:
    embeddings: list[list[float]]
    model: str
    provider: str
    slot: str = "primary"

    def to_dict(self) -> dict[str, Any]:
        return {
            "embeddings": self.embeddings,
            "model": self.model,
            "provider": self.provider,
            "slot": self.slot,
        }


@dataclass
class RerankResult:
    scores: list[float]
    indices: list[int]
    model: str
    provider: str
    slot: str = "primary"

    def to_dict(self) -> dict[str, Any]:
        return {
            "scores": self.scores,
            "indices": self.indices,
            "model": self.model,
            "provider": self.provider,
            "slot": self.slot,
        }


@dataclass
class OcrResult:
    text: str
    model: str
    provider: str
    slot: str = "primary"

    def to_dict(self) -> dict[str, Any]:
        return {
            "text": self.text,
            "model": self.model,
            "provider": self.provider,
            "slot": self.slot,
        }


@dataclass
class ProviderSpec:
    provider: str
    model: str
    slot: str


class LLMUnavailable(Exception):
    """三层 fallback 全部失败，或无可用 slot / API Key。"""

    def __init__(self, module: str, message: str = "", *, attempts: list[str] | None = None) -> None:
        self.module = module
        self.attempts = attempts or []
        detail = message or f"module={module} 无可用 LLM"
        super().__init__(detail)


@dataclass
class StreamChunk:
    text: str
    done: bool = False
    usage: Usage | None = None


@dataclass
class RouterTrace:
    lines: list[str] = field(default_factory=list)

    def add(self, line: str) -> None:
        self.lines.append(line)
