"""
MLL Intake Multi-Agent Orchestrator
LangGraph-based coordination of all specialized agents
"""

import uuid
from typing import Dict, Any, Optional, List, TypedDict
from datetime import datetime
from loguru import logger

from config import settings
from models.schemas import (
    AgentType, AgentResponse, ConversationState, Message,
    ClassificationResult, Intent, RequestType
)

# Import all agents
from agents.supervisor import SupervisorAgent
from agents.classifier import ClassifierAgent
from agents.intake import IntakeAgent
from agents.network import NetworkAgent
from agents.incident import IncidentAgent
from agents.access import AccessAgent
from agents.status import StatusAgent
from agents.knowledge import KnowledgeAgent

# Import utilities
from utils.semantic_classifier import classify_request, initialize_classifier
from utils.data_loader import data_loader
from utils.embeddings import embeddings_manager, initialize_embeddings


class AgentState(TypedDict):
    """State passed between agents in the graph"""
    messages: List[Dict[str, str]]
    current_agent: str
    next_agent: Optional[str]
    classification: Optional[Dict[str, Any]]
    collected_fields: Dict[str, Any]
    session_id: str
    user_input: str
    response: str
    completed: bool
    error: Optional[str]


class MultiAgentOrchestrator:
    """
    Multi-Agent Orchestrator using LangGraph-style state management
    Coordinates all 8 specialized agents
    """
    
    def __init__(self):
        self.llm_client = None
        self.agents: Dict[AgentType, Any] = {}
        self.sessions: Dict[str, ConversationState] = {}
        self._initialized = False
    
    def initialize(self) -> bool:
        """Initialize the orchestrator and all agents"""
        try:
            # Initialize LLM client
            if not self._init_llm_client():
                logger.error("Failed to initialize LLM client")
                return False
            
            # Initialize all agents
            self._init_agents()
            
            # Initialize classifier
            initialize_classifier()
            
            # Load data
            logger.info("Loading historical data...")
            data_loader.load_all()
            
            # Initialize embeddings
            logger.info("Initializing embeddings...")
            documents = data_loader.get_documents_for_embedding()
            if documents:
                initialize_embeddings(documents)
            else:
                logger.warning("No documents available for embedding")
                initialize_embeddings()
            
            self._initialized = True
            logger.info("Multi-Agent Orchestrator initialized successfully")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize orchestrator: {e}")
            return False
    
    def _init_llm_client(self) -> bool:
        """Initialize Anthropic LLM client"""
        try:
            import anthropic
            
            if not settings.validate_api_key():
                logger.error("Invalid or missing API key")
                return False
            
            self.llm_client = anthropic.Anthropic(
                api_key=settings.model.api_key
            )
            logger.info("Anthropic client initialized")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize Anthropic client: {e}")
            return False
    
    def _init_agents(self):
        """Initialize all specialized agents"""
        agent_classes = {
            AgentType.SUPERVISOR: SupervisorAgent,
            AgentType.CLASSIFIER: ClassifierAgent,
            AgentType.INTAKE: IntakeAgent,
            AgentType.NETWORK: NetworkAgent,
            AgentType.INCIDENT: IncidentAgent,
            AgentType.ACCESS: AccessAgent,
            AgentType.STATUS: StatusAgent,
            AgentType.KNOWLEDGE: KnowledgeAgent
        }
        
        for agent_type, agent_class in agent_classes.items():
            agent = agent_class()
            agent.initialize(self.llm_client)
            self.agents[agent_type] = agent
            logger.info(f"Initialized {agent_type.value} agent")
    
    def create_session(self) -> str:
        """Create a new conversation session"""
        session_id = str(uuid.uuid4())
        self.sessions[session_id] = ConversationState(
            session_id=session_id,
            messages=[],
            current_intent=Intent.UNKNOWN,
            collected_fields={},
            missing_fields=[],
            active_agent=AgentType.SUPERVISOR,
            created_tickets=[]
        )
        logger.info(f"Created session: {session_id}")
        return session_id
    
    def get_session(self, session_id: str) -> Optional[ConversationState]:
        """Get existing session or create new one"""
        if session_id not in self.sessions:
            self.sessions[session_id] = ConversationState(
                session_id=session_id,
                messages=[],
                current_intent=Intent.UNKNOWN,
                collected_fields={},
                missing_fields=[],
                active_agent=AgentType.SUPERVISOR,
                created_tickets=[]
            )
        return self.sessions.get(session_id)
    
    def process_message(
        self, 
        session_id: str, 
        user_message: str
    ) -> str:
        """
        Process a user message through the multi-agent system
        """
        if not self._initialized:
            if not self.initialize():
                return "I apologize, but I'm having trouble initializing. Please try again in a moment."
        
        # Get or create session
        state = self.get_session(session_id)
        if not state:
            return "Session error. Please refresh and try again."
        
        # Add user message to history
        state.messages.append(Message(
            role="user",
            content=user_message,
            timestamp=datetime.now()
        ))
        
        try:
            # Run the agent pipeline
            response = self._run_agent_pipeline(user_message, state)
            
            # Add assistant response to history
            state.messages.append(Message(
                role="assistant",
                content=response,
                timestamp=datetime.now()
            ))
            
            return response
            
        except Exception as e:
            logger.error(f"Error processing message: {e}")
            return f"I encountered an error while processing your request. Please try again."
    
    def _run_agent_pipeline(
        self, 
        user_message: str, 
        state: ConversationState
    ) -> str:
        """Run the multi-agent pipeline"""
        
        # Step 1: Classify the request
        classification = classify_request(user_message, state.collected_fields)
        state.classification = classification
        
        # Step 2: Determine which agent should handle this
        target_agent = self._route_to_agent(classification, state)
        state.active_agent = target_agent
        
        # Step 3: Get response from target agent
        agent = self.agents.get(target_agent)
        if not agent:
            agent = self.agents[AgentType.SUPERVISOR]
        
        # Process with the agent
        agent_response = agent.process(user_message, state)
        
        # Step 4: Update state with agent's updates
        if agent_response.state_updates:
            for key, value in agent_response.state_updates.items():
                if hasattr(state, key):
                    setattr(state, key, value)
                elif key == 'collected_fields':
                    state.collected_fields.update(value)
        
        # Step 5: Check if we need to chain to another agent
        if agent_response.next_agent and agent_response.next_agent != target_agent:
            next_agent = self.agents.get(agent_response.next_agent)
            if next_agent:
                next_response = next_agent.process(user_message, state)
                if next_response.message and next_response.message != agent_response.message:
                    return f"{agent_response.message}\n\n{next_response.message}"
        
        return agent_response.message
    
    def _route_to_agent(
        self, 
        classification: ClassificationResult, 
        state: ConversationState
    ) -> AgentType:
        """
        Determine which agent should handle the request
        
        CRITICAL: Maintains agent continuity during active workflows!
        """
        
        # RULE 1: If there's an active workflow awaiting confirmation, continue with that agent
        if state.awaiting_confirmation:
            logger.debug(f"Continuing with {state.active_agent} (awaiting confirmation)")
            return state.active_agent
        
        # RULE 2: If there's an active agent AND we're collecting fields, continue with that agent
        # This prevents re-routing in the middle of a workflow!
        if state.active_agent and state.active_agent != AgentType.SUPERVISOR:
            # Check if the agent still has fields to collect
            if state.missing_fields:
                logger.debug(f"Continuing with {state.active_agent} (collecting fields: {state.missing_fields})")
                return state.active_agent
            
            # Check if we've been in an active workflow (have collected some fields)
            workflow_agents = [AgentType.INTAKE, AgentType.INCIDENT, AgentType.NETWORK, AgentType.ACCESS]
            if state.active_agent in workflow_agents and state.collected_fields:
                # Still in active workflow - check if workflow is complete
                if not state.completed:
                    logger.debug(f"Continuing with {state.active_agent} (active workflow)")
                    return state.active_agent
        
        # RULE 3: For new conversations, route based on intent
        intent_routing = {
            Intent.CREATE_REQUEST: self._route_create_request,
            Intent.REPORT_INCIDENT: lambda c, s: AgentType.INCIDENT,
            Intent.CHECK_STATUS: lambda c, s: AgentType.STATUS,
            Intent.GET_ACCESS: lambda c, s: AgentType.ACCESS,
            Intent.GET_HELP: lambda c, s: AgentType.KNOWLEDGE,
            Intent.MODIFY_REQUEST: lambda c, s: AgentType.STATUS,
        }
        
        route_func = intent_routing.get(classification.intent)
        if route_func:
            target = route_func(classification, state)
            logger.debug(f"Routing to {target} based on intent {classification.intent}")
            return target
        
        return AgentType.SUPERVISOR
    
    def _route_create_request(
        self, 
        classification: ClassificationResult, 
        state: ConversationState
    ) -> AgentType:
        """Route create request to appropriate agent"""
        
        if classification.request_type == RequestType.NETWORK:
            if state.collected_fields.get('parent_ritm'):
                return AgentType.NETWORK
            return AgentType.INTAKE
        
        if classification.request_type == RequestType.ACCESS:
            return AgentType.ACCESS
        
        if classification.request_type == RequestType.INCIDENT:
            return AgentType.INCIDENT
        
        return AgentType.INTAKE
    
    def get_conversation_history(self, session_id: str) -> List[Dict[str, str]]:
        """Get conversation history for a session"""
        state = self.sessions.get(session_id)
        if not state:
            return []
        
        return [
            {"role": msg.role, "content": msg.content}
            for msg in state.messages
        ]
    
    def clear_session(self, session_id: str):
        """Clear a session's state"""
        if session_id in self.sessions:
            del self.sessions[session_id]
            logger.info(f"Cleared session: {session_id}")
    
    def get_session_summary(self, session_id: str) -> Dict[str, Any]:
        """Get summary of session state"""
        state = self.sessions.get(session_id)
        if not state:
            return {}
        
        return {
            'session_id': session_id,
            'message_count': len(state.messages),
            'current_intent': state.current_intent.value if state.current_intent else 'unknown',
            'active_agent': state.active_agent.value if state.active_agent else 'supervisor',
            'collected_fields': list(state.collected_fields.keys()),
            'created_tickets': state.created_tickets,
            'completed': state.completed
        }


# Global orchestrator instance
orchestrator = MultiAgentOrchestrator()


def initialize_system() -> bool:
    """Initialize the complete system"""
    return orchestrator.initialize()


def process_message(session_id: str, message: str) -> str:
    """Process a message through the orchestrator"""
    return orchestrator.process_message(session_id, message)


def create_session() -> str:
    """Create a new session"""
    return orchestrator.create_session()
