import os
from typing import TypedDict, Dict, Any
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI

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

# 2. LLM instance
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)

# 3. Nodes

def draft_answer(state: ReflectState) -> Dict[str, Any]:
    prompt = (
        "Write a concise answer (5–10 sentences) to the following question:\n"
        f"Question: {state['question']}\n"
        "Answer:"
    )
    response = llm.invoke(prompt)
    state["draft"] = response.content.strip()
    return {"draft": state["draft"]}


def reflect(state: ReflectState) -> Dict[str, Any]:
    prompt = (
        "You are a critical reviewer of the draft answer.\n"
        "Evaluate the draft for completeness, specificity, and lack of filler.\n"
        "Provide a verdict: 'ok' if the answer is satisfactory, otherwise 'needs_revision'.\n"
        "If revision is needed, give 2–3 concise points for improvement.\n"
        f"Draft: {state['draft']}\n"
        "Verdict and critique:"
    )
    response = llm.invoke(prompt)
    text = response.content.strip()
    lines = text.splitlines()
    verdict_line = lines[0].lower().strip()
    verdict = "ok" if "ok" in verdict_line else "needs_revision"
    critique = "\n".join(lines[1:]).strip()
    state["verdict"] = verdict
    state["critique"] = critique
    return {"verdict": verdict, "critique": critique}


def rewrite(state: ReflectState) -> Dict[str, Any]:
    prompt = (
        "Rewrite the draft answer incorporating the following critique points.\n"
        "Keep the answer concise (5–10 sentences).\n"
        f"Critique: {state['critique']}\n"
        f"Original Draft: {state['draft']}\n"
        "Revised Answer:"
    )
    response = llm.invoke(prompt)
    state["draft"] = response.content.strip()
    state["round"] += 1
    return {"draft": state["draft"], "round": state["round"]}

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

# Edges
builder.set_entry_point("draft_answer")
builder.add_edge("draft_answer", "reflect")
builder.add_conditional_edges(
    "reflect",
    lambda x: x["verdict"],
    {
        "ok": END,
        "needs_revision": "rewrite"
    }
)
builder.add_edge("rewrite", "reflect")

# Max rounds guard
@builder.before_node("rewrite")
def check_rounds(state: ReflectState) -> ReflectState:
    if state["round"] >= state["max_rounds"]:
        state["verdict"] = "ok"
    return state

graph = builder.compile()

# 5. Demo execution
if __name__ == "__main__":
    question = "Объясни студенту разницу между tool и resource в MCP"
    initial_state: ReflectState = {
        "question": question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2
    }
    result = graph.invoke(initial_state)
    print("\n--- Final Draft ---\n")
    print(result["draft"])
    print("\n--- Critique ---\n")
    print(result["critique"])
    print("\n--- Verdict ---\n")
    print(result["verdict"])