from typing import Any, Dict, List, Optional
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from config import Config


class AgentState:
    """State definition for the agent workflow"""
    messages: List[BaseMessage]
    current_step: str
    context: Dict[str, Any]
    
    def __init__(self, messages: List[BaseMessage] = None, current_step: str = "start", context: Dict[str, Any] = None):
        self.messages = messages or []
        self.current_step = current_step
        self.context = context or {}


class BaseAgent:
    """Base agent class with common functionality"""
    
    def __init__(self, model_name: str = Config.DEFAULT_MODEL):
        self.model = ChatOpenAI(
            model=model_name,
            api_key=Config.OPENAI_API_KEY,
            temperature=0.1,
            max_retries=Config.MAX_RETRIES,
            timeout=Config.TIMEOUT
        )
        self.graph = self._build_graph()
    
    def _build_graph(self) -> StateGraph:
        """Build the state graph for the agent"""
        workflow = StateGraph(AgentState)
        
        # Add nodes and edges
        workflow.add_node("process_input", self._process_input)
        workflow.add_node("generate_response", self._generate_response)
        
        # Define edges
        workflow.add_edge("process_input", "generate_response")
        workflow.add_edge("generate_response", END)
        
        # Set entry point
        workflow.set_entry_point("process_input")
        
        return workflow.compile()
    
    def _process_input(self, state: AgentState) -> AgentState:
        """Process user input"""
        # This is a base implementation - subclasses should override
        state.current_step = "processed"
        return state
    
    def _generate_response(self, state: AgentState) -> AgentState:
        """Generate response using the LLM"""
        if state.messages:
            response = self.model.invoke(state.messages)
            state.messages.append(response)
        return state
    
    def run(self, user_input: str, **kwargs) -> Dict[str, Any]:
        """Run the agent with user input"""
        initial_state = AgentState(
            messages=[HumanMessage(content=user_input)],
            context=kwargs
        )
        
        result = self.graph.invoke(initial_state)
        
        return {
            "messages": result.messages,
            "context": result.context,
            "final_step": result.current_step
        }