"""
LangGraph planning agent example.
"""

from typing import TypedDict, List
import os

# State definition
class PlanningState(TypedDict):
    task: str
    plan: List[str] | None
    current_step: int
    results: List[str]

# LLM setup (OpenAI)
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2, api_key=os.getenv("OPENAI_API_KEY"))

# Planning node
async def planning(state: PlanningState) -> PlanningState:
    prompt = (
        f"Разбей задачу '{state['task']}' на 3–6 конкретных шагов.\n"
        "Ответ в виде нумерованного списка, без лишних слов."
    )
    response = await llm.agenerate([prompt])
    text = response.generations[0][0].text.strip()
    # Parse numbered list
    steps: List[str] = []
    for line in text.splitlines():
        if line.lstrip().startswith("1") or line.lstrip()[0].isdigit():
            step = line.split('.', 1)[-1].strip()
            if step:
                steps.append(step)
    return {
        "task": state["task"],
        "plan": steps,
        "current_step": 0,
        "results": [],
    }

# Execution node
async def execution(state: PlanningState) -> PlanningState:
    idx = state["current_step"]
    step_text = state["plan"][idx]
    prompt = f"Выполни шаг {idx+1}: {step_text}.\nОтвет в виде короткого абзаца." 
    response = await llm.agenerate([prompt])
    result = response.generations[0][0].text.strip()
    new_results = state["results"] + [result]
    return {
        "task": state["task"],
        "plan": state["plan"],
        "current_step": idx + 1,
        "results": new_results,
    }

# Condition node
from langgraph import StateGraph, END

def should_continue(state: PlanningState):
    if state["current_step"] >= len(state["plan"]):
        return "finish"
    return "execute"

# Graph definition
workflow = StateGraph(PlanningState)
workflow.add_node("planning", planning)
workflow.add_node("execution", execution)
workflow.set_entry_point("planning")
workflow.add_conditional_edges(
    "execution",
    should_continue,
    {
        "execute": "execution",
        "finish": END,
    },
)

# Run example
if __name__ == "__main__":
    task = "Сравни Python и JavaScript"
    initial_state: PlanningState = {"task": task, "plan": None, "current_step": 0, "results": []}
    result = workflow.invoke(initial_state)
    print("План:")
    for i, step in enumerate(result["plan"]):
        print(f"{i+1}. {step}")
    print("\nШаги: ")
    for i, res in enumerate(result["results"]):
        print(f"[Шаг {i+1}] {res}\n")
    print("Итог:")
    final = llm.invoke("\n\nСводка всех результатов: " + "\n".join(result["results"]))
    print(final)