import os
import asyncio
from typing import List, Dict, Any

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, BaseMessage
from langchain.tools import tool

from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# ----------------------------------------------------------------------
# Configuration
# ----------------------------------------------------------------------
# LLM - OpenRouter (free tier).  The API key must be stored in the environment.
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.7,
)

# Backend for the agents - a simple composite that allows file operations
# and execution of shell commands inside a sandboxed workspace.
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# ----------------------------------------------------------------------
# Sub-agent: price generator
# ----------------------------------------------------------------------
def _create_price_subagent() -> Any:
    """
    Creates a lightweight sub-agent that, given a product and a city,
    returns a markdown table with a plausible price and a store name.
    The sub-agent re-uses the same LLM and backend as the main agent.
    """
    subagent = create_deep_agent(
        model=llm,
        tools=[],  # No additional tools are required for price generation
        backend=backend,
        system_prompt=(
            "You are a price-estimation sub-agent. "
            "Given a product name and a city, generate a realistic price "
            "in Russian rubles and suggest a typical store. "
            "Return the result as a markdown table with columns: "
            "`Продукт`, `Цена (руб.)`, `Магазин`."
        ),
    )
    return subagent

_price_subagent = _create_price_subagent()

@tool
def get_price(product: str, city: str) -> str:
    """
    Estimate the price of a product in a given city.
    The function creates a sub-agent that returns a markdown table:
    | Продукт | Цена (руб.) | Магазин |
    """
    # Build the prompt for the sub-agent
    prompt = HumanMessage(
        content=f"Продукт: {product}\nГород: {city}\nСгенерируй цену."
    )
    # Invoke the sub-agent asynchronously and wait for the result
    result = asyncio.run(
        _price_subagent.ainvoke(
            {"messages": [prompt]},
            {"configurable": {"thread_id": f"price-{product}-{city}"}},
        )
    )
    # The sub-agent returns a list of messages; the last one contains the table
    final_message = result["messages"][-1]
    return final_message.content if isinstance(final_message, BaseMessage) else str(final_message)

# ----------------------------------------------------------------------
# Main agent: shopping list planner
# ----------------------------------------------------------------------
main_agent = create_deep_agent(
    model=llm,
    tools=[get_price],
    backend=backend,
    system_prompt="Ты помощник по планированию покупок.",
)

def format_message(msg: BaseMessage) -> str:
    """
    Convert a LangChain message to a readable string.
    Handles normal text messages and tool calls.
    """
    if hasattr(msg, "content") and msg.content:
        return msg.content
    # Tool call representation
    if hasattr(msg, "tool_calls") and msg.tool_calls:
        call = msg.tool_calls[0]
        name = call["name"]
        args = ", ".join(f"{k}={v!r}" for k, v in call["args"].items())
        return f"{name}({args})"
    return str(msg)

async def main() -> None:
    user_query = (
        "Помоги составить список покупок: молоко, хлеб, яблоки. "
        "Я нахожусь в Казани."
    )
    result = await main_agent.ainvoke(
        {"messages": [HumanMessage(content=user_query)]},
        {"configurable": {"thread_id": "shopping-session-1"}},
    )
    # Print the whole conversation chain
    for i, msg in enumerate(result["messages"], start=1):
        print(f"--- Message {i} ---")
        print(format_message(msg))
        print()

if __name__ == "__main__":
    asyncio.run(main())