from langchain_openai import ChatOpenAI
from langchain.agents import create_agent
from langchain.tools import tool
import re
import json

# Initialize main LLM
llm = ChatOpenAI(
    base_url="http://localhost:11434/v1",
    api_key="ollama",
    model="<название модели в LM Studio>",
    temperature=0.7,
)

# Define get_price tool with subagent
@tool(name="get_price", description="Get price for a product in a city.")
def get_price(product: str, city: str) -> str:
    # For demonstration, use a subagent that simply returns a static price
    sub_llm = ChatOpenAI(
        base_url="http://localhost:11434/v1",
        api_key="ollama",
        model="<название модели в LM Studio>",
        temperature=0.2,
    )
    sub_agent = create_agent(
        llm=sub_llm,
        tools=[],
        system_prompt="You are a simple price lookup tool. Return price as 'price: <amount>, store: <store>'.",
        verbose=False,
    )
    response = sub_agent.invoke({"input": f"Give price for {product} in {city}."})
    # The sub_agent will return a dict with messages; extract content
    if isinstance(response, dict) and "messages" in response:
        # Take the first assistant message content
        for msg in response["messages"]:
            if msg.get("role") == "assistant" and msg.get("content"):
                return msg["content"].strip()
    # Fallback to static response if sub_agent fails
    return "price: 100, store: SuperMarket"


def main():
    # Create main agent
    agent = create_agent(
        llm=llm,
        tools=[get_price],
        system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цен.",
        verbose=False,
    )
    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    result = agent.invoke({"input": user_query})
    messages = result.get("messages", [])
    # Print each message and tool calls
    for msg in messages:
        if msg.get("role") == "assistant":
            if msg.get("content"):
                print(msg["content"])
            if msg.get("tool_calls"):
                for call in msg["tool_calls"]:
                    name = call.get("name")
                    args = call.get("arguments") or call.get("function", {}).get("arguments")
                    print(f"Tool call: {name} with args {args}")
    # Build table from tool calls
    price_entries = []
    total = 0.0
    for msg in messages:
        if msg.get("role") == "assistant" and msg.get("tool_calls"):
            for call in msg["tool_calls"]:
                if call.get("name") == "get_price":
                    arg_str = call.get("arguments") or call.get("function", {}).get("arguments")
                    try:
                        args = json.loads(arg_str)
                        product = args.get("product")
                    except Exception:
                        product = "unknown"
                    result_text = call.get("function", {}).get("arguments", "")
                    m = re.search(r"price:\s*([0-9]+)\s*,\s*store:\s*([A-Za-z0-9 ]+)", result_text, re.IGNORECASE)
                    if m:
                        price = float(m.group(1))
                        store = m.group(2).strip()
                        price_entries.append((product, price, store))
                        total += price
    print("\nТаблица цен:")
    print("{:<15} {:<10} {:<15}".format("Товар", "Цена", "Магазин"))
    for prod, price, store in price_entries:
        print("{:<15} {:<10} {:<15}".format(prod, f"{price} руб", store))
    print(f"\nОбщая стоимость: {total} руб")

if __name__ == "__main__":
    main()