"""
LangGraph planning agent example.

This script demonstrates a simple LangGraph agent that:
1. Takes an input task description.
2. Uses an LLM to split the task into 3‑6 concrete steps (JSON list).
3. Executes each step sequentially, collecting results.
4. Returns a final summary of all results.

Requirements:
- python >= 3.10
- langgraph
- langchain-openai (or langchain-ollama)

Run with:
    python main.py "Compare Python and JavaScript"
"""

from __future__ import annotations

import json
import os
from typing import TypedDict, List

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

# ---------- 1. Define the state ---------------------------------
class PlanningState(TypedDict):
    task: str
    plan: List[str] | None
    current_step: int
    results: List[str]

# ---------- 2. LLM for planning ---------------------------------
# The user should set OPENAI_API_KEY in environment.
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

planning_prompt = (
    '''You are a helpful assistant that plans tasks.
Given the following task description, break it into 3‑6 concrete steps.
Return only a JSON array of strings. Example:
["Step 1: ...", "Step 2: ..."]'''
)

# ---------- 3. Planning node -----------------------------------
async def planning(state: PlanningState) -> PlanningState:
    task = state["task"]
    # Call LLM to get plan
    response = await llm.agenerate([planning_prompt + f"\nTask: {task}"])
    text = response.generations[0][0].text.strip()
    try:
        plan = json.loads(text)
        if not isinstance(plan, list):
            raise ValueError
    except Exception:
        # Fallback: split by newlines
        plan = [line for line in text.split("\n") if line]
    return {
        "task": task,
        "plan": plan,
        "current_step": 0,
        "results": [],
    }

# ---------- 4. Execution node -----------------------------------
async def execution(state: PlanningState) -> PlanningState:
    idx = state["current_step"]
    step_text = state["plan"][idx]
    # For demo, just echo the step as result.
    result = f"Result of {step_text}"
    new_results = state["results"].copy()
    new_results.append(result)
    return {
        "task": state["task"],
        "plan": state["plan"],
        "current_step": idx + 1,
        "results": new_results,
    }

# ---------- 5. Condition node -----------------------------------
def should_continue(state: PlanningState) -> str:
    if state["current_step"] >= len(state["plan"]):
        return "finish"
    return "execute"

# ---------- 6. 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,
    },
)
graph = builder.compile()

# ---------- 7. Demo runner -------------------------------------
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python main.py '<task description>'")
        sys.exit(1)
    task_desc = sys.argv[1]
    # Run graph
    result = graph.invoke({"task": task_desc})
    plan = result["plan"]
    results = result["results"]
    print("\nTask:", task_desc)
    print("\nPlan:\n", "\n".join(f"{i+1}. {step}" for i, step in enumerate(plan)))
    print("\nResults:\n", "\n".join(results))
    # Final summary via LLM
    final_prompt = (
        f"Given the following results: {json.dumps(results)}\nProvide a concise summary."
    )
    final_resp = llm.invoke(final_prompt)
    print("\nFinal Summary:\n", final_resp.content.strip())