"""
AIECS Client - Main API for programmatic usage of AI Execute Services
"""

import asyncio
import logging
from typing import Dict, Any, Optional, List
from contextlib import asynccontextmanager

# Import core components
from aiecs.config.config import get_settings, validate_required_settings
from aiecs.domain.task.task_context import TaskContext
from aiecs.tools import discover_tools, list_tools, get_tool
from aiecs.llm.client_factory import LLMClientFactory

logger = logging.getLogger(__name__)


class AIECS:
    """
    Main AIECS client for programmatic usage
    
    This class provides a high-level API for:
    - Executing tasks with AI providers
    - Managing tool orchestration
    - Configuration management
    
    Two operation modes:
    - Simple mode: Tools and basic AI functionality (no database/Celery)
    - Full mode: Complete infrastructure (requires database/Redis)
    """
    
    def __init__(self, config: Optional[Dict[str, Any]] = None, mode: str = "simple"):
        """
        Initialize AIECS client
        
        Args:
            config: Optional configuration override
            mode: Operation mode - "simple" or "full"
        """
        self.settings = get_settings()
        self.config = config or {}
        self.mode = mode
        
        # Core components (loaded based on mode)
        self.db_manager = None
        self.task_manager = None
        self.operation_executor = None
        
        # State
        self._initialized = False
        self._tools_discovered = False
    
    async def initialize(self):
        """Initialize AIECS services based on mode"""
        if self._initialized:
            return
            
        logger.info(f"Initializing AIECS client in {self.mode} mode...")
        
        try:
            # Always discover tools
            if not self._tools_discovered:
                discover_tools("aiecs.tools")
                self._tools_discovered = True
                logger.info("Tools discovered and registered")
            
            if self.mode == "simple":
                # Simple mode: only tools, no database/Celery
                logger.info("Simple mode: tools only")
                
            elif self.mode == "full":
                # Full mode: with database and task queue
                try:
                    # Check configuration first
                    validate_required_settings("database")
                    
                    # Initialize database connection
                    from aiecs.infrastructure.persistence.database_manager import DatabaseManager
                    self.db_manager = DatabaseManager()
                    await self.db_manager.init_connection_pool()
                    logger.info("Database connection pool established")
                    
                    # Initialize task manager
                    from aiecs.infrastructure.messaging.celery_task_manager import CeleryTaskManager
                    celery_config = {
                        "broker_url": self.settings.celery_broker_url,
                        "backend_url": self.settings.celery_broker_url,
                    }
                    self.task_manager = CeleryTaskManager(celery_config)
                    logger.info("Task manager initialized")
                    
                    # Initialize operation executor
                    from aiecs.application.executors.operation_executor import OperationExecutor
                    from aiecs.tools.tool_executor import ToolExecutor
                    from aiecs.utils.execution_utils import ExecutionUtils
                    
                    tool_executor = ToolExecutor()
                    execution_utils = ExecutionUtils()
                    self.operation_executor = OperationExecutor(
                        tool_executor=tool_executor,
                        execution_utils=execution_utils,
                        config=self.config
                    )
                    
                except Exception as e:
                    logger.warning(f"Full mode initialization failed: {e}")
                    logger.info("Falling back to simple mode")
                    self.mode = "simple"
            
            self._initialized = True
            logger.info(f"AIECS client initialized successfully in {self.mode} mode")
            
        except Exception as e:
            logger.error(f"Failed to initialize AIECS: {e}")
            raise
    
    async def execute(self, context: TaskContext) -> Dict[str, Any]:
        """
        Execute a task with the given context
        
        Args:
            context: TaskContext with task parameters
            
        Returns:
            Task execution result
        """
        if not self._initialized:
            await self.initialize()
        
        if self.mode == "simple":
            # Simple mode: direct execution without queue
            logger.info("Executing task in simple mode (direct execution)")
            
            # Basic validation
            try:
                validate_required_settings("llm")
            except ValueError as e:
                return {
                    "status": "failed",
                    "error": f"LLM configuration required: {e}",
                    "mode": "simple"
                }
            
            # TODO: Implement direct task execution without Celery
            # For now, return a placeholder
            return {
                "status": "completed", 
                "result": "Simple mode execution (placeholder)",
                "mode": "simple",
                "context_id": context.user_id
            }
        
        elif self.mode == "full":
            if not self.task_manager:
                raise RuntimeError("Task manager not initialized in full mode")
            
            # Submit task to queue and get result
            task_id = await self.task_manager.submit_task(
                context=context,
                task_type="task"
            )
            
            # Wait for task completion
            result = await self._wait_for_task_completion(task_id)
            return result
        
        else:
            raise ValueError(f"Unknown mode: {self.mode}")
    
    async def execute_tool(self, tool_name: str, operation: str, params: Dict[str, Any]) -> Any:
        """
        Execute a specific tool operation directly
        
        Args:
            tool_name: Name of the tool
            operation: Operation to execute  
            params: Parameters for the operation
            
        Returns:
            Operation result
        """
        if not self._initialized:
            await self.initialize()
        
        # Get tool and execute directly (works in both modes)
        tool = get_tool(tool_name)
        if not tool:
            raise ValueError(f"Tool '{tool_name}' not found")
        
        # Prepare parameters with operation
        tool_params = {**params, "op": operation}
        
        # Execute tool
        return await tool.execute(tool_params)
    
    async def get_available_tools(self) -> List[Dict[str, Any]]:
        """Get list of all available tools"""
        if not self._tools_discovered:
            discover_tools("aiecs.tools")
            self._tools_discovered = True
        
        return list_tools()
    
    async def get_tool(self, tool_name: str):
        """Get a specific tool instance"""
        if not self._tools_discovered:
            discover_tools("aiecs.tools")
            self._tools_discovered = True
        
        return get_tool(tool_name)
    
    async def _wait_for_task_completion(self, task_id: str, timeout: int = 300) -> Dict[str, Any]:
        """
        Wait for task completion with timeout
        
        Args:
            task_id: Task ID to wait for
            timeout: Maximum wait time in seconds
            
        Returns:
            Task result
        """
        if not self.task_manager:
            raise RuntimeError("Task manager not initialized")
        
        start_time = asyncio.get_event_loop().time()
        
        while True:
            status = await self.task_manager.get_task_status(task_id)
            
            if status.get("status") in ["completed", "failed", "cancelled"]:
                return status
            
            # Check timeout
            if asyncio.get_event_loop().time() - start_time > timeout:
                raise TimeoutError(f"Task {task_id} did not complete within {timeout} seconds")
            
            # Wait before checking again
            await asyncio.sleep(1)
    
    async def close(self):
        """Close AIECS client and cleanup resources"""
        logger.info("Shutting down AIECS client...")
        
        if self.mode == "full" and self.db_manager:
            try:
                await self.db_manager.close_connection_pool()
                logger.info("Database connection pool closed")
            except Exception as e:
                logger.error(f"Error closing database: {e}")
        
        # Close all LLM clients
        try:
            await LLMClientFactory.close_all()
            logger.info("LLM clients closed")
        except Exception as e:
            logger.error(f"Error closing LLM clients: {e}")
        
        self._initialized = False
        logger.info("AIECS client shutdown complete")
    
    @asynccontextmanager
    async def session(self):
        """Context manager for AIECS session"""
        await self.initialize()
        try:
            yield self
        finally:
            await self.close()


# Convenience functions for quick usage
async def create_aiecs_client(config: Optional[Dict[str, Any]] = None, mode: str = "simple") -> AIECS:
    """
    Create and initialize an AIECS client
    
    Args:
        config: Optional configuration override
        mode: Operation mode - "simple" or "full"
        
    Returns:
        Initialized AIECS client
    """
    client = AIECS(config, mode)
    await client.initialize()
    return client

async def create_simple_client(config: Optional[Dict[str, Any]] = None) -> AIECS:
    """Create a simple AIECS client (tools only, no database/Celery)"""
    return await create_aiecs_client(config, "simple")

async def create_full_client(config: Optional[Dict[str, Any]] = None) -> AIECS:
    """Create a full AIECS client (with database and Celery)"""
    return await create_aiecs_client(config, "full")
