#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# DESIGN DECISION: deepagents is required by the course assignment to build the agent.
# NECESSITY: The assignment explicitly requires using create_deep_agent from deepagents; without it the agent cannot be instantiated.
# OPTIMALITY: Using deepagents ensures consistent agent behavior and simplifies tool integration; alternative frameworks would violate the course constraints.
# ALTERNATIVES CONSIDERED: Using plain langgraph without deepagents would miss the required framework; manually handling tool calls would increase boilerplate.

import os
import asyncio
import argparse
from typing import TypedDict, Annotated

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.graph import StateGraph, START, END
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field

# LLM configuration - 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(),
    ]
)

# ---------- LangGraph components ----------

class ReflectState(TypedDict):
    question: str
    draft: str
    critique: str
    verdict: str  # ok | needs_revision
    round: int
    max_rounds: int

class CritiqueOutput(BaseModel):
    verdict: str = Field(description="ok or needs_revision")
    critique: str = Field(description="2-3 bullet points of critique")

critique_parser = PydanticOutputParser(pydantic_object=CritiqueOutput)

def draft_answer(state: dict) -> dict:
    prompt = f"Write a brief answer (5-10 sentences) to the following question:\n\n{state['question']}"
    result = llm.invoke([HumanMessage(content=prompt)])
    state["draft"] = result.content.strip()
    return state

def reflect(state: dict) -> dict:
    prompt = (
        f"You are a critic. Evaluate the following draft answer for completeness, specificity, and lack of fluff.\n\nDraft:\n{state['draft']}\n\n"
        "Respond with JSON containing 'verdict' ('ok' or 'needs_revision') and 'critique' (2-3 bullet points)."
    )
    result = llm.invoke([HumanMessage(content=prompt)])
    critique = critique_parser.parse(result.content)
    state["verdict"] = critique.verdict
    state["critique"] = critique.critique
    return state

def rewrite(state: dict) -> dict:
    prompt = (
        f"Rewrite the draft answer to address the following critique:\n\n{state['critique']}\n\n"
        "Keep the answer brief (5-10 sentences)."
    )
    result = llm.invoke([HumanMessage(content=prompt)])
    state["draft"] = result.content.strip()
    state["round"] += 1
    return state

def build_graph() -> StateGraph:
    graph = StateGraph(ReflectState)
    graph.add_node("draft_answer", draft_answer)
    graph.add_node("reflect", reflect)
    graph.add_node("rewrite", rewrite)

    graph.add_edge(START, "draft_answer")
    graph.add_edge("draft_answer", "reflect")

    def decide_next(state: dict) -> str:
        if state["verdict"] == "needs_revision" and state["round"] < state["max_rounds"]:
            return "rewrite"
        return END

    graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", "END": END})
    graph.add_edge("rewrite", "reflect")

    return graph.compile()

def run_graph(question: str) -> str:
    graph = build_graph()
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    final_state = graph.run(initial_state)
    return final_state["draft"]

# ---------- DeepAgents tool ----------

@tool
def run_graph_tool(question: str) -> str:
    """Run the LangGraph to produce a refined answer."""
    return run_graph(question)

# ---------- Agent ----------

agent = create_deep_agent(
    model=llm,
    tools=[run_graph_tool],
    backend=backend,
    system_prompt="You are a helpful agent that answers questions by running the run_graph tool.",
)

# ---------- CLI ----------

async def main():
    parser = argparse.ArgumentParser(description="Answer a question with self-reflection.")
    parser.add_argument("question", nargs="*", help="The question to answer.")
    args = parser.parse_args()
    if not args.question:
        question = input("Enter your question: ").strip()
    else:
        question = " ".join(args.question).strip()

    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=question)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    # The agent will return the final answer in the last message
    final_message = result["messages"][-1].content
    print("\nAnswer:\n")
    print(final_message)

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