from typing import Dict, List
from pydantic import Field
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage
from orchestrator_maze_implementation.visualization.maze_visualization import display_maze_state
from orchestrator_maze_implementation.state.maze_state import MazeState


def maze_agent(state: MazeState):
    """
    Maze-solving agent that uses the LLM with tools - just the model part.
    """
    print(f"\nModel turn (message count: {len(state['messages'])})")
    
    # Display maze state with 0.5 second delay
    display_maze_state(state)
    
    # Get the first agent's maze wrapper and agent_id
    agent_id = list(state["maze_wrappers"].keys())[0]
    maze_wrapper = state["maze_wrappers"][agent_id]
    
    # Get current agent position
    agent_pos = maze_wrapper.get_agent_position()
    
    # Build messages - continue from previous conversation or start new
    if not state["messages"]:
        # First turn - initialize with system prompt
        system_prompt = SystemMessage(content="""You are an intelligent maze-solving agent
Your goal is to find the maze exit (E).You may get your current view via a tool call.
Maze symbols: O=Open path, W=Wall, E=Exit, A=Agent position, D=Marked dead end
You can use the tool to get your current view of the maze.
""")
        
        user_message = HumanMessage(content=f"""
You are navigating a maze. Your current position is {agent_pos}.

""")
        
        messages = [system_prompt, user_message]
    else:
        # Continue previous conversation
        messages = list(state["messages"])[-8:]
        
        # Add current turn context
        current_turn_message = HumanMessage(content=f"""
Continue navigating the maze. Your current position is {agent_pos}.
Focus on finding the exit (E) systematically.
""")
        
        messages.append(current_turn_message)
   
    # Get LLM response with tools
    try:
        print("\nMessages:")
        for msg in messages:
            print(f"  {msg.type}: {msg.content}...")

        
        response = llm_with_tools.invoke(messages)
        print(f"Agent thinking:| {response.content} {response.tool_calls}|")
        
        # Return the response as messages
        return {"messages": [response]}
        
    except Exception as e:
        print(f"Error getting LLM response: {e}")
        return {"messages": []}
