"""LangGraph-агент с рефлексией и доработкой (без try/except)."""
import os
import asyncio
from typing import TypedDict

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END

load_dotenv()

# ── LLM ──────────────────────────────────────────────────────────────────
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

# ── Nodes ─────────────────────────────────────────────────────────────────

async def draft_answer(state: ReflectState) -> dict:
    """Первичный ответ на вопрос (5–10 предложений)."""
    prompt = (
        "Напиши краткий, конкретный ответ (5–10 предложений) на вопрос.\n\n"
        f"Вопрос: {state['question']}"
    )
    response = await llm.ainvoke([{"role": "user", "content": prompt}])
    return {"draft": response.content.strip(), "round": 0}


async def reflect(state: ReflectState) -> dict:
    """LLM-критик: оценивает черновик и выносит структурированный вердикт.

    Промпт требует строгого формата ответа — вердикт парсится напрямую,
    без try/except и без запасных значений при ошибке парсинга.
    """
    prompt = (
        "Ты — строгий критик учебных ответов. Оцени черновик по трём критериям:\n"
        "1. Полнота — все ключевые аспекты раскрыты\n"
        "2. Конкретика — есть примеры или точные определения\n"
        "3. Отсутствие воды — нет пустых фраз\n\n"
        f"Черновик:\n{state['draft']}\n\n"
        "Ответь строго по шаблону (без лишнего текста до первой строки):\n"
        "VERDICT: ok\n"
        "CRITIQUE:\n- ...\n\n"
        "или:\n\n"
        "VERDICT: needs_revision\n"
        "CRITIQUE:\n- замечание 1\n- замечание 2\n- замечание 3"
    )
    response = await llm.ainvoke([{"role": "user", "content": prompt}])
    text = response.content.strip()

    # Парсим вердикт из строки «VERDICT: ok» / «VERDICT: needs_revision»
    verdict = "needs_revision"
    for line in text.splitlines():
        stripped = line.strip()
        if stripped.upper().startswith("VERDICT:"):
            value = stripped.split(":", 1)[1].strip().lower()
            verdict = "ok" if value == "ok" else "needs_revision"
            break

    return {"critique": text, "verdict": verdict}


async def rewrite(state: ReflectState) -> dict:
    """Перерабатывает черновик с учётом замечаний критика."""
    prompt = (
        "Улучши ответ, учитывая замечания критика. Сохрани объём 5–10 предложений.\n\n"
        f"Замечания:\n{state['critique']}\n\n"
        f"Текущий черновик:\n{state['draft']}"
    )
    response = await llm.ainvoke([{"role": "user", "content": prompt}])
    return {"draft": response.content.strip(), "round": state["round"] + 1}

# ── Router ────────────────────────────────────────────────────────────────

def route_after_reflect(state: ReflectState) -> str:
    """Возвращает следующий узел после reflect."""
    if state["verdict"] == "ok":
        return "end"
    if state["round"] >= state["max_rounds"]:
        return "end"
    return "rewrite"

# ── 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")
builder.add_conditional_edges(
    "reflect",
    route_after_reflect,
    {"end": END, "rewrite": "rewrite"},
)
builder.add_edge("rewrite", "reflect")   # после rewrite всегда → reflect
                                          # лимит раундов проверяет route_after_reflect

graph = builder.compile()

# ── Demo ──────────────────────────────────────────────────────────────────

async def main() -> None:
    question = "Объясни студенту разницу между tool и resource в MCP."
    print(f"Вопрос: {question}")
    print("=" * 60)

    init: ReflectState = {
        "question":   question,
        "draft":      "",
        "critique":   "",
        "verdict":    "",
        "round":      0,
        "max_rounds": 2,
    }

    async for chunk in graph.astream(init, stream_mode="updates"):
        node = next(iter(chunk))
        data = chunk[node]
        if node == "draft_answer":
            print(f"\n[Черновик]\n{data.get('draft','')}")
        elif node == "reflect":
            print(f"\n[Критик] Вердикт: {data.get('verdict','')}")
            print(data.get("critique", ""))
        elif node == "rewrite":
            print(f"\n[Переработка, раунд {data.get('round','')}]\n{data.get('draft','')}")

    final = await graph.ainvoke(init)
    print(f"\n{'='*60}\n[Финальный ответ]\n{final['draft']}")


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