from __future__ import annotations

from typing import Any, Dict, List, Literal, Optional

from pydantic import BaseModel, ConfigDict, Field, model_validator


class SDKModel(BaseModel):
    model_config = ConfigDict(extra="forbid", validate_assignment=True)


class Message(SDKModel):
    role: Literal["system", "user", "assistant", "tool"]
    content: str
    name: Optional[str] = None


class ImageInput(SDKModel):
    url: Optional[str] = None
    base64_data: Optional[str] = Field(default=None, alias="base64")
    mime_type: Optional[str] = None

    @model_validator(mode="after")
    def validate_source(self) -> "ImageInput":
        if bool(self.url) == bool(self.base64_data):
            raise ValueError("Exactly one of 'url' or 'base64' must be provided.")
        return self


class RagChunk(SDKModel):
    content: str
    source_id: Optional[str] = None
    score: Optional[float] = None
    metadata: Optional[Dict[str, Any]] = None


class Inputs(SDKModel):
    messages: Optional[List[Message]] = None
    images: Optional[List[ImageInput]] = None
    rag_context: Optional[List[RagChunk]] = None
    json_schema: Optional[Dict[str, Any]] = None


class Outputs(SDKModel):
    raw: Optional[str] = None
    parsed: Optional[Any] = None


class Config(SDKModel):
    scoring_tags: List[str] = Field(default_factory=list)


class TracePayload(SDKModel):
    project_id: str
    trace_id: str
    inputs: Inputs
    outputs: Outputs
    expected: Optional[Any] = None
    metadata: Optional[Dict[str, Any]] = None
    config: Config = Field(default_factory=Config)
