from typing import Optional, List, Dict, Any, Literal, Callable, TYPE_CHECKING
from pydantic import BaseModel, Field
from enum import Enum

if TYPE_CHECKING:
    from modules.datamodels.datamodelExtraction import ContentPart

# Operation Types
class OperationTypeEnum(str, Enum):
    GENERAL = "general"
    PLAN = "plan"
    ANALYSE = "analyse"
    GENERATE = "generate"
    EXTRACT = "extract"
    WEB_RESEARCH = "webResearch"
    IMAGE_ANALYSE = "imageAnalyse"
    IMAGE_GENERATE = "imageGenerate"


# Processing Modes
class ProcessingModeEnum(str, Enum):
    BASIC = "basic"
    ADVANCED = "advanced"
    DETAILED = "detailed"

# Priority Levels
class PriorityEnum(str, Enum):
    SPEED = "speed"
    QUALITY = "quality"
    COST = "cost"
    BALANCED = "balanced"


# Model Capabilities Enumeration
class ModelCapabilitiesEnum(str, Enum):
    # Text generation capabilities
    TEXT_GENERATION = "text_generation"
    CHAT = "chat"
    REASONING = "reasoning"
    ANALYSIS = "analysis"
    
    # Image capabilities
    IMAGE_ANALYSE = "imageAnalyse"
    IMAGE_GENERATE = "imageGenerate"
    VISION = "vision"
    MULTIMODAL = "multimodal"
    ART = "art"
    VISUAL_CREATION = "visual_creation"
    
    # Web capabilities
    WEB_SEARCH = "web_search"
    WEB_CRAWLING = "web_crawling"
    CONTENT_EXTRACTION = "content_extraction"
    TEXT_EXTRACTION = "text_extraction"
    INFORMATION_RETRIEVAL = "information_retrieval"
    URL_DISCOVERY = "url_discovery"
    MAPPING = "mapping"
    
    # Research capabilities
    RESEARCH = "research"
    QUESTION_ANSWERING = "question_answering"
    INFORMATION_GATHERING = "information_gathering"
    NEWS = "news"
    CURRENT_EVENTS = "current_events"

    
class AiModel(BaseModel):
    """Enhanced AI model definition with dynamic capabilities."""
    
    # Core identification
    name: str = Field(description="Unique model identifier")
    displayName: str = Field(description="Human-readable model name")
    connectorType: str = Field(description="Type of connector (openai, anthropic, perplexity, tavily, etc.)")
    
    # Token and context limits
    maxTokens: int = Field(description="Maximum tokens this model can generate")
    contextLength: int = Field(description="Maximum context length this model can handle")
    
    # Cost information
    costPer1kTokensInput: float = Field(default=0.0, description="Cost per 1000 input tokens")
    costPer1kTokensOutput: float = Field(default=0.0, description="Cost per 1000 output tokens")
    
    # Performance ratings
    speedRating: int = Field(ge=1, le=10, description="Speed rating (1-10, higher = faster)")
    qualityRating: int = Field(ge=1, le=10, description="Quality rating (1-10, higher = better)")

    # Function reference (not serialized)
    functionCall: Optional[Callable] = Field(default=None, exclude=True, description="Function to call for this model")
    calculatePriceUsd: Optional[Callable] = Field(default=None, exclude=True, description="Function to calculate price in USD")
    
    # Selection criteria
    capabilities: List[ModelCapabilitiesEnum] = Field(description="List of model capabilities. See ModelCapabilitiesEnum enum for available values.")
    priority: PriorityEnum = Field(default=PriorityEnum.BALANCED, description="Default priority for this model. See PriorityEnum for available values.")
    processingMode: ProcessingModeEnum = Field(default=ProcessingModeEnum.BASIC, description="Default processing mode. See ProcessingModeEnum for available values.")
    operationTypes: List[OperationTypeEnum] = Field(default=[], description="Operation types this model should avoid")
    minContextLength: Optional[int] = Field(default=None, description="Minimum context length required")
    isAvailable: bool = Field(default=True, description="Whether model is currently available")
    
    # Metadata
    version: Optional[str] = Field(default=None, description="Model version")
    lastUpdated: Optional[str] = Field(default=None, description="Last update timestamp")
    
    class Config:
        arbitraryTypesAllowed = True  # Allow Callable type


class SelectionRule(BaseModel):
    """A rule for model selection."""
    name: str = Field(description="Rule name identifier")
    condition: str = Field(description="Description of when this rule applies")
    weight: float = Field(description="Weight for scoring (higher = more important)")
    operationTypes: List[OperationTypeEnum] = Field(description="Operation types this rule applies to")
    priority: PriorityEnum = Field(default=PriorityEnum.BALANCED, description="Priority level for this rule")
    capabilities: List[ModelCapabilitiesEnum] = Field(default=[], description="Required capabilities for this rule")
    minQualityRating: Optional[int] = Field(default=None, description="Minimum quality rating")
    maxCost: Optional[float] = Field(default=None, description="Maximum cost threshold")
    minContextLength: Optional[int] = Field(default=None, description="Minimum context length required")


class AiCallOptions(BaseModel):
    """Options for centralized AI processing with clear operation types and tags."""
    operationType: OperationTypeEnum = Field(default=OperationTypeEnum.GENERAL, description="Type of operation")
    priority: PriorityEnum = Field(default=PriorityEnum.BALANCED, description="Priority level")
    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")
    processingMode: ProcessingModeEnum = Field(default=ProcessingModeEnum.BASIC, description="Processing mode")
    resultFormat: Optional[str] = Field(default=None, description="Expected result format: txt, json, csv, xml, etc.")
    
    safetyMargin: float = Field(default=0.1, ge=0.0, le=0.5, description="Safety margin for token limits (0.0-0.5)")
    capabilities: Optional[List[ModelCapabilitiesEnum]] = 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)
    contentParts: Optional[List['ContentPart']] = None  # NEW: Content parts for model-aware chunking


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

    content: str = Field(description="AI response content")
    modelName: str = Field(description="Selected model name")
    priceUsd: float = Field(default=0.0, description="Calculated price in USD")
    processingTime: float = Field(default=0.0, description="Duration in seconds")
    bytesSent: int = Field(default=0, description="Input data size in bytes")
    bytesReceived: int = Field(default=0, description="Output data size in bytes")
    errorCount: int = Field(default=0, description="0 for success, 1+ for errors")


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"
    )


class AiModelCall(BaseModel):
    """Standardized input for AI model calls."""
    
    messages: List[Dict[str, Any]] = Field(description="Messages in OpenAI format (role, content)")
    model: Optional[AiModel] = Field(default=None, description="The AI model being called")
    options: Dict[str, Any] = Field(default_factory=dict, description="Additional model-specific options")
    
    class Config:
        arbitraryTypesAllowed = True


class AiModelResponse(BaseModel):
    """Standardized output from AI model calls."""
    
    content: str = Field(description="The AI response content")
    success: bool = Field(default=True, description="Whether the call was successful")
    error: Optional[str] = Field(default=None, description="Error message if success=False")
    
    # Optional metadata that models can include
    modelId: Optional[str] = Field(default=None, description="Model identifier used")
    processingTime: Optional[float] = Field(default=None, description="Processing time in seconds")
    tokensUsed: Optional[Dict[str, int]] = Field(default=None, description="Token usage (input, output, total)")
    metadata: Optional[Dict[str, Any]] = Field(default=None, description="Additional model-specific metadata")
    
    class Config:
        arbitraryTypesAllowed = True