models.py

from typing import List, Optional, Dict, Any, Union, Literal
from pydantic import BaseModel, Field, field_validator, model_validator, PrivateAttr
from datetime import datetime
import uuid
import logging

logger = logging.getLogger(__name__)


class ImageUrl(BaseModel):
    """Image URL object for OpenAI format."""
    url: str  # Can be a base64 data URL or HTTP/HTTPS URL
    detail: Optional[Literal["low", "high", "auto"]] = "auto"


class TextContentPart(BaseModel):
    """Text content part for multimodal messages."""
    type: Literal["text"]
    text: str


class ImageContentPart(BaseModel):
    """Image content part for multimodal messages."""
    type: Literal["image_url"]
    image_url: ImageUrl


# Union type for content parts - can be text or image
ContentPart = Union[TextContentPart, ImageContentPart]


class Message(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: Union[str, List[ContentPart]]
    name: Optional[str] = None
    # Store original content for image processing (private attr won't be serialized)
    _original_content: Optional[Union[str, List[Dict]]] = PrivateAttr(default=None)
    
    @model_validator(mode='after')
    def normalize_content(self):
        """Convert array content to string for Claude Code compatibility.
        
        Note: This only extracts text parts for now. Image processing happens
        separately in the image_handler module to maintain separation of concerns.
        """
        # Store original content before normalization for image processing
        if isinstance(self.content, list):
            # Store the raw dict representation for image handler
            original_list = []
            for part in self.content:
                if isinstance(part, (TextContentPart, ImageContentPart)):
                    # Convert Pydantic models to dicts
                    original_list.append(part.model_dump())
                elif isinstance(part, dict):
                    original_list.append(part)
            self._original_content = original_list
            
            # Extract text from content parts and concatenate
            text_parts = []
            has_images = False
            
            for part in self.content:
                # Handle Pydantic model instances
                if isinstance(part, TextContentPart):
                    text_parts.append(part.text)
                elif isinstance(part, ImageContentPart):
                    has_images = True
                    # Images will be processed separately by ImageHandler
                # Handle dict representations (from JSON parsing)
                elif isinstance(part, dict):
                    if part.get("type") == "text":
                        text_parts.append(part.get("text", ""))
                    elif part.get("type") == "image_url":
                        has_images = True
                        # Images will be processed separately
            
            # Log if images were found (for debugging)
            if has_images:
                logger.debug(f"Message contains image content parts (will be processed separately)")
            
            # Join all text parts with spaces (not newlines, to preserve original formatting)
            self.content = " ".join(text_parts) if text_parts else ""
        else:
            # For string content, store as-is
            self._original_content = self.content
            
        return self
    
    def model_dump(self, **kwargs):
        """Override model_dump to include original content for image processing."""
        data = super().model_dump(**kwargs)
        # Use original content if available
        if self._original_content is not None:
            data['content'] = self._original_content
        return data
    
    def dict(self, **kwargs):
        """Override dict to include original content for image processing."""
        data = super().dict(**kwargs)
        # Use original content if available
        if self._original_content is not None:
            data['content'] = self._original_content
        return data


class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[Message]
    temperature: Optional[float] = Field(default=1.0, ge=0, le=2)
    top_p: Optional[float] = Field(default=1.0, ge=0, le=1)
    n: Optional[int] = Field(default=1, ge=1)
    stream: Optional[bool] = False
    stop: Optional[Union[str, List[str]]] = None
    max_tokens: Optional[int] = None
    presence_penalty: Optional[float] = Field(default=0, ge=-2, le=2)
    frequency_penalty: Optional[float] = Field(default=0, ge=-2, le=2)
    logit_bias: Optional[Dict[str, float]] = None
    user: Optional[str] = None
    
    @field_validator('n')
    @classmethod
    def validate_n(cls, v):
        if v > 1:
            raise ValueError("Claude Code SDK does not support multiple choices (n > 1). Only single response generation is supported.")
        return v
    
    def log_unsupported_parameters(self):
        """Log warnings for parameters that are not supported by Claude Code SDK."""
        warnings = []
        
        if self.temperature != 1.0:
            warnings.append(f"temperature={self.temperature} is not supported by Claude Code SDK and will be ignored")
        
        if self.top_p != 1.0:
            warnings.append(f"top_p={self.top_p} is not supported by Claude Code SDK and will be ignored")
            
        if self.max_tokens is not None:
            warnings.append(f"max_tokens={self.max_tokens} is not supported by Claude Code SDK and will be ignored. Consider using max_turns to limit conversation length")
        
        if self.presence_penalty != 0:
            warnings.append(f"presence_penalty={self.presence_penalty} is not supported by Claude Code SDK and will be ignored")
            
        if self.frequency_penalty != 0:
            warnings.append(f"frequency_penalty={self.frequency_penalty} is not supported by Claude Code SDK and will be ignored")
            
        if self.logit_bias:
            warnings.append(f"logit_bias is not supported by Claude Code SDK and will be ignored")
            
        if self.stop:
            warnings.append(f"stop sequences are not supported by Claude Code SDK and will be ignored")
        
        for warning in warnings:
            logger.warning(f"OpenAI API compatibility: {warning}")
    
    def to_claude_options(self) -> Dict[str, Any]:
        """Convert OpenAI request parameters to Claude Code SDK options."""
        # Log warnings for unsupported parameters
        self.log_unsupported_parameters()
        
        options = {}
        
        # Direct mappings
        if self.model:
            options['model'] = self.model
            
        # Use user field for session identification if provided
        if self.user:
            # Could be used for analytics/logging or session tracking
            logger.info(f"Request from user: {self.user}")
        
        return options


class Choice(BaseModel):
    index: int
    message: Message
    finish_reason: Optional[Literal["stop", "length", "content_filter", "null"]] = None


class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class ChatCompletionResponse(BaseModel):
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:8]}")
    object: Literal["chat.completion"] = "chat.completion"
    created: int = Field(default_factory=lambda: int(datetime.now().timestamp()))
    model: str
    choices: List[Choice]
    usage: Optional[Usage] = None
    system_fingerprint: Optional[str] = None


class StreamChoice(BaseModel):
    index: int
    delta: Dict[str, Any]
    finish_reason: Optional[Literal["stop", "length", "content_filter", "null"]] = None


class ChatCompletionStreamResponse(BaseModel):
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex[:8]}")
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    created: int = Field(default_factory=lambda: int(datetime.now().timestamp()))
    model: str
    choices: List[StreamChoice]
    system_fingerprint: Optional[str] = None


class ErrorDetail(BaseModel):
    message: str
    type: str
    param: Optional[str] = None
    code: Optional[str] = None


class ErrorResponse(BaseModel):
    error: ErrorDetail