"""
Agent with memory and confirmation of tool calls.

This script demonstrates how to create an agent using LangGraph with:
- Memory (MemorySaver) to keep conversation history.
- Interrupt-before to pause before calling a tool.
- Rich console for pretty printing.

Run:
    python agent_with_memory.py

Make sure to install dependencies:
    pip install -r requirements.txt
"""

import json
from typing import Any, Dict, List, Optional

# Rich console for pretty printing
from rich.console import Console

# LangGraph components
from langgraph import AgentBuilder
from langgraph.checkpoint.memory import MemorySaver

# LangChain components
from langchain_ollama import ChatOllama
from langchain_ollama import OllamaEmbeddings

# Tool definition
from langchain.tools import BaseTool

console = Console()

# ----- Tool definition -----
class GetPriceTool(BaseTool):
    name: str = "get_price"
    description: str = "Get the price of a product for a given city and date. Returns a string."

    def _run(self, city: str, date: str) -> str:
        # Dummy implementation – in real life, query an API
        return f"The price in {city} on {date} is $42."

    def _arun(self, city: str, date: str) -> str:  # async version
        return self._run(city, date)

# ----- Agent creation -----

def create_agent(
    model: Any,
    tools: List[BaseTool],
    system_prompt: str,
    checkpointer: MemorySaver,
    interrupt_before: Optional[List[str]] = None,
) -> Any:
    """Build and return a LangGraph agent.

    Parameters
    ----------
    model: The language model (e.g., ChatOllama).
    tools: List of tools the agent can use.
    system_prompt: System prompt for the agent.
    checkpointer: MemorySaver instance for conversation memory.
    interrupt_before: List of nodes to interrupt before (e.g., ['tools']).
    """
    builder = AgentBuilder(
        model=model,
        tools=tools,
        system_prompt=system_prompt,
        checkpointer=checkpointer,
    )
    if interrupt_before:
        builder = builder.with_interrupt_before(interrupt_before)
    return builder.build()

# ----- Conversation loop -----

def ask_and_run(
    agent: Any,
    user_input: Optional[Dict[str, Any]],
    config: Dict[str, Any],
) -> None:
    """Handle streaming from the agent, including pauses for tool confirmation.

    Parameters
    ----------
    agent: The LangGraph agent.
    user_input: The user message dict or None to resume.
    config: Configuration dict with thread_id.
    """
    # Prepare the input for the agent
    input_data = user_input if user_input is not None else None

    # Stream the agent's response
    for chunk in agent.stream(
        input_data, config=config, stream_mode=["messages", "updates"]
    ):
        chunk_type, chunk_data = chunk

        # Handle message chunks (token streaming)
        if chunk_type == "messages":
            # chunk_data is a list of messages; we print the last message content
            if chunk_data:
                last_msg = chunk_data[-1]
                if last_msg.get("role") == "assistant":
                    console.print(f"[bold cyan]Agent:[/bold cyan] {last_msg.get("content", "")}")

        # Handle updates (tool calls, etc.)
        if chunk_type == "updates":
            # chunk_data contains the updated state; we can inspect tool calls
            pass  # For simplicity, we ignore updates here

        # Detect interrupt before tool call
        if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
            # Agent is pausing before a tool call
            state = agent.get_state(config)
            # The last message should contain the tool call
            last_msg = state.values["messages"][-1]
            tool_call = last_msg["tool_calls"][0]
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            console.print(f"\n[bold yellow]Agent wants to call tool:[/bold yellow] {tool_name}({json.dumps(tool_args)})")
            answer = input("Разрешить? (Y/n): ")
            if answer.lower().strip() == "y":
                console.print("[green]Tool call allowed. Resuming...[/green]")
                # Recursively call ask_and_run to resume
                ask_and_run(agent, None, config)
            else:
                console.print("[red]Tool call cancelled.[/red]")
                break

if __name__ == "__main__":
    # Initialize the model and embeddings
    llm = ChatOllama(model="llama2")
    embeddings = OllamaEmbeddings(model="llama2")

    # Create memory saver
    memory = MemorySaver()

    # Define system prompt
    system_prompt = "You are a helpful assistant. Use the get_price tool to answer queries about prices."

    # Create the tool
    get_price_tool = GetPriceTool()

    # Build the agent
    agent = create_agent(
        model=llm,
        tools=[get_price_tool],
        system_prompt=system_prompt,
        checkpointer=memory,
        interrupt_before=["tools"],
    )

    # Conversation loop
    thread_id = "thread-1"
    config = {"configurable": {"thread_id": thread_id}}

    console.print("[bold green]Agent ready. Type 'exit' to quit.[/bold green]")
    while True:
        user_input = input("\nВы: ")
        if user_input.lower().strip() == "exit":
            console.print("[bold magenta]Goodbye![/bold magenta]")
            break
        # Wrap user input into a message dict
        message = {"messages": [{"role": "human", "content": user_input}]}
        ask_and_run(agent, message, config)