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
from langgraph.checkpoint.memory import MemorySaver
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.types import Command

# LLM via 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.0,
)

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

# Example tool – get_weather
@tool
def get_weather(city: str, date: str) -> str:
    """Return the weather for a given city and date."""
    # Dummy implementation – replace with real API call if needed
    return f"Weather in {city} on {date} is sunny and 25°C."

# Create the agent with Human‑in‑the‑Loop middleware
agent = create_deep_agent(
    model=llm,
    tools=[get_weather],
    backend=backend,
    system_prompt="You are a helpful assistant.",
    checkpointer=MemorySaver(),
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={"get_weather": True},
            description_prefix="Подтвердите вызов инструмента",
        ),
    ],
)

async def main():
    config = {"configurable": {"thread_id": "session-1"}}
    user_input = input("Вы: ")

    # Initial invoke
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=user_input)]},
        config,
    )

    # Process Human‑in‑the‑Loop interrupts
    while "__interrupt__" in result:
        interrupt = result["__interrupt__"][0].value
        action_requests = interrupt.get("action_requests", [])
        review_configs = interrupt.get("review_configs", {})

        decisions = []
        print("\n--- Подтверждение ---")
        for action in action_requests:
            name = action.get("name")
            args = action.get("args", {})
            description = action.get("description", "")
            print(f"Инструмент: {name}")
            print(f"Аргументы: {args}")
            if description:
                print(f"Описание: {description}")
            allowed = review_configs.get(name, {}).get("allowed_decisions", ["approve", "reject"])
            while True:
                choice = input("a = approve, r = reject: ").strip().lower()
                if choice == "a" and "approve" in allowed:
                    decisions.append({"type": "approve"})
                    break
                elif choice == "r" and "reject" in allowed:
                    reason = input("Сообщение для агента (причина отказа): ").strip()
                    decisions.append({"type": "reject", "message": reason})
                    break
                else:
                    print("Неверный ввод. Попробуйте снова.")

        # Resume agent after decisions
        result = await agent.ainvoke(
            Command(resume={"decisions": decisions}),
            config,
        )

    # Final answer
    final_message = result["messages"][-1].content
    print("\nАгент:", final_message)

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