# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
import logging
import asyncio
import uuid
import base64
from typing import Dict, Any, List, Union, Tuple, Optional
from dataclasses import dataclass
import time

logger = logging.getLogger(__name__)

from modules.aicore.aicoreModelRegistry import modelRegistry
from modules.aicore.aicoreModelSelector import modelSelector
from modules.datamodels.datamodelAi import (
    AiModel,
    AiCallOptions, 
    AiCallRequest, 
    AiCallResponse,
    OperationTypeEnum,
    AiModelCall,
    AiModelResponse,
)
from modules.datamodels.datamodelExtraction import ContentPart, MergeStrategy


# Dynamic model registry - models are now loaded from connectors via aicore system


@dataclass(slots=True)
class AiObjects:
    """Centralized AI interface: dynamically discovers and uses AI models. Includes web functionality."""

    def __post_init__(self) -> None:
        # Auto-discover and register all available connectors
        self._discoverAndRegisterConnectors()
    
    def _discoverAndRegisterConnectors(self):
        """Auto-discover and register all available AI connectors."""
        logger.info("Auto-discovering AI connectors...")
        
        # Use the model registry's built-in discovery mechanism
        discoveredConnectors = modelRegistry.discoverConnectors()
        
        # Register each discovered connector
        for connector in discoveredConnectors:
            modelRegistry.registerConnector(connector)
            logger.info(f"Registered connector: {connector.getConnectorType()}")
        
        logger.info(f"Total connectors registered: {len(discoveredConnectors)}")
        logger.info("All AI connectors registered with dynamic model registry")

    @classmethod
    async def create(cls) -> "AiObjects":
        """Create AiObjects instance with auto-discovered connectors."""
        # No need to manually create connectors - they're auto-discovered
        return cls()

    def _selectModel(self, prompt: str, context: str, options: AiCallOptions) -> str:
        """Select the best model using dynamic model selection system. Returns displayName (unique identifier)."""
        # Get available models from the dynamic registry
        availableModels = modelRegistry.getAvailableModels()
        
        if not availableModels:
            logger.error("No models available in the registry")
            raise ValueError("No AI models available")
        
        # Use the dynamic model selector
        selectedModel = modelSelector.selectModel(prompt, context, options, availableModels)
        
        if not selectedModel:
            logger.error("No suitable model found for the given criteria")
            raise ValueError("No suitable AI model found")
        
        logger.info(f"Selected model: {selectedModel.name} ({selectedModel.displayName})")
        return selectedModel.displayName


    # AI for Extraction, Processing, Generation
    async def callWithTextContext(self, request: AiCallRequest) -> AiCallResponse:
        """Call AI model for traditional text/context calls with fallback mechanism."""
        prompt = request.prompt
        context = request.context or ""
        options = request.options

        # Input bytes will be calculated inside _callWithModel

        # Generation parameters are handled inside _callWithModel

        # Get failover models for this operation type
        availableModels = modelRegistry.getAvailableModels()
        
        # Filter by allowedProviders if specified (from workflow config)
        allowedProviders = getattr(options, 'allowedProviders', None) if options else None
        if allowedProviders:
            filteredModels = [m for m in availableModels if m.connectorType in allowedProviders]
            if filteredModels:
                logger.info(f"Filtered models by allowedProviders {allowedProviders}: {len(filteredModels)} models (from {len(availableModels)})")
                availableModels = filteredModels
            else:
                logger.warning(f"No models match allowedProviders {allowedProviders}, using all {len(availableModels)} available models")
        
        failoverModelList = modelSelector.getFailoverModelList(prompt, context, options, availableModels)
        
        if not failoverModelList:
            errorMsg = f"No suitable models found for operation {options.operationType}"
            logger.error(errorMsg)
            return AiCallResponse(
                content=errorMsg,
                modelName="error",
                priceCHF=0.0,
                processingTime=0.0,
                bytesSent=0,
                bytesReceived=0,
                errorCount=1
            )

        # Try each model in failover sequence
        lastError = None
        for attempt, model in enumerate(failoverModelList):
            try:
                logger.info(f"Attempting AI call with model: {model.name} (attempt {attempt + 1}/{len(failoverModelList)})")
                
                # Call the model directly - no truncation or compression here
                response = await self._callWithModel(model, prompt, context, options)
                
                logger.info(f"✅ AI call successful with model: {model.name}")
                return response
                
            except Exception as e:
                lastError = e
                logger.warning(f"❌ AI call failed with model {model.name}: {str(e)}")
                
                # If this is not the last model, try the next one
                if attempt < len(failoverModelList) - 1:
                    logger.info(f"🔄 Trying next failover model...")
                    continue
                else:
                    # All models failed
                    logger.error(f"💥 All {len(failoverModelList)} models failed for operation {options.operationType}")
                    break

        # All failover attempts failed - return error response
        errorMsg = f"All AI models failed for operation {options.operationType}. Last error: {str(lastError)}"
        logger.error(errorMsg)
        return AiCallResponse(
            content=errorMsg,
            modelName="error",
            priceCHF=0.0,
            processingTime=0.0,
            bytesSent=0,
            bytesReceived=0,
            errorCount=1
        )

    def _createErrorResponse(self, errorMsg: str, inputBytes: int, outputBytes: int) -> AiCallResponse:
        """Create an error response."""
        return AiCallResponse(
            content=errorMsg,
            modelName="error",
            priceCHF=0.0,
            processingTime=0.0,
            bytesSent=inputBytes,
            bytesReceived=outputBytes,
            errorCount=1
        )

    async def _callWithModel(self, model: AiModel, prompt: str, context: str, options: AiCallOptions = None) -> AiCallResponse:
        """Call a specific model and return the response."""
        # Calculate input bytes from prompt and context
        inputBytes = len((prompt + context).encode('utf-8'))
        
        # Replace <TOKEN_LIMIT> placeholder with model's maxTokens value
        if "<TOKEN_LIMIT>" in prompt:
            if model.maxTokens > 0:
                tokenLimit = str(model.maxTokens)
                modelPrompt = prompt.replace("<TOKEN_LIMIT>", tokenLimit)
                logger.debug(f"Replaced <TOKEN_LIMIT> with {tokenLimit} for model {model.name}")
            else:
                raise ValueError(f"Model {model.name} has invalid maxTokens ({model.maxTokens}). Cannot set token limit.")
        else:
            modelPrompt = prompt
        
        # Update messages array with replaced content
        messages = []
        if context:
            messages.append({"role": "system", "content": f"Context from documents:\n{context}"})
        messages.append({"role": "user", "content": modelPrompt})
        
        # Start timing
        startTime = time.time()
        
        # Call the model's function directly - completely generic
        if model.functionCall:
            # Create standardized call object
            modelCall = AiModelCall(
                messages=messages,
                model=model,
                options=options or {}
            )
            
            # Log before calling model
            contextSize = len(context.encode('utf-8')) if context else 0
            promptSize = len(modelPrompt.encode('utf-8')) if modelPrompt else 0
            totalInputSize = contextSize + promptSize
            logger.debug(f"Calling model {model.name} with {len(messages)} messages, context size: {contextSize} bytes, prompt size: {promptSize} bytes, total input: {totalInputSize} bytes")
            
            # Call the model with standardized interface
            modelResponse = await model.functionCall(modelCall)
            
            # Log after successful call
            logger.debug(f"Model {model.name} returned successfully")
            
            # Extract content from standardized response
            if not modelResponse.success:
                raise ValueError(f"Model call failed: {modelResponse.error}")
            content = modelResponse.content
        else:
            raise ValueError(f"Model {model.name} has no function call defined")

        # Calculate timing and output bytes
        endTime = time.time()
        processingTime = endTime - startTime
        outputBytes = len(content.encode("utf-8"))
        
        # Calculate price using model's own price calculation method
        priceCHF = model.calculatepriceCHF(processingTime, inputBytes, outputBytes)
        
        return AiCallResponse(
            content=content,
            modelName=model.name,
            priceCHF=priceCHF,
            processingTime=processingTime,
            bytesSent=inputBytes,
            bytesReceived=outputBytes,
            errorCount=0
        )


    # Utility methods
    async def listAvailableModels(self, connectorType: str = None) -> List[Dict[str, Any]]:
        """List available models, optionally filtered by connector type."""
        models = modelRegistry.getAvailableModels()
        if connectorType:
            return [model.model_dump() for model in models if model.connectorType == connectorType]
        return [model.model_dump() for model in models]

    async def getModelInfo(self, displayName: str) -> Dict[str, Any]:
        """Get information about a specific model by displayName."""
        model = modelRegistry.getModel(displayName)
        if not model:
            raise ValueError(f"Model with displayName '{displayName}' not found")
        return model.model_dump()

    async def getModelsByTag(self, tag: str) -> List[str]:
        """Get model displayNames that have a specific tag. Returns displayNames (unique identifiers)."""
        models = modelRegistry.getModelsByTag(tag)
        return [model.displayName for model in models]