from __future__ import annotations

from enum import StrEnum
from typing import Any

from pydantic import (  # noqa: TC002 — Pydantic 模型
    BaseModel,
    Field,
    field_validator,
    model_validator,
)


class SupportedFormats(StrEnum):
    OPENAI = "openai"
    ANTHROPIC = "anthropic"


_BASE_CAPABILITIES = {
    ("openai", "gpt"): {"chat", "tool_use", "structured_output"},
    ("openai", "gpt-4"): {"chat", "tool_use", "structured_output", "vision", "reasoning"},
    ("anthropic", "claude"): {
        "chat",
        "tool_use",
        "structured_output",
        "vision",
        "reasoning",
        "long_context",
    },
    ("deepseek", "deepseek"): {"chat", "tool_use", "structured_output", "reasoning"},
    ("deepseek", "deepseek-v4"): {
        "chat",
        "tool_use",
        "structured_output",
        "reasoning",
        "long_context",
    },
}


def derive_capabilities(provider: str, model: str) -> set[str]:
    """根据 provider + model 前缀自动推导能力。

    调用者可以在调用此函数后，在 ModelConfig 上覆盖个别标记。
    """
    for (prov_prefix, model_prefix), caps in _BASE_CAPABILITIES.items():
        if provider.startswith(prov_prefix) and model.startswith(model_prefix):
            return set(caps)
    return {"chat"}


class ModelEndpointConfig(BaseModel):
    name: str
    model: str
    base_url: str | None = None
    base_url_env: str | None = None
    api_key: str | None = None
    api_key_env: str | None = None
    weight: int = 100
    priority: int = 100
    rpm_limit: int | None = None
    tpm_limit: int | None = None
    max_concurrency: int | None = None
    request_timeout_seconds: int | None = None
    status: str = "active"
    extra_params: dict[str, Any] = Field(default_factory=dict)


class ModelConfig(BaseModel):
    name: str
    display_name: str | None = None
    provider: str
    model: str
    format: SupportedFormats  # "openai" | "anthropic"
    chat_format: str | None = None
    tags: list[str] = Field(default_factory=list)
    supports_tools: bool | None = None
    supports_vision: bool | None = None
    supports_reasoning_effort: bool | None = None
    supports_thinking: bool | None = None
    supports_streaming: bool | None = True
    supports_structured_output: bool | None = None
    context_window: int = 32000
    response_reserve: int = 4096
    max_output_tokens: int | None = None
    max_thinking_tokens: int | None = None
    reasoning_effort: str | None = None
    timeout_seconds: int | None = 60
    dimensions: int | None = None  # 向量维度（embedding 模型专用）
    routing_strategy: str = "weighted_round_robin"
    endpoint_pool_enabled: bool = False
    weight: int = 100
    endpoints: list[ModelEndpointConfig] = Field(default_factory=list)

    @model_validator(mode="after")
    def _validate_embedding_dimensions(self) -> ModelConfig:
        """embedding 模型必须声明 dimensions。"""
        if "embedding" in self.tags and self.dimensions is None:
            raise ValueError(
                f"Embedding 模型 '{self.name}' 必须设置 dimensions（如 dimensions: 1024）"
            )
        return self

    def resolved_capabilities(self) -> dict[str, bool]:
        base = derive_capabilities(self.provider, self.model)
        return {
            "chat": True,
            "tool_use": self.supports_tools
            if self.supports_tools is not None
            else ("tool_use" in base),
            "vision": self.supports_vision
            if self.supports_vision is not None
            else ("vision" in base),
            "reasoning": self.supports_reasoning_effort
            if self.supports_reasoning_effort is not None
            else ("reasoning" in base),
            "thinking": self.supports_thinking
            if self.supports_thinking is not None
            else ("reasoning" in base),
            "streaming": self.supports_streaming if self.supports_streaming is not None else True,
            "structured_output": self.supports_structured_output
            if self.supports_structured_output is not None
            else ("structured_output" in base),
        }


class ParameterMapping(BaseModel):
    provider_param: str
    min: float | None = None
    max: float | None = None
    default: Any = None
    transform: str | None = None  # "multiply_1000", "int", "str"
    alias_for: str | None = None  # map from another param name
    handler: str | None = None  # named handler for complex param transform


class ProviderConfig(BaseModel):
    name: str
    format: SupportedFormats
    base_url_env: str
    api_key_env: str
    base_url: str | None = None
    timeout_seconds: int = 60
    max_retries: int = 1
    parameter_mappings: dict[str, ParameterMapping] = Field(default_factory=dict)
    parameter_handlers: list[str] = Field(
        default_factory=list
    )  # ordered global handlers per provider
    extra_headers: dict[str, str] = Field(default_factory=dict)
    supports_streaming: bool = True

    @field_validator("format")
    @classmethod
    def _validate_format(cls, v: str) -> str:
        allowed = {e.value for e in SupportedFormats}
        if v not in allowed:
            raise ValueError(f"format must be one of {allowed}, got '{v}'")
        return v