"""LangGraph с узлом рефлексии: draft_answer -> reflect -> rewrite (при необходимости)."""
import os
from typing import Literal
from typing_extensions import TypedDict
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import StateGraph, START, END

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.5,
)

class ReflectState(TypedDict):
    """Состояние графа рефлексии."""
    question: str
    draft: str
    critique: str
    verdict: str    # ok | needs_revision
    round: int
    max_rounds: int


def draft_answer(state: ReflectState) -> dict:
    """Узел 1: генерирует первичный ответ на вопрос (5-10 предложений)."""
    response = llm.invoke([
        SystemMessage(content=(
            "Ты эксперт. Напиши развёрнутый ответ на вопрос в 5-10 предложений. "
            "Будь конкретен, приводи примеры, избегай воды."
        )),
        HumanMessage(content=state["question"]),
    ])
    print(f"\n[draft_answer] Черновик:\n{response.content}\n")
    return {"draft": response.content, "round": 0}


def reflect(state: ReflectState) -> dict:
    """Узел 2: LLM-критик оценивает качество ответа и выставляет вердикт."""
    prompt = (
        f'''Оцени следующий ответ на вопрос: "{state["question"]}"

Ответ:
{state["draft"]}

Оцени по трём критериям:
1. Полнота — охвачены ли все аспекты вопроса?
2. Конкретика — есть ли конкретные примеры?
3. Отсутствие воды — нет ли лишних слов?

Выведи строго в формате:
VERDICT: ok
или
VERDICT: needs_revision
CRITIQUE: <2-3 конкретных замечания>'''
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    text = response.content

    verdict = "needs_revision"
    if "VERDICT: ok" in text or "verdict: ok" in text.lower():
        verdict = "ok"

    critique = ""
    if "CRITIQUE:" in text:
        critique = text.split("CRITIQUE:")[-1].strip()

    print(f"[reflect] Вердикт: {verdict}")
    if critique:
        print(f"[reflect] Замечания: {critique[:200]}")

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


def rewrite(state: ReflectState) -> dict:
    """Узел 3: улучшает черновик с учётом замечаний критика."""
    prompt = (
        f'''Улучши следующий ответ на вопрос: "{state["question"]}"

Текущий ответ:
{state["draft"]}

Замечания критика:
{state["critique"]}

Напиши улучшенный ответ в 5-10 предложений, устранив все замечания.'''
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    new_round = state["round"] + 1
    print(f"\n[rewrite] Раунд {new_round}:\n{response.content}\n")
    return {"draft": response.content, "round": new_round}


def route_after_reflect(state: ReflectState) -> Literal["rewrite", "__end__"]:
    """Роутер после узла reflect: продолжить или завершить."""
    if state["verdict"] == "ok":
        print("[route] Вердикт ok — завершаем.")
        return "__end__"
    if state["round"] >= state["max_rounds"]:
        print(f"[route] Достигнут лимит раундов ({state['max_rounds']}) — завершаем.")
        return "__end__"
    print(f"[route] needs_revision, раунд {state['round']} < {state['max_rounds']} — rewrite.")
    return "rewrite"

# ── Сборка графа ──────────────────────────────────────────────────────────────
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,
    {"rewrite": "rewrite", "__end__": END},
)
builder.add_edge("rewrite", "reflect")

graph = builder.compile()