import os
from typing import TypedDict, Literal

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI


# ---------- 1. Состояние ----------
class ReflectState(TypedDict):
    question: str
    draft: str
    critique: str
    verdict: Literal["ok", "needs_revision"]
    round: int
    max_rounds: int


# ---------- 2. Узлы ----------
def draft_answer(state: ReflectState) -> dict:
    """Первичный ответ на вопрос."""
    llm = ChatOpenAI(model="gpt-4o-mini")
    prompt = f"Напиши краткий ответ (5–10 предложений) на следующий вопрос:\n\n{state['question']}"
    response = llm.invoke(prompt)
    state["draft"] = response.content.strip()
    return {"draft": state["draft"]}


def reflect(state: ReflectState) -> dict:
    """LLM‑критик оценивает ответ."""
    llm = ChatOpenAI(model="gpt-4o-mini")
    prompt = (
        f"Оцени следующий ответ на вопрос \"{state['question']}\":\n\n"
        f"{state['draft']}\n\n"
        "Дай verdict: ok или needs_revision. "
        "Если needs_revision, перечисли 2–3 конкретных замечания."
    )
    response = llm.invoke(prompt)
    text = response.content.strip()
    # Разделяем verdict и critique
    if "needs_revision" in text:
        state["verdict"] = "needs_revision"
        parts = text.split("needs_revision", 1)[-1].strip()
        state["critique"] = parts
    else:
        state["verdict"] = "ok"
        state["critique"] = ""
    return {"verdict": state["verdict"], "critique": state["critique"]}


def rewrite(state: ReflectState) -> dict:
    """Учитываем critique и переписываем ответ."""
    llm = ChatOpenAI(model="gpt-4o-mini")
    prompt = (
        f"Перепиши ответ, учитывая следующие замечания:\n\n"
        f"{state['critique']}\n\n"
        f"Оригинальный ответ:\n{state['draft']}"
    )
    response = llm.invoke(prompt)
    state["draft"] = response.content.strip()
    state["round"] += 1
    return {"draft": state["draft"], "round": state["round"]}


# ---------- 3. Граф ----------
builder = StateGraph(ReflectState)

builder.add_node("draft_answer", draft_answer)
builder.add_node("reflect", reflect)
builder.add_node("rewrite", rewrite)

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")

graph = builder.compile()


# ---------- 4. CLI ----------
def main():
    import argparse

    parser = argparse.ArgumentParser(description="LangGraph рефлексия")
    parser.add_argument(
        "--question",
        type=str,
        required=True,
        help="Вопрос, на который нужно ответить.",
    )
    args = parser.parse_args()

    # Инициализируем состояние
    state: ReflectState = {
        "question": args.question,
        "draft": "",
        "critique": "",
        "verdict": "",
        "round": 0,
        "max_rounds": 2,
    }

    # Запускаем граф
    result = graph.invoke(state)

    # Выводим финальный ответ и логи
    print("\n=== Финальный ответ ===")
    print(result["draft"])
    print("\n=== Лог рефлексии ===")
    print(f"Round: {result['round']}")
    print(f"Verdict: {result['verdict']}")
    if result["critique"]:
        print("Critique:")
        print(result["critique"])


if __name__ == "__main__":
    # Проверяем наличие ключа OpenAI
    if "OPENAI_API_KEY" not in os.environ:
        raise RuntimeError(
            "Пожалуйста, установите переменную окружения OPENAI_API_KEY."
        )
    main()