"""Embedding API 数据模型"""

from typing import Union, List, Optional, Dict, Any
from pydantic import Field, field_validator
from .base import OctenBaseModel
from ..utils.constants import EMBEDDING_MODELS, INPUT_TYPES


class EmbeddingRequest(OctenBaseModel):
    """Embedding API 请求模型

    Attributes:
        input: 输入文本（字符串或字符串列表）
        model: 模型名称（octen-embedding-0.6b/4b/8b）
        dimension: 向量维度
        input_type: 输入类型（query/document）
        truncation: 是否截断超长输入
    """

    input: Union[str, List[str]] = Field(
        ...,
        description="输入文本",
    )
    model: Optional[str] = Field(
        None,
        description="模型: octen-embedding-0.6b/4b/8b",
    )
    dimension: Optional[int] = Field(
        None,
        ge=1,
        description="向量维度",
    )
    input_type: Optional[str] = Field(
        None,
        description="输入类型: query/document",
    )
    truncation: Optional[bool] = Field(
        True,
        description="是否截断超长输入",
    )

    @field_validator("model")
    @classmethod
    def validate_model(cls, v: Optional[str]) -> Optional[str]:
        """验证模型名称"""
        if v is not None:
            valid_models = set(EMBEDDING_MODELS.values())
            if v not in valid_models:
                raise ValueError(f"model 必须是 {valid_models} 之一")
        return v

    @field_validator("input_type")
    @classmethod
    def validate_input_type(cls, v: Optional[str]) -> Optional[str]:
        """验证输入类型"""
        if v is not None and v not in INPUT_TYPES:
            raise ValueError(f"input_type 必须是 {INPUT_TYPES} 之一")
        return v


class EmbeddingObject(OctenBaseModel):
    """单个嵌入对象

    Attributes:
        index: 在批次中的索引
        embedding: 向量表示
    """

    index: int = Field(..., description="索引")
    embedding: List[float] = Field(..., description="向量")

    @property
    def dimension(self) -> int:
        """获取向量维度"""
        return len(self.embedding)


class EmbeddingResponse(OctenBaseModel):
    """Embedding API 响应模型

    Attributes:
        code: 响应码（0 表示成功）
        msg: 响应消息
        data: 响应数据
        meta: 元数据（包含 usage、warning 等）
    """

    code: int = Field(..., description="响应码")
    msg: str = Field(..., description="响应消息")
    data: Dict[str, Any] = Field(..., description="响应数据")
    meta: Dict[str, Any] = Field(..., description="元数据")

    @property
    def results(self) -> List[Dict[str, Any]]:
        """获取嵌入结果列表"""
        return self.data.get("results", [])

    @property
    def model(self) -> Optional[str]:
        """获取实际使用的模型"""
        return self.data.get("model")

    @property
    def usage(self) -> Optional[Dict[str, Any]]:
        """获取使用量信息"""
        return self.meta.get("usage")

    @property
    def warning(self) -> Optional[str]:
        """获取警告信息"""
        return self.meta.get("warning")

    def get_embeddings(self) -> List[List[float]]:
        """获取所有嵌入向量

        Returns:
            嵌入向量列表
        """
        return [result["embedding"] for result in self.results if "embedding" in result]

    def get_first_embedding(self) -> Optional[List[float]]:
        """获取第一个嵌入向量（单输入时使用）

        Returns:
            第一个嵌入向量，如果没有则返回 None
        """
        embeddings = self.get_embeddings()
        return embeddings[0] if embeddings else None
