"""
Human‑in‑the‑Loop example using LangChain `HumanInTheLoopMiddleware`.

This repository demonstrates how to pause an agent when a tool is about to be called, ask the user for approval (or rejection), and then resume execution.

The script contains three independent examples:
1. Simple weather query – single tool call.
2. Multiple tool calls in one conversation – shows that the loop continues until all tools are approved.
3. Rejection path – demonstrates how a rejected action is handled by the agent.

Run with:
    python -m venv .venv && source .venv/bin/activate
    pip install -r requirements.txt
    export JOURNAL_MCP_PAT=YOUR_BROJS_TOKEN
    python main.py
"""

import os
import json
from typing import List, Dict

from langchain_openai import ChatOpenAI
from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command

# ---------------------------------------------------------------------------
# 1. LLM configuration – BroJS only
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
    api_key=os.getenv("JOURNAL_MCP_PAT"),
    temperature=0.0,
)

# ---------------------------------------------------------------------------
# 2. Simple tool – get_weather (mocked for demo purposes)
# ---------------------------------------------------------------------------
def get_weather(city: str, date: str = "today") -> str:
    """Return a fabricated weather report.

    Parameters
    ----------
    city: str
        Name of the city.
    date: str, optional
        Date for which to fetch the forecast. Defaults to ``today``.
    """
    return f"The weather in {city} on {date} is sunny with a high of 25°C."

# ---------------------------------------------------------------------------
# 3. Agent construction – middleware pauses before calling get_weather
# ---------------------------------------------------------------------------
memory = MemorySaver()
agent = create_agent(
    model=llm,
    tools=[get_weather],
    system_prompt="You are a helpful assistant that can provide weather information.",
    middleware=[
        HumanInTheLoopMiddleware(
            interrupt_on={"get_weather": True},  # allow approve / reject / edit
            description_prefix="Please confirm the tool call:",
        ),
    ],
    checkpointer=memory,
)

# ---------------------------------------------------------------------------
# Helper: run a single turn with HIL loop
# ---------------------------------------------------------------------------
def run_turn(user_msg: str, thread_id: str) -> str:
    """Run one user message through the agent.

    The function handles all pauses caused by the middleware and asks the
    user for approval or rejection.  It returns the final assistant reply.
    """
    config = {"configurable": {"thread_id": thread_id}}
    result: Dict = agent.invoke({"messages": [{"role": "human", "content": user_msg}]}, config=config)

    while "__interrupt__" in result:
        interrupt_value = result["__interrupt__"][0].value
        action_requests = interrupt_value.get("action_requests", [])
        review_configs = interrupt_value.get("review_configs", {})

        decisions: List[Dict] = []
        for idx, act in enumerate(action_requests):
            name = act.get("name")
            args = act.get("args", {})
            description = act.get("description", "")
            allowed = review_configs.get(name, {}).get("allowed_decisions", ["approve", "reject", "edit"])

            print(f"\n--- Tool call {idx + 1} ---")
            print(f"Name: {name}")
            print(f"Arguments: {json.dumps(args)}")
            if description:
                print(f"Description: {description}")
            opts = ["a" for _ in allowed if "approve" in _]
            opts += ["r" for _ in allowed if "reject" in _]
            opts += ["e" for _ in allowed if "edit" in _]
            opt_str = "/".join(opts)
            choice = input(f"Choose {opt_str}: ").strip().lower()

            if choice == "a":
                decisions.append({"type": "approve"})
            elif choice == "r":
                msg = input("Reason for rejection: ")
                decisions.append({"type": "reject", "message": msg})
            elif choice == "e":
                new_args_raw = input("Enter edited arguments as JSON: ")
                try:
                    new_args = json.loads(new_args_raw)
                except Exception:
                    print("Invalid JSON – falling back to original.")
                    new_args = args
                decisions.append({"type": "edit", "edited_action": {"name": name, "args": new_args}})
            else:
                print("Unrecognised choice – defaulting to reject.")
                decisions.append({"type": "reject", "message": "User did not provide valid input."})

        result = agent.invoke(Command(resume={"decisions": decisions}), config=config)

    final_msg = result["messages"][-1].content
    return final_msg

# ---------------------------------------------------------------------------
# Demo – three independent examples
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    print("=== Example 1: Simple weather query ===")
    reply = run_turn("Какая погода в Казани сегодня?", thread_id="session-1")
    print(f"Assistant: {reply}\n")

    print("=== Example 2: Multiple tool calls in one conversation ===")
    reply = run_turn(
        "Сколько будет в Казани завтра и как погода в Москве сегодня?", thread_id="session-2"
    )
    print(f"Assistant: {reply}\n")

    print("=== Example 3: Rejection path ===")
    reply = run_turn("Погода в Лондоне на завтра, пожалуйста.", thread_id="session-3")
    print(f"Assistant: {reply}\n")

    print("All examples finished.\n")