from typing import Optional, List, Dict, Any, Literal
from pydantic import BaseModel, Field


# Operation Types
class OperationType:
    GENERAL = "general"
    GENERATE_PLAN = "generate_plan"
    ANALYSE_CONTENT = "analyse_content"
    GENERATE_CONTENT = "generate_content"
    WEB_RESEARCH = "web_research"
    IMAGE_ANALYSIS = "image_analysis"
    IMAGE_GENERATION = "image_generation"


# Processing Modes
class ProcessingMode:
    BASIC = "basic"
    ADVANCED = "advanced"
    DETAILED = "detailed"


# Priority Levels
class Priority:
    SPEED = "speed"
    QUALITY = "quality"
    COST = "cost"
    BALANCED = "balanced"


# Model Tags
class ModelTags:
    # Core capabilities
    TEXT = "text"
    CHAT = "chat"
    REASONING = "reasoning"
    ANALYSIS = "analysis"
    IMAGE = "image"
    VISION = "vision"
    MULTIMODAL = "multimodal"
    WEB = "web"
    SEARCH = "search"
    CRAWL = "crawl"
    EXTRACT = "extract"
    CONTENT = "content"
    INFORMATION = "information"
    
    # Quality indicators
    HIGH_QUALITY = "high_quality"
    FAST = "fast"
    COST_EFFECTIVE = "cost_effective"
    GENERAL = "general"
    
    # Specialized capabilities
    IMAGE_GENERATION = "image_generation"
    ART = "art"
    VISUAL = "visual"
    VARIATIONS = "variations"
    API = "api"
    INFO = "info"
    MODELS = "models"


# Operation Type to Required Tags Mapping
OPERATION_TAG_MAPPING = {
    OperationType.GENERAL: [ModelTags.TEXT, ModelTags.CHAT, ModelTags.REASONING],
    OperationType.GENERATE_PLAN: [ModelTags.TEXT, ModelTags.REASONING, ModelTags.ANALYSIS],
    OperationType.ANALYSE_CONTENT: [ModelTags.TEXT, ModelTags.ANALYSIS, ModelTags.REASONING],
    OperationType.GENERATE_CONTENT: [ModelTags.TEXT, ModelTags.CHAT, ModelTags.REASONING],
    OperationType.WEB_RESEARCH: [ModelTags.TEXT, ModelTags.ANALYSIS, ModelTags.REASONING],
    OperationType.IMAGE_ANALYSIS: [ModelTags.IMAGE, ModelTags.VISION, ModelTags.MULTIMODAL],
    OperationType.IMAGE_GENERATION: [ModelTags.IMAGE_GENERATION, ModelTags.ART, ModelTags.VISUAL],
}


# Processing Mode to Priority Mapping
PROCESSING_MODE_PRIORITY_MAPPING = {
    ProcessingMode.BASIC: Priority.SPEED,
    ProcessingMode.ADVANCED: Priority.BALANCED,
    ProcessingMode.DETAILED: Priority.QUALITY,
}


class ModelCapabilities(BaseModel):
    """Model capabilities and characteristics for dynamic selection."""
    
    name: str = Field(description="Model name/identifier")
    maxTokens: int = Field(description="Maximum token limit for this model")
    capabilities: List[str] = Field(description="List of capabilities: text, image, vision, reasoning, analysis, etc.")
    costPerToken: float = Field(default=0.0, description="Cost per token (if available)")
    processingTime: float = Field(default=1.0, description="Average processing time multiplier")
    isAvailable: bool = Field(default=True, description="Whether model is currently available")


class AiCallOptions(BaseModel):
    """Options for centralized AI processing with clear operation types and tags."""

    operationType: str = Field(default="general", description="Type of operation: general, generate_plan, analyse_content, generate_content, web_research")
    priority: str = Field(default="balanced", description="speed|quality|cost|balanced")
    compressPrompt: bool = Field(default=True, description="Whether to compress the prompt")
    compressContext: bool = Field(default=True, description="If False: process each chunk; If True: summarize and work on summary")
    processDocumentsIndividually: bool = Field(default=True, description="If True, process each document separately; else pool docs")
    maxContextBytes: Optional[int] = Field(default=None, description="Hard cap for extracted context size passed to the model")
    maxCost: Optional[float] = Field(default=None, description="Max cost budget")
    maxProcessingTime: Optional[int] = Field(default=None, description="Max processing time in seconds")
    requiredTags: Optional[List[str]] = Field(default=None, description="Required model tags for selection")
    processingMode: str = Field(default="basic", description="Processing mode: basic, advanced, detailed")
    resultFormat: Optional[str] = Field(default=None, description="Expected result format: txt, json, csv, xml, etc.")
    
    # New fields for dynamic strategy
    callType: Literal["planning", "text"] = Field(default="text", description="Call type: planning or text")
    safetyMargin: float = Field(default=0.1, ge=0.0, le=0.5, description="Safety margin for token limits (0.0-0.5)")
    modelCapabilities: Optional[List[str]] = Field(default=None, description="Required model capabilities for filtering")
    
    # Model generation parameters
    temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0, description="Temperature for response generation (0.0-2.0, lower = more consistent)")
    maxTokens: Optional[int] = Field(default=None, ge=1, le=32000, description="Maximum tokens in response")
    maxParts: Optional[int] = Field(default=1000, ge=1, le=1000, description="Maximum number of continuation parts to fetch")


class AiCallRequest(BaseModel):
    """Centralized AI call request payload for interface use."""

    prompt: str = Field(description="The user prompt")
    context: Optional[str] = Field(default=None, description="Optional external context (e.g., extracted docs)")
    options: AiCallOptions = Field(default_factory=AiCallOptions)


class AiCallResponse(BaseModel):
    """Standardized AI call response."""

    content: str = Field(description="AI response content")
    modelName: str = Field(description="Selected model name")
    usedTokens: Optional[int] = Field(default=None, description="Estimated used tokens")
    costEstimate: Optional[float] = Field(default=None, description="Estimated cost of the call")


class EnhancedAiCallOptions(AiCallOptions):
    """Enhanced options for improved document processing with chunk mapping."""
    
    # Parallel processing
    enableParallelProcessing: bool = Field(
        default=True, 
        description="Enable parallel processing of chunks"
    )
    maxConcurrentChunks: int = Field(
        default=5, 
        ge=1, 
        le=20, 
        description="Maximum number of chunks to process concurrently"
    )
    
    # Chunk mapping
    preserveChunkMetadata: bool = Field(
        default=True, 
        description="Preserve chunk metadata during processing"
    )
    chunkSeparator: str = Field(
        default="\n\n---\n\n", 
        description="Separator between chunks in merged output"
    )