"""Embedding API data models"""

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 request model

    Attributes:
        input: Input text (string or list of strings)
        model: Model name (octen-embedding-0.6b/4b/8b)
        dimension: Vector dimension
        input_type: Input type（query/document）
        truncation: Whether to truncate overly long input
    """

    input: Union[str, List[str]] = Field(
        ...,
        description="Input text",
    )
    model: Optional[str] = Field(
        None,
        description="Model: octen-embedding-0.6b/4b/8b",
    )
    dimension: Optional[int] = Field(
        None,
        ge=1,
        description="Vector dimension",
    )
    input_type: Optional[str] = Field(
        None,
        description="Input type: query/document",
    )
    truncation: Optional[bool] = Field(
        True,
        description="Whether to truncate overly long input",
    )

    @field_validator("model")
    @classmethod
    def validate_model(cls, v: Optional[str]) -> Optional[str]:
        """validateModel name"""
        if v is not None:
            valid_models = set(EMBEDDING_MODELS.values())
            if v not in valid_models:
                raise ValueError(f"model must be one of {valid_models}")
        return v

    @field_validator("input_type")
    @classmethod
    def validate_input_type(cls, v: Optional[str]) -> Optional[str]:
        """validateInput type"""
        if v is not None and v not in INPUT_TYPES:
            raise ValueError(f"input_type must be one of {INPUT_TYPES}")
        return v


class EmbeddingObject(OctenBaseModel):
    """Single embedding object

    Attributes:
        index: Index in batch
        embedding: Vector representation
    """

    index: int = Field(..., description="Index")
    embedding: List[float] = Field(..., description="Vector")

    @property
    def dimension(self) -> int:
        """GetVector dimension"""
        return len(self.embedding)


class EmbeddingResponse(OctenBaseModel):
    """Embedding API response model

    Attributes:
        code: Response code (0 for success)
        msg: Response message
        data: Response data
        meta: Metadata (includes usage, warning, etc.)
    """

    code: int = Field(..., description="Response code")
    msg: str = Field(..., description="Response message")
    data: Dict[str, Any] = Field(..., description="Response data")
    meta: Dict[str, Any] = Field(..., description="Metadata")

    @property
    def results(self) -> List[Dict[str, Any]]:
        """GetEmbeddingResult/ResultsList"""
        return self.data.get("results", [])

    @property
    def model(self) -> Optional[str]:
        """Get actually used model"""
        return self.data.get("model")

    @property
    def usage(self) -> Optional[Dict[str, Any]]:
        """GetUsage information"""
        return self.meta.get("usage")

    @property
    def warning(self) -> Optional[str]:
        """Get warning messages"""
        return self.meta.get("warning")

    def get_embeddings(self) -> List[List[float]]:
        """Get all embedding vectors

        Returns:
            EmbeddingVectorList
        """
        return [result["embedding"] for result in self.results if "embedding" in result]

    def get_first_embedding(self) -> Optional[List[float]]:
        """Get first embedding vector (for single input)

        Returns:
            First embedding vector, or None if not available
        """
        embeddings = self.get_embeddings()
        return embeddings[0] if embeddings else None
