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

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command

# DESIGN DECISION: use OpenRouter LLM via ChatOpenAI
# NECESSITY: the assignment explicitly requires OpenRouter and forbids Ollama.
# OPTIMALITY: ChatOpenAI works with OpenAI compatible API, easy to configure base_url.
# ALTERNATIVES CONSIDERED: using other providers (e.g., Anthropic) - rejected because not allowed.

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,
)

# Simple example tool that returns a fake weather.
@tool
def get_weather(city: str) -> str:
    """Return a short weather description for the given city."""
    # In a real scenario this would call an external API.
    return f"The weather in {city} is sunny with a mild temperature."

# Memory saver is required for the middleware pause to be persisted.
memory = MemorySaver()

# DESIGN DECISION: use HumanInTheLoopMiddleware with interrupt_on for get_weather
# NECESSITY: required by the task to pause before tool execution.
# OPTIMALITY: middleware handles the interrupt generation and resume logic.
# ALTERNATIVES CONSIDERED: manual interrupt handling - more code, less reusable.

agent = create_agent(
    model=llm,
    tools=[get_weather],
    system_prompt="You are a helpful assistant.",
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={
                "get_weather": True,
            },
            description_prefix="Подтвердите вызов инструмента",
        ),
    ],
    checkpointer=memory,
)

def display_action_requests(action_requests: List[Dict[str, Any]]) -> None:
    for idx, action in enumerate(action_requests, start=1):
        print(f"\nAction {idx}:")
        print(f"  name: {action.get('name')}")
        print(f"  args: {action.get('args')}")
        description = action.get('description')
        if description:
            print(f"  description: {description}")

def collect_decisions(action_requests: List[Dict[str, Any]],
                      review_configs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    decisions: List[Dict[str, Any]] = []
    for action, config in zip(action_requests, review_configs):
        allowed = config.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
            if choice == "r" and "reject" in allowed:
                message = input("Message for rejection (sent to model): ").strip()
                decisions.append({"type": "reject", "message": message})
                break
            print("Invalid choice. Please enter a valid option.")
    return decisions

async def run_conversation():
    thread_id = "session-1"
    config = {"configurable": {"thread_id": thread_id}}

    # initial user message
    user_input = input("You: ").strip()
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=user_input)]},
        config,
    )

    # Loop while middleware interrupts
    while "__interrupt__" in result:
        interrupt_value = result["__interrupt__"][0].value
        action_requests = interrupt_value["action_requests"]
        review_configs = interrupt_value["review_configs"]

        print("\n--- Confirmation required ---")
        display_action_requests(action_requests)

        decisions = collect_decisions(action_requests, review_configs)

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

    # final response
    final_message = result["messages"][-1].content
    print(f"\nAssistant: {final_message}")

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