#!/usr/bin/env python
# main.py
# LangGraph planning agent example

import os
from typing import TypedDict, List, Dict, Any

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_chat_agent
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage

# ---------- LLM ----------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

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

# ---------- Planning node ----------
planning_prompt = (
    "You are a helpful assistant. Given the following task, break it into 3-6 concrete steps. "
    "Return the steps as a numbered list or JSON array. The steps should be short, actionable, and in Russian."
)

async def planning(state: PlanningState) -> PlanningState:
    task = state["task"]
    messages = [HumanMessage(content=f"{planning_prompt}\nTask: {task}")]
    response = await llm.ainvoke(messages)
    text = response.content.strip()

    # Try to parse JSON first
    plan: List[str] | None = None
    try:
        import json
        data = json.loads(text)
        if isinstance(data, list):
            plan = [str(item).strip() for item in data]
    except Exception:
        pass

    # If JSON parsing failed, try to extract numbered list
    if plan is None:
        import re
        lines = re.findall(r"\d+\.\s*(.+)", text)
        if lines:
            plan = [line.strip() for line in lines]

    if plan is None:
        raise ValueError("Could not parse plan from LLM response")

    return {
        "task": task,
        "plan": plan,
        "current_step": 0,
        "results": [],
    }

# ---------- Execution node ----------
async def execution(state: PlanningState) -> PlanningState:
    plan = state["plan"]
    idx = state["current_step"]
    if plan is None or idx >= len(plan):
        return state

    step = plan[idx]
    # Execute the step: ask LLM to produce result for this step
    messages = [HumanMessage(content=f"Task: {state['task']}\nStep {idx+1}: {step}\nProvide the result for this step.")]
    response = await llm.ainvoke(messages)
    result = response.content.strip()

    new_results = state["results"].copy()
    new_results.append(result)

    return {
        "task": state["task"],
        "plan": plan,
        "current_step": idx + 1,
        "results": new_results,
    }

# ---------- Should continue ----------
async def should_continue(state: PlanningState) -> str:
    if state["current_step"] >= len(state["plan"]):
        return "finish"
    return "execute"

# ---------- Build graph ----------
builder = StateGraph(PlanningState)
builder.add_node("planning", planning)
builder.add_node("execution", execution)
builder.add_conditional_edges(
    "planning",
    lambda _: "execute",
)
builder.add_conditional_edges(
    "execution",
    should_continue,
    {
        "execute": "execution",
        "finish": END,
    },
)
builder.set_entry_point("planning")
graph = builder.compile()

# ---------- Run example ----------
if __name__ == "__main__":
    task = "Сравни Python и JavaScript"
    initial_state: PlanningState = {
        "task": task,
        "plan": None,
        "current_step": 0,
        "results": [],
    }
    result = graph.invoke(initial_state)

    print(f"\nЗадача: {task}\n")
    print("План:")
    for i, step in enumerate(result["plan"], 1):
        print(f"{i}. {step}")
    print("\n[Шаги]")
    for i, res in enumerate(result["results"], 1):
        print(f"[Шаг {i}] {res}\n")
    print("Итог: " + "\n".join(result["results"]))