"""Data models for AI Bridge."""

from dataclasses import dataclass
from enum import Enum
from typing import Optional, Dict, Any

from .registry import VENDORS


# Dynamically generate VendorType enum from registry
VendorType = Enum("VendorType", {v.upper(): v for v in VENDORS})


@dataclass
class Usage:
    """Token usage statistics."""
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    total_tokens: Optional[int] = None
    
    @classmethod
    def from_openai(cls, usage) -> "Usage":
        """Create from OpenAI SDK usage object."""
        if not usage:
            return cls()
        return cls(
            prompt_tokens=getattr(usage, "prompt_tokens", None),
            completion_tokens=getattr(usage, "completion_tokens", None),
            total_tokens=getattr(usage, "total_tokens", None),
        )
    
    @classmethod
    def from_gemini(cls, metadata) -> "Usage":
        """Create from Gemini SDK usage metadata."""
        if not metadata:
            return cls()
        prompt = getattr(metadata, "prompt_token_count", None)
        completion = getattr(metadata, "candidates_token_count", None)
        # 优先使用 SDK 返回的 total，没有则手动计算
        total = getattr(metadata, "total_token_count", None)
        if total is None and prompt is not None and completion is not None:
            total = prompt + completion
        return cls(
            prompt_tokens=prompt,
            completion_tokens=completion,
            total_tokens=total,
        )
    
    @classmethod
    def from_dashscope(cls, usage: Optional[Dict]) -> "Usage":
        """Create from Dashscope SDK usage dict."""
        if not usage:
            return cls()
        return cls(
            prompt_tokens=usage.get("input_tokens"),
            completion_tokens=usage.get("output_tokens"),
            total_tokens=usage.get("total_tokens"),
        )

    def __add__(self, other: "Usage") -> "Usage":
        """
        Add two Usage objects together for aggregated statistics.

        Args:
            other: Another Usage object to add.

        Returns:
            New Usage object with combined token counts.

        Example:
            >>> usage1 = Usage(prompt_tokens=100, completion_tokens=50)
            >>> usage2 = Usage(prompt_tokens=200, completion_tokens=100)
            >>> total = usage1 + usage2
            >>> total.total_tokens  # None (not auto-calculated)
            >>> total.prompt_tokens
            300
        """
        return Usage(
            prompt_tokens=self._add_values(self.prompt_tokens, other.prompt_tokens),
            completion_tokens=self._add_values(self.completion_tokens, other.completion_tokens),
            total_tokens=self._add_values(self.total_tokens, other.total_tokens),
        )

    def __radd__(self, other: int) -> "Usage":
        """Support sum() with Usage as start value: 0 + usage."""
        if other == 0:
            return self
        return NotImplemented

    @staticmethod
    def _add_values(a: Optional[int], b: Optional[int]) -> Optional[int]:
        """Safely add two optional integers."""
        if a is None and b is None:
            return None
        if a is None:
            return b
        if b is None:
            return a
        return a + b


@dataclass
class Response:
    """
    AI response wrapper.
    
    Attributes:
        content: Raw AI response content, unmodified.
        usage: Token usage statistics.
        raw: Original SDK response object (for advanced users).
    """
    content: str
    usage: Usage
    raw: Optional[Any] = None
    
    def __str__(self) -> str:
        return self.content
