import os
from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
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,
)

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

# ---------- Structured output models ----------
class CritiqueModel(BaseModel):
    verdict: str = Field(..., description="ok or needs_revision")
    remarks: list[str] = Field(..., description="2-3 bullet points of critique")

class RewriteModel(BaseModel):
    draft: str = Field(..., description="Rewritten draft answer")

critique_parser = PydanticOutputParser(pydantic_object=CritiqueModel)
rewrite_parser = PydanticOutputParser(pydantic_object=RewriteModel)

# ---------- Nodes ----------
async def draft_answer(state: ReflectState) -> ReflectState:
    prompt = (
        f"Write a concise answer (5–10 sentences) to the following question:\n\n"
        f"Question: {state['question']}"
    )
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    state['draft'] = response.content.strip()
    return state

async def reflect(state: ReflectState) -> ReflectState:
    prompt = (
        f"You are a critic evaluating the draft answer for completeness, concreteness, and absence of filler.\n"
        f"Draft: {state['draft']}\n"
        f"Provide a verdict ("ok" or "needs_revision") and 2–3 bullet points of critique.\n"
        f"Return the result in JSON format: {{\"verdict\": "...", \"remarks\": ["...", ...]}}"
    )
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    parsed = critique_parser.parse(response.content)
    state['critique'] = "\n".join(parsed.remarks)
    state['verdict'] = parsed.verdict
    return state

async def rewrite(state: ReflectState) -> ReflectState:
    prompt = (
        f"Rewrite the draft answer to address the following critique:\n"
        f"Critique: {state['critique']}\n"
        f"Original draft: {state['draft']}\n"
        f"Provide the rewritten draft only.\n"
        f"Return JSON: {{\"draft\": "..."}}"
    )
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    parsed = rewrite_parser.parse(response.content)
    state['draft'] = parsed.draft
    state['round'] += 1
    return state

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

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

# Conditional edges after reflect
builder.add_conditional_edges(
    "reflect",
    lambda x: x['verdict'] == "ok",
    {"ok": END, "needs_revision": "rewrite"},
)

# After rewrite go back to reflect
builder.add_edge("rewrite", "reflect")

# If max rounds exceeded, end
builder.add_conditional_edges(
    "reflect",
    lambda x: x['round'] >= x['max_rounds'] and x['verdict'] == "needs_revision",
    {True: END, False: END},  # both lead to END, but loop already handled
)

graph = builder.compile()

# ---------- Demo ----------
async def main():
    initial_state: ReflectState = {
        "question": "Объясни студенту разницу между tool и resource в MCP",
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }
    final_state = await graph.ainvoke(initial_state)
    print("\n--- Final Draft ---")
    print(final_state['draft'])
    print("\n--- Critique ---")
    print(final_state['critique'])
    print("\n--- Verdict ---")
    print(final_state['verdict'])
    print("\n--- Rounds Used ---")
    print(final_state['round'])

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