"""
Демонстрация LangGraph-агента с рефлексией.

Граф: START -> draft_answer -> reflect -> (ok: END | needs_revision: rewrite -> reflect)
max_rounds по умолчанию 2.

Запуск:
    python main.py
"""
import asyncio
import os
from dotenv import load_dotenv

load_dotenv()

from graph import graph, ReflectState  # noqa: E402

async def run(question: str, max_rounds: int = 2) -> ReflectState:
    """Запускает граф и возвращает финальное состояние."""
    print(f"Вопрос: {question}")
    print(f"max_rounds: {max_rounds}")
    print("=" * 60)

    state = await graph.ainvoke({
        "question":   question,
        "draft":      "",
        "critique":   "",
        "verdict":    "",
        "round":      0,
        "max_rounds": max_rounds,
    })
    return state

def main() -> None:
    # Демо 1: основной вопрос из задания
    state = asyncio.run(run(
        question="Объясни студенту разницу между tool и resource в MCP",
        max_rounds=2,
    ))
    print("\n" + "=" * 60)
    print(f"Итоговый ответ (раундов доработки: {state['round']}):")
    print(state["draft"])
    print(f"Финальный вердикт критика: {state['verdict']}")

    print("\n" + "=" * 60)

    # Демо 2: другой вопрос
    state2 = asyncio.run(run(
        question="Что такое LangGraph и чем он отличается от обычной цепочки LangChain?",
        max_rounds=2,
    ))
    print(f"Итоговый ответ (раундов доработки: {state2['round']}):")
    print(state2["draft"])

    print("\n" + "=" * 60)

    # Демо 3: короткий вопрос — проверяем что лимит соблюдается
    state3 = asyncio.run(run(
        question="Что такое StateGraph в LangGraph?",
        max_rounds=1,
    ))
    print(f"Итоговый ответ (max_rounds=1, раундов: {state3['round']}):")
    print(state3["draft"])


if __name__ == "__main__":
    main()