import json
from typing import Dict, Any

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


# ------------------------------------------------------------
# 1. Определяем состояние графа
# ------------------------------------------------------------
class PlannerState:
    """Состояние агента: цель и список шагов."""
    goal: str = ""
    steps: list[str] | None = None


# ------------------------------------------------------------
# 2. Функция планирования (использует LLM)
# ------------------------------------------------------------
def plan_steps(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Принимает цель и генерирует список шагов в формате JSON.
    Возвращает обновлённое состояние с полем `steps`.
    """
    goal = state["goal"]
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)

    prompt = f"""
    Ты планировщик задач. Тебе дана цель:
    {goal}

    Сгенерируй список конкретных шагов, которые нужно выполнить,
    чтобы достичь этой цели. Каждый шаг должен быть короткой фразой.
    Выведите результат в виде JSON-массива строк, например:
    ["Шаг 1", "Шаг 2", ...]
    """
    response = llm.invoke(prompt)
    try:
        steps = json.loads(response.content.strip())
        if not isinstance(steps, list):
            raise ValueError
    except Exception:
        # Если LLM не вернул корректный JSON – создаём простую fallback‑планировку
        steps = [f"Шаг {i+1} для выполнения цели: {goal}" for i in range(3)]

    return {"steps": steps}


# ------------------------------------------------------------
# 3. Создаём граф планировщика
# ------------------------------------------------------------
builder = StateGraph(PlannerState)

# Входной узел – задаёт цель и сразу переходит к планированию
def start_node(state: PlannerState) -> Dict[str, Any]:
    return {"goal": state.goal}

builder.add_node("start", start_node)
builder.add_node("plan_steps", plan_steps)

# Переходы
builder.set_entry_point("start")
builder.add_edge("start", "plan_steps")
builder.add_edge("plan_steps", END)

planner_graph = builder.compile()


# ------------------------------------------------------------
# 4. Функция запуска планировщика
# ------------------------------------------------------------
def run_planner(goal: str) -> list[str]:
    """
    Запускает граф планировщика для заданной цели.
    Возвращает список шагов.
    """
    result = planner_graph.invoke({"goal": goal})
    return result.get("steps", [])


# ------------------------------------------------------------
# 5. Тестовый запуск (можно удалить при интеграции)
# ------------------------------------------------------------
if __name__ == "__main__":
    sample_goal = "Составить план подготовки к экзамену по Python"
    steps = run_planner(sample_goal)
    print("\nПлан действий:")
    for i, step in enumerate(steps, 1):
        print(f"{i}. {step}")