from typing import Dict, Any, List
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langgraph.graph import StateGraph, END
from .base_agent import BaseAgent, AgentState


class ConversationAgent(BaseAgent):
    """Agent for handling multi-turn conversations"""
    
    def __init__(self, model_name: str = None):
        super().__init__(model_name)
        self.conversation_history: List[BaseMessage] = []
    
    def _build_graph(self) -> StateGraph:
        """Build conversation-specific graph"""
        workflow = StateGraph(AgentState)
        
        # Add nodes
        workflow.add_node("analyze_intent", self._analyze_intent)
        workflow.add_node("generate_response", self._generate_response)
        workflow.add_node("update_context", self._update_context)
        
        # Define edges
        workflow.add_edge("analyze_intent", "generate_response")
        workflow.add_edge("generate_response", "update_context")
        workflow.add_edge("update_context", END)
        
        # Set entry point
        workflow.set_entry_point("analyze_intent")
        
        return workflow.compile()
    
    def _analyze_intent(self, state: AgentState) -> AgentState:
        """Analyze user intent and conversation context"""
        # Simple intent analysis - can be enhanced with more sophisticated logic
        user_message = state.messages[-1] if state.messages else None
        
        if user_message and isinstance(user_message, HumanMessage):
            content = user_message.content.lower()
            
            # Basic intent detection
            if any(word in content for word in ["hello", "hi", "hey", "greetings"]):
                state.context["intent"] = "greeting"
            elif any(word in content for word in ["help", "assist", "support"]):
                state.context["intent"] = "help"
            elif "?" in content:
                state.context["intent"] = "question"
            else:
                state.context["intent"] = "general"
        
        state.current_step = "intent_analyzed"
        return state
    
    def _generate_response(self, state: AgentState) -> AgentState:
        """Generate response considering conversation history"""
        # Combine conversation history with current message
        all_messages = self.conversation_history + state.messages
        
        if all_messages:
            response = self.model.invoke(all_messages)
            state.messages.append(response)
        
        state.current_step = "response_generated"
        return state
    
    def _update_context(self, state: AgentState) -> AgentState:
        """Update conversation context and history"""
        # Add the conversation to history (excluding system messages)
        for message in state.messages:
            if isinstance(message, (HumanMessage, AIMessage)):
                self.conversation_history.append(message)
        
        # Limit conversation history to avoid token limits
        if len(self.conversation_history) > 10:
            self.conversation_history = self.conversation_history[-10:]
        
        state.current_step = "context_updated"
        return state
    
    def run(self, user_input: str, **kwargs) -> Dict[str, Any]:
        """Run conversation with history management"""
        initial_state = AgentState(
            messages=[HumanMessage(content=user_input)],
            context=kwargs
        )
        
        result = self.graph.invoke(initial_state)
        
        return {
            "messages": result.messages,
            "context": result.context,
            "conversation_history": self.conversation_history,
            "final_step": result.current_step
        }