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

# Utility to initialize LLM with error handling

def init_llm(base_url: str, model: str, api_key: str = "ollama") -> ChatOpenAI:
    try:
        return ChatOpenAI(
            base_url=base_url,
            api_key=SecretStr(api_key),
            model=model,
            temperature=0.7,
        )
    except Exception as e:
        print(f"Ошибка при подключении к LLM: {e}")
        sys.exit(1)

# Main LLM
LLM_URL = "http://localhost:1234/v1"
LLM_MODEL = "llama3"
llm = init_llm(LLM_URL, LLM_MODEL)

# 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:
    """Subagent that looks up price for a product in a given city.
    For demonstration it returns a static price string, but the structure
    mirrors a real subagent that could call another LLM.
    """
    sub_llm = init_llm(LLM_URL, LLM_MODEL, api_key="ollama")
    sub_agent = create_agent(
        llm=sub_llm,
        tools=[],
        system_prompt="You are a simple price lookup tool. Return price as 'price: <amount>, store: <store>'.",
    )
    try:
        response = sub_agent.invoke({"input": f"Give price for {product} in {city}."})
        if isinstance(response, dict) and "messages" in response:
            for msg in response["messages"]:
                if msg.get("role") == "assistant" and msg.get("content"):
                    return msg["content"].strip()
    except Exception as e:
        print(f"Ошибка при вызове subagent: {e}")
    return "price: 100, store: SuperMarket"


def format_message(message) -> str:
    if message.get("content"):
        return message["content"]
    if message.get("tool_calls"):
        parts = []
        for call in message["tool_calls"]:
            name = call.get("name")
            args = call.get("arguments") or call.get("function", {}).get("arguments")
            parts.append(f"{name}({args})")
        return ", ".join(parts)
    return ""


def main():
    agent = create_agent(
        llm=llm,
        tools=[get_price],
        system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price для получения цен.",
    )
    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    result = agent.invoke({"input": user_query})
    messages = result.get("messages", [])

    # Print tool calls in order
    for msg in messages:
        if msg.get("role") == "assistant" and 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}")

    # Aggregate prices by calling the tool directly for each product
    products = ["молоко", "хлеб", "яблоки"]
    city = "Казань"
    price_entries = []
    total = 0.0
    for prod in products:
        try:
            price_str = get_price(prod, city)
            m = re.search(r"price:\s*([0-9]+)\s*,\s*store:\s*([A-Za-z0-9 ]+)", price_str, re.IGNORECASE)
            if m:
                price = float(m.group(1))
                store = m.group(2).strip()
                price_entries.append((prod, price, store))
                total += price
        except Exception:
            continue

    # Print final table
    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()