"""
Simple hierarchical AI agent for shopping list using LangChain.

The script demonstrates:
* Connection to a local LLM via OpenAI-compatible API.
* A tool that internally creates a sub‑agent to generate realistic prices.
* A top‑level agent that orchestrates the price queries and aggregates results.
"""

from __future__ import annotations

import json
from typing import Dict, Any

from langchain_openai import ChatOpenAI
from langchain.tools import tool
from langchain.agents import create_agent
from pydantic import SecretStr

# ---------------------------------------------------------------------------
# 1. LLM configuration – adjust model name to your LM Studio instance.
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
    model="gpt-4o-mini",  # replace with the actual model name in LM Studio
    base_url="http://localhost:1234/v1",
    api_key=SecretStr("fake"),
    temperature=0.7,
)

# ---------------------------------------------------------------------------
# 2. Sub‑agent that generates a price table for a single product.
# ---------------------------------------------------------------------------
@tool
def get_price(product: str, city: str) -> str:
    """Return a realistic price table for *product* in *city*.

    The function internally creates a small LangChain agent that asks the LLM to
    produce a markdown table with columns Product | Price (руб.) | Store.
    """
    # Create a tiny agent that only has the task of generating a price row.
    sub_agent = create_agent(
        model=llm,
        tools=[],
        system_prompt=(
            f"You are an assistant that provides realistic prices for products in {city}. "
            "Respond with a markdown table containing columns: Product, Price (руб.), Store."
        ),
    )

    # Ask the sub‑agent to produce the price.
    response = sub_agent.invoke(
        {
            "messages": [
                {"role": "human", "content": f"Provide a price for {product} in {city}."}
            ]
        }
    )

    # Extract the final message content.
    messages = response.get("messages", [])
    if not messages:
        return f"| {product} | N/A | Unknown |
"
    last_msg = messages[-1]
    content = last_msg.get("content") or ""
    # Ensure the table has a header row.
    if "|" in content and "Product" not in content:
        content = f"| Product | Price (руб.) | Store |\n{content}"
    return content

# ---------------------------------------------------------------------------
# 3. Main agent that orchestrates price queries for a shopping list.
# ---------------------------------------------------------------------------
shopping_agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="You are an assistant that helps plan a shopping list and calculates total cost.",
)

def main() -> None:
    user_query = (
        "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    )
    result = shopping_agent.invoke({"messages": [{"role": "human", "content": user_query}]})

    # Pretty‑print all messages.
    for msg in result.get("messages", []):
        if msg.get("content"):
            print(msg["content"])
        elif msg.get("tool_calls"):
            call = msg["tool_calls"][0]
            print(f"{call['name']}({json.dumps(call['args'], ensure_ascii=False)})")

if __name__ == "__main__":
    main()