import os
import asyncio
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 и backend -----------------------------------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

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

# --- Tool with sub‑agent -----------------------------------------------------
@tool
def get_price(query: str) -> str:
    """Return a fake price table for a product in a city.
    Internally a sub‑agent is used to generate the answer.
    """
    # Создаём суб‑агента, который просто отвечает на запрос
    sub_agent = create_deep_agent(
        model=llm,
        tools=[],
        backend=backend,
        system_prompt="You are a price lookup assistant.",
    )
    # Запускаем суб‑агента в отдельном цикле событий
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        result = loop.run_until_complete(
            sub_agent.ainvoke(
                {"messages": [HumanMessage(content=f"Provide price for {query}")]},
                {"configurable": {"thread_id": f"price-{query}"}},
            )
        )
    finally:
        loop.close()
    return result["messages"][-1].content

# --- Главный агент ----------------------------------------------------------
agent = create_deep_agent(
    model=llm,
    tools=[get_price],
    backend=backend,
    system_prompt="You are a helpful agent. Use tools when necessary.",
)

# --- Stream ----------------------------------------------
stream = agent.stream(
    {
        "messages": [HumanMessage(content="What is the price of milk in Kazan?")],
    },
    stream_mode=["messages", "updates"],
)

# --- Handlers ------------------------------------------
step = 1

def format_chunk_message(chunk):
    global step
    message, meta = chunk
    if meta.get("langgraph_step") != step:
        step = meta.get("langgraph_step")
        print("\n--- --- ---\n")
    if message.content:
        print(message.content, end="", flush=True)


def format_message(message):
    if message.content:
        return message.content
    # tool call representation
    if message.tool_calls:
        call = message.tool_calls[0]
        return f"{call['name']}({call['args']})"
    return ""

for chunk in stream:
    chunk_type, chunk_data = chunk
    if chunk_type == "messages":
        format_chunk_message(chunk_data)
    elif chunk_type == "updates":
        if chunk_data.get("model"):
            last_msg = chunk_data["model"]["messages"][-1]
            print(format_message(last_msg), end="", flush=True)

# --- End of script ----------------------------------------------------------