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

from pydantic import SecretStr
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend

# ----------------------------------------------------------------------
# LLM configuration (OpenRouter)
# ----------------------------------------------------------------------
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 sub-agents (allows file operations and shell commands)
# ----------------------------------------------------------------------
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# ----------------------------------------------------------------------
# Sub-agent that generates a realistic price table for a product
# ----------------------------------------------------------------------
def create_price_subagent() -> Any:
    """
    Returns a deep agent that, given a product and a city, produces a markdown
    table with product, price and store. The prompt forces the model to fabricate
    plausible data based on typical market prices.
    """
    system_prompt = (
        "You are a price-generation sub-agent. Given a product name and a city, "
        "return a markdown table with columns: Продукт, Цена (руб.), Магазин. "
        "Fabricate realistic prices based on typical Russian market data. "
        "Do not add any extra commentary, only the table."
    )
    subagent = create_deep_agent(
        model=llm,
        tools=[],  # no external tools needed for this simple sub-agent
        backend=backend,
        system_prompt=system_prompt,
    )
    return subagent

price_subagent = create_price_subagent()

# ----------------------------------------------------------------------
# Tool that calls the sub-agent
# ----------------------------------------------------------------------
@tool
def get_price(product: str, city: str) -> str:
    """
    Generate a realistic price for the given product in the specified city.
    Returns a markdown table with columns: Продукт, Цена (руб.), Магазин.
    """
    # Build the prompt for the sub-agent
    prompt = f"Продукт: {product}\nГород: {city}"
    # Invoke the sub-agent synchronously (deepagents also supports async,
    # but a simple sync call keeps the example straightforward)
    result = asyncio.run(
        price_subagent.ainvoke(
            {"messages": [HumanMessage(content=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]
    if isinstance(final_message, AIMessage):
        return final_message.content
    elif isinstance(final_message, ToolMessage):
        return final_message.content
    else:
        return str(final_message)

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

def format_message(msg: Any) -> str:
    """Human-readable representation of a message or tool call."""
    if isinstance(msg, (HumanMessage, AIMessage)):
        return msg.content
    if isinstance(msg, ToolMessage):
        return f"{msg.name}({msg.args}) -> {msg.content}"
    # Fallback for generic dict-like messages
    if hasattr(msg, "tool_calls") and msg.tool_calls:
        call = msg.tool_calls[0]
        return f"{call['name']}({call['args']})"
    return str(msg)

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

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