from typing import Literal, Union
from pydantic import BaseModel, field_validator
from fastapi import UploadFile
from ombra.storage import RehydratedFile # Import RehydratedFile
from enum import Enum


class Model(str, Enum):
    """Supported LLM models."""
    # OpenAI
    GPT_4O = "gpt-4o"
    GPT_4O_MINI = "gpt-4o-mini"
    
    # Anthropic
    CLAUDE_3_5_SONNET = "claude-3-5-sonnet-20240620"
    CLAUDE_3_HAIKU = "claude-3-haiku-20240307"
    
    # Google (Gemini)
    GEMINI_1_5_PRO = "gemini/gemini-1.5-pro"
    GEMINI_1_5_FLASH = "gemini/gemini-1.5-flash"
    GEMINI_2_5_FLASH = "gemini/gemini-2.5-flash"
    GEMINI_2_5_FLASH_LITE = "gemini/gemini-2.5-flash-lite" # Added based on user feedback
    
    # Meta
    LLAMA_3_1_70B = "llama-3.1-70b-versatile" # Example provider mapping (groq/others)


class Role(str, Enum):
    """Message roles."""
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    FUNCTION = "function"
    TOOL = "tool"


class Usage(BaseModel):
    """Token usage statistics."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class LLMResponse(BaseModel):
    """Standardized response from an LLM call."""
    content: str | None = None
    role: Role = Role.ASSISTANT
    usage: Usage
    model: str
    finish_reason: str | None = None


class ContentPart(BaseModel):
    """A part of a message content (text or file)."""
    type: Literal["text", "file"] = "text"
    text: str | None = None
    file: Union[UploadFile, RehydratedFile] | None = None # Allow RehydratedFile as well
    
    class Config:
        arbitrary_types_allowed = True


class Message(BaseModel):
    """A message in the conversation."""
    role: Role
    content: str | list[Union[str, ContentPart]] 

    class Config:
        arbitrary_types_allowed = True

    @field_validator('content', mode='before')
    @classmethod
    def validate_content(cls, v):
        if isinstance(v, list):
            new_list = []
            for item in v:
                # Check for both UploadFile and RehydratedFile
                if isinstance(item, (UploadFile, RehydratedFile)):
                    # Auto-convert raw file to ContentPart
                    new_list.append(ContentPart(type="file", file=item))
                else:
                    new_list.append(item)
            return new_list
        return v
