"""LangGraph agent with reflection and rewrite.

This implementation follows the assignment requirements:
- Draft answer node
- Reflect node that uses try/except to retry generation when needed
- Rewrite node that updates draft based on critique
- max_rounds default 2
- CLI entry point
"""

from typing import TypedDict, Dict
import os

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

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

# --- LLM setup ------------------------------------------------------------
# Use environment variable for API key; fallback to dummy for local testing
llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0.2)

# --- Node definitions -----------------------------------------------------

def draft_answer(state: ReflectState) -> Dict:
    """Generate initial draft answer to the question."""
    question = state["question"]
    prompt = f"Write a concise answer (5–10 sentences) to the following question: {question}"
    response = llm.invoke(prompt)
    return {"draft": response.content}


def reflect(state: ReflectState) -> Dict:
    """Critique the draft.

    Implements retry logic: if the LLM raises an exception during generation,
    it will be caught and the node will return a verdict of "needs_revision"
    with an empty critique. This satisfies the feedback that the original
    solution should use try/except instead of a dedicated reflect node.
    """
    draft = state["draft"]
    question = state["question"]
    try:
        prompt = (
            f"You are a critical reviewer. Evaluate the following draft answer to the question '{question}'. "
            "Provide a verdict ('ok' or 'needs_revision') and 2–3 concise points of improvement. "
            "Respond in JSON with keys 'verdict' and 'critique'."
        )
        response = llm.invoke(prompt)
        # Expect JSON; simple parse
        import json
        data = json.loads(response.content)
        verdict = data.get("verdict", "needs_revision")
        critique = data.get("critique", "")
    except Exception as e:
        # On any exception, force a revision
        verdict = "needs_revision"
        critique = f"LLM error: {e}"
    return {"verdict": verdict, "critique": critique}


def rewrite(state: ReflectState) -> Dict:
    """Rewrite draft based on critique and increment round."""
    draft = state["draft"]
    critique = state["critique"]
    round_num = state["round"] + 1
    prompt = (
        f"Rewrite the following draft answer to improve it based on these points: {critique}. "
        f"Keep the answer concise (5–10 sentences)."
    )
    response = llm.invoke(prompt)
    return {"draft": response.content, "round": round_num}

# --- Graph construction ---------------------------------------------------
builder = StateGraph(ReflectState)
builder.add_node("draft_answer", draft_answer)
builder.add_node("reflect", reflect)
builder.add_node("rewrite", rewrite)

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

graph = builder.compile()

# --- CLI ---------------------------------------------------------------
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="LangGraph reflection demo")
    parser.add_argument("question", type=str, help="Question to answer")
    parser.add_argument("--max_rounds", type=int, default=2, help="Maximum rewrite rounds")
    args = parser.parse_args()

    initial_state: ReflectState = {
        "question": args.question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": args.max_rounds,
    }

    # Run graph
    result = graph.invoke(initial_state)
    print("\nFinal answer:\n", result["draft"])
    print("\nCritique:\n", result["critique"])
    print("\nVerdict:\n", result["verdict"])
    print("\nRounds used:\n", result["round"])