"""
LangGraph Reflective Agent
=========================

This repository contains a small demo of a LangGraph agent that:

* Generates an initial answer to a question.
* Critiques the answer using a separate node.
* If the critique indicates "needs_revision", rewrites the answer up to ``max_rounds`` times.

The implementation follows the specification from the assignment and is fully runnable with
``pip install -r requirements.txt``.
"""

import os
from typing import TypedDict, Annotated
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from dotenv import load_dotenv

# Load environment variables (JOURNAL_MCP_PAT must be set)
load_dotenv()

# LLM configuration – BroJS endpoint
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,
)

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

# ---------- Node implementations ---------------------------------------
async def draft_answer(state: ReflectState) -> dict:
    """Generate a concise answer (5–10 sentences)."""
    prompt = (
        f"Write a short answer (5-10 sentences) to the following question:\n\n{state['question']}"
    )
    response = await llm.ainvoke([{"role": "user", "content": prompt}])
    state["draft"] = response.content.strip()
    return {"draft": state["draft"]}

async def reflect(state: ReflectState) -> dict:
    """Critique the draft and decide if revision is needed."""
    critique_prompt = (
        f"You are a critical reviewer. Evaluate the following answer for completeness, specificity, and lack of filler.\n\nAnswer:\n{state['draft']}\n\nProvide verdict (ok / needs_revision) followed by 2-3 bullet points of feedback."
    )
    response = await llm.ainvoke([{"role": "user", "content": critique_prompt}])
    # Parse verdict and critique
    text = response.content.strip()
    if "needs_revision" in text.lower():
        state["verdict"] = "needs_revision"
    else:
        state["verdict"] = "ok"
    state["critique"] = text
    return {"critique": state["critique"], "verdict": state["verdict"]}

async def rewrite(state: ReflectState) -> dict:
    """Rewrite the draft incorporating critique feedback."""
    rewrite_prompt = (
        f"You are revising an answer based on the following critique. Update the answer to improve it, keeping it concise (5-10 sentences).\n\nCritique:\n{state['critique']}\n\nOriginal Answer:\n{state['draft']}"
    )
    response = await llm.ainvoke([{"role": "user", "content": rewrite_prompt}])
    state["draft"] = response.content.strip()
    state["round"] += 1
    return {"draft": state["draft"], "round": state["round"]}

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

# Define transitions
builder.set_entry_point("draft_answer")
builder.add_edge("draft_answer", "reflect")
# From reflect: if ok -> END, else if needs_revision and round < max_rounds -> rewrite
builder.add_conditional_edges(
    "reflect",
    lambda x: x["verdict"] == "ok",
    {"ok": END},
)
builder.add_conditional_edges(
    "reflect",
    lambda x: x["verdict"] == "needs_revision" and x["round"] < x["max_rounds"],
    {"needs_revision": "rewrite"},
)
# If needs_revision but round >= max_rounds -> END
builder.add_edge("reflect", END, condition=lambda _: True)  # fallback

# Add rewrite to reflect loop
builder.add_edge("rewrite", "reflect")

graph = builder.compile(checkpointer=MemorySaver())

# ---------- Demo execution ----------------------------------------------
async def run_demo(question: str, max_rounds: int = 2):
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": max_rounds,
    }
    result = await graph.ainvoke(initial_state)
    # Extract final answer
    final_answer = result.get("draft", "")
    print("\n=== Final Answer ===")
    print(final_answer)
    return final_answer

if __name__ == "__main__":
    import asyncio

    examples = [
        "Explain the difference between a tool and a resource in MCP.",
        "What is the capital of France?",
        "Describe how to set up a virtual environment in Python 3.10.",
    ]
    for q in examples:
        print("\nQuestion:", q)
        asyncio.run(run_demo(q))
"""