import os
from typing import Literal, Optional
from langchain_core.messages import BaseMessage
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.prebuilt import tools_condition
from langgraph.checkpoint.base import BaseCheckpointSaver
from pydantic import BaseModel, Field
from dotenv import load_dotenv
from .agent_node import agent_node
from .tools_node import tool_node

load_dotenv()

class AgentConfig(BaseModel):
    """Configuration for the agent graph with validation."""
    model_name: str = Field(default_factory=lambda: os.getenv("MODEL_NAME", "gpt-5-mini"), description="LLM model to use")
    model_provider: str = Field(default_factory=lambda: os.getenv("MODEL_PROVIDER", "openai"), description="Model provider")
    enable_checkpointing: bool = Field(default=True, description="Enable state persistence")
    max_iterations: int = Field(default=10, description="Maximum agent iterations")
    recursion_limit: int = Field(default_factory=lambda: int(os.getenv("LANGRAPH_RECURSION_LIMIT", "100")), description="Maximum recursion limit for LangGraph")

def create_agent_graph(
    checkpointer: Optional[BaseCheckpointSaver] = None,
    config: Optional[AgentConfig] = None
):
    """
    Create the ReAct agent using LangGraph's built-in components with typed state.
    
    Args:
        checkpointer: Optional checkpointer for state persistence
        config: Optional agent configuration
        
    Returns:
        Compiled LangGraph with built-in state management
    """
    if config is None:
        config = AgentConfig()
    
    # Initialize StateGraph with built-in MessagesState (typed)
    workflow = StateGraph(MessagesState)
    
    # Add nodes using built-in components with validation
    workflow.add_node("agent", agent_node)
    workflow.add_node("tools", tool_node)
    
    # Add edges using built-in routing patterns
    workflow.add_edge(START, "agent")
    
    # Use built-in tools_condition with proper typing
    workflow.add_conditional_edges(
        "agent",
        tools_condition,  # Built-in condition with proper message type checking
        {
            "tools": "tools",
            END: END,
        }
    )
    
    # Tools return to agent for continued reasoning (ReAct pattern)
    workflow.add_edge("tools", "agent")
    
    # Compile with built-in error handling and validation
    compile_config = {}
    if checkpointer and config.enable_checkpointing:
        compile_config["checkpointer"] = checkpointer
    
    return workflow.compile(**compile_config)