import uuid
from typing import Any, Dict, List

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
from langchain.tools import tool


def main() -> None:
    llm = ChatOpenAI(
        model="gpt-3.5-turbo",
        temperature=0,
        # For a local model, uncomment and adjust:
        # base_url="http://localhost:11434/v1",
        # api_key="ollama",
    )

    @tool("get_weather")
    def get_weather(city: str, date: str) -> str:
        """Return a mocked weather string for the given city and date."""
        return f"Weather in {city} on {date}"

    tools = [get_weather]

    memory = MemorySaver()

    agent = create_agent(
        model=llm,
        tools=tools,
        system_prompt="You are a helpful assistant",
        middleware=[
            HumanInTheLoopMiddleware(
                interrupt_on={"get_weather": True},
                description_prefix="Confirm tool usage:",
            ),
        ],
        checkpointer=memory,
    )

    thread_id = str(uuid.uuid4())
    config: Dict[str, Any] = {"configurable": {"thread_id": thread_id}}

    user_message = input("Enter your question: ")

    result = agent.invoke(
        {"messages": [{"role": "human", "content": user_message}]},
        config=config,
    )

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

        decisions: List[Dict[str, Any]] = []

        for idx, action in enumerate(action_requests):
            name = action.get("name", "<unknown>")
            args = action.get("args", {})
            description = action.get("description", "")
            print(f"\nAction {idx + 1}: {name}")
            if description:
                print(f"Description: {description}")
            print(f"Arguments: {args}")
            while True:
                choice = input("Decision (a=approve, r=reject): ").strip().lower()
                if choice == "a":
                    decisions.append({"type": "approve"})
                    break
                elif choice == "r":
                    reason = input("Reason for rejection: ")
                    decisions.append({"type": "reject", "message": reason})
                    break
                else:
                    print("Invalid input. Please enter 'a' or 'r'.")

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

    final_content = result["messages"][-1]["content"]
    print("\nAgent response:")
    print(final_content)


if __name__ == "__main__":
    main()