import logging
from typing import Dict, Any, Optional
from datetime import datetime, UTC
import asyncio

from modules.methods import MethodBase, MethodResult
from modules.interfaces.serviceChatModel import AgentTask, AgentAction, AgentResult, Action, TaskStatus, ActionStatus

logger = logging.getLogger(__name__)

class ServiceContainer:
    """Service container with improved state management"""
    
    def __init__(self):
        self.state = {
            'status': TaskStatus.PENDING,
            'retryCount': 0,
            'retryMax': 3,
            'timeout': 300,  # 5 minutes
            'lastError': None,
            'lastErrorTime': None
        }
        self.methods: Dict[str, MethodBase] = {}
        self.tasks: Dict[str, Any] = {}  # Will be populated with AgentTask instances
        
    def register_method(self, method: MethodBase) -> None:
        """Register a method in the container"""
        self.methods[method.name] = method
        logger.info(f"Registered method: {method.name}")
        
    async def execute_task(self, task: Any) -> None:  # task: AgentTask
        """Execute task with improved error handling and timeout"""
        try:
            # Check for timeout
            if (datetime.now(UTC) - datetime.fromisoformat(task.startedAt)).seconds > self.state['timeout']:
                task.status = TaskStatus.TIMEOUT
                return
                
            # Execute actions
            for action in task.actionList:
                if not task.can_execute_action(action):
                    if not task.get_auth_data(action.auth_source):
                        action.status = ActionStatus.FAILED
                        task.error = f"Missing authentication for {action.auth_source}"
                    else:
                        action.status = ActionStatus.DEPENDENCY_FAILED
                    continue
                    
                try:
                    # Get method
                    method = self.methods.get(action.method)
                    if not method:
                        raise ValueError(f"Unknown method: {action.method}")
                        
                    # Validate parameters
                    if not await method.validate_parameters(action.action, action.parameters):
                        raise ValueError(f"Invalid parameters for {action.method}:{action.action}")
                        
                    # Get auth data if needed
                    auth_data = None
                    if action.auth_source and action.auth_source != "local":
                        auth_data = task.get_auth_data(action.auth_source)
                        if not auth_data:
                            raise ValueError(f"Missing authentication data for {action.auth_source}")
                        
                    # Execute with timeout
                    result = await asyncio.wait_for(
                        method.execute(action.action, action.parameters, auth_data),
                        timeout=action.timeout or 60
                    )
                    
                    if result.success:
                        action.status = ActionStatus.SUCCESS
                    else:
                        if self._should_retry(result.data.get('error')):
                            action.retryCount += 1
                            if action.retryCount > action.retryMax:
                                action.status = ActionStatus.FAILED
                                if action.rollback_on_failure:
                                    await method.rollback(action.action, action.parameters, auth_data)
                            else:
                                action.status = ActionStatus.RETRY
                        else:
                            action.status = ActionStatus.FAILED
                            if action.rollback_on_failure:
                                await method.rollback(action.action, action.parameters, auth_data)
                                
                except asyncio.TimeoutError:
                    action.status = ActionStatus.TIMEOUT
                except Exception as e:
                    action.status = ActionStatus.FAILED
                    if action.rollback_on_failure:
                        await method.rollback(action.action, action.parameters, auth_data)
                        
            # Update task status
            if task.has_failed():
                task.status = TaskStatus.FAILED
            elif task.is_complete():
                task.status = TaskStatus.SUCCESS
                task.finishedAt = datetime.now(UTC).isoformat()
                
        except Exception as e:
            task.status = TaskStatus.FAILED
            task.error = str(e)
            
    def _should_retry(self, error: str) -> bool:
        """Determine if error is retryable"""
        retryable_errors = [
            "AI down",
            "Document not found",
            "Content extraction failed",
            "Network error",
            "Temporary failure"
        ]
        return any(err in error for err in retryable_errors)
        
    def get_method(self, name: str) -> Optional[MethodBase]:
        """Get a method by name"""
        return self.methods.get(name)
        
    def get_available_methods(self) -> Dict[str, Dict[str, Any]]:
        """Get catalog of available methods and their actions"""
        return {
            name: {
                "description": method.description,
                "actions": method.actions
            }
            for name, method in self.methods.items()
        }