from typing import TypedDict, List, Optional
import json
import asyncio
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_conditional_node

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

async def planning(state: PlanningState, llm: ChatOpenAI) -> PlanningState:
    """LLM generates a JSON plan with a list of steps."""
    prompt = (
        "Разбей задачу на 3–6 конкретных шагов. Возвращай JSON с ключом \"plan\".\n\n"
        f"Задача: {state['task']}"
    )
    response = await llm.ainvoke(prompt)
    # response may be a string or a ChatResult
    text = response if isinstance(response, str) else response.content
    plan: List[str] = []
    try:
        data = json.loads(text)
        plan = data.get("plan", [])
    except Exception:
        # Fallback: parse numbered list
        for line in text.splitlines():
            line = line.strip()
            if not line:
                continue
            # Remove leading number or bullet
            if line[0].isdigit():
                parts = line.split('.', 1)
                if len(parts) == 2:
                    step = parts[1].strip()
                else:
                    step = line
                plan.append(step)
            elif line[0] in ('-','*'):
                plan.append(line[1:].strip())
    return {
        **state,
        "plan": plan,
        "current_step": 0,
        "results": []
    }

async def execution(state: PlanningState) -> PlanningState:
    """Execute a single step and record the result."""
    if state["plan"] is None or state["current_step"] >= len(state["plan"]):
        return state
    step = state["plan"][state["current_step"]]
    result = f"Шаг {state['current_step']+1}: {step}"
    new_results = state["results"] + [result]
    return {
        **state,
        "results": new_results,
        "current_step": state["current_step"] + 1
    }

def should_continue(state: PlanningState) -> str:
    """Decide whether to finish or execute another step."""
    if state["current_step"] >= len(state["plan"] or []):
        return "finish"
    return "execute"

async def planning_node(state: PlanningState, llm: ChatOpenAI) -> PlanningState:
    return await planning(state, llm)

def create_agent(llm: ChatOpenAI):
    workflow = StateGraph(PlanningState)
    # Add nodes
    workflow.add_node("planning", lambda state: planning_node(state, llm))
    workflow.add_node("execution", execution)
    workflow.add_conditional_node("should_continue", should_continue, {
        "execute": "execution",
        "finish": END
    })
    # Set entry point and edges
    workflow.set_entry_point("planning")
    workflow.add_edge("planning", "execution")
    workflow.add_edge("execution", "should_continue")
    return workflow.compile()

async def run_agent(task: str, llm: ChatOpenAI) -> PlanningState:
    agent = create_agent(llm)
    initial_state: PlanningState = {
        "task": task,
        "plan": None,
        "current_step": 0,
        "results": []
    }
    result = await agent.ainvoke(initial_state)
    return result