import asyncio
import os
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# --- LLM configuration -----------------------------------------------------
# Connect to the local LM Studio server. Replace '<model_name>' with the exact
# name of the model you have loaded in LM Studio.
llm = ChatOpenAI(
    model='<model_name>',
    base_url='http://localhost:1234/v1',
    api_key=os.getenv('OPENAI_API_KEY', 'fake'),
    temperature=0.7,
)

# --- Backend ---------------------------------------------------------------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# --- Sub‑agent tool --------------------------------------------------------
@tool
def get_price(product: str, city: str) -> str:
    """Return a realistic price for a product in a given city.

    The function internally creates a sub‑agent that asks the LLM to generate
    a price table. The sub‑agent is a lightweight wrapper around the same
    LLM instance to keep the example simple.
    """
    # Create a sub‑agent that only has the task of generating a price table.
    sub_agent = create_deep_agent(
        model=llm,
        tools=[],
        backend=backend,
        system_prompt=f"You are a market analyst. Provide a realistic price for {product} in {city}. Output a markdown table with columns: Продукт, Цена (руб.), Магазин.",
    )
    # Invoke the sub‑agent with a simple prompt.
    result = asyncio.run(
        sub_agent.ainvoke(
            {"messages": [HumanMessage(content=f"Generate price for {product} in {city}")]},
            {"configurable": {"thread_id": f"price-{product}-{city}"}},
        )
    )
    # Return the content of the last message (the table).
    return result["messages"][-1].content

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

# --- Helper to pretty‑print the conversation ------------------------------
from langchain_core.messages import BaseMessage

def format_message(msg: BaseMessage) -> str:
    if hasattr(msg, "content") and msg.content:
        return msg.content
    if hasattr(msg, "tool_calls") and msg.tool_calls:
        call = msg.tool_calls[0]
        return f"{call['name']}({call['args']})"
    return ""

# --- Main entry point ------------------------------------------------------
async def main():
    user_prompt = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    result = await main_agent.ainvoke(
        {"messages": [HumanMessage(content=user_prompt)]},
        {"configurable": {"thread_id": "shopping-session"}},
    )
    # Print all messages in order
    for msg in result["messages"]:
        print(format_message(msg))
        print("---")

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