"""LangGraph reflection agent example.

The agent writes a short answer, critiques it, and rewrites if needed.
"""

from typing import TypedDict, Dict
import sys

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI

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

# LLM
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# Draft node
async def draft_answer(state: ReflectState) -> Dict:
    prompt = f"Write a short answer (5–10 sentences) to the following question: {state['question']}"
    response = await llm.ainvoke(prompt)
    state['draft'] = response.content
    return state

# Critique node
async def reflect(state: ReflectState) -> Dict:
    prompt = (
        f"You are a critic. Evaluate the following draft answer for completeness, specificity, and lack of filler.\n"
        f"Draft: {state['draft']}\n"
        f"Give a verdict: 'ok' or 'needs_revision'.\n"
        f"If needs_revision, provide 2–3 points of critique."
    )
    response = await llm.ainvoke(prompt)
    # Simple parsing: first line verdict, rest critique
    lines = response.content.strip().splitlines()
    verdict = lines[0].strip().lower()
    critique = "\n".join(lines[1:]).strip()
    state['verdict'] = verdict
    state['critique'] = critique
    return state

# Rewrite node
async def rewrite(state: ReflectState) -> Dict:
    prompt = (
        f"Rewrite the draft answer taking into account the following critique: {state['critique']}\n"
        f"Original draft: {state['draft']}"
    )
    response = await llm.ainvoke(prompt)
    state['draft'] = response.content
    state['round'] += 1
    return state

# Build graph
builder = StateGraph(ReflectState)
builder.add_node("draft_answer", draft_answer)
builder.add_node("reflect", reflect)
builder.add_node("rewrite", rewrite)

builder.set_entry_point("draft_answer")
builder.add_edge("draft_answer", "reflect")
builder.add_conditional_edges(
    "reflect",
    lambda state: "END" if state["verdict"] == "ok" else "rewrite",
)
builder.add_edge("rewrite", "reflect")

graph = builder.compile()

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python main.py <question>")
        sys.exit(1)
    question = sys.argv[1]
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    result = graph.invoke(initial_state)
    print("\nFinal answer:\n", result["draft"])
    print("\nCritique:\n", result["critique"])
    print("\nVerdict:\n", result["verdict"])