"""
LangGraph planning agent example.

Run with:
    python main.py "Compare Python and JavaScript"
"""
import json, sys
from typing import TypedDict, List
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_chat_agent
from langchain_openai import ChatOpenAI

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

# 2. LLM for planning and execution
llm = ChatOpenAI(temperature=0)

# 3. Planning node – split task into steps
async def planning(state: PlanningState) -> PlanningState:
    prompt = (
        "You are a helpful assistant that splits a user task into a numbered list of 3-6 concrete steps.
        Return only the JSON array of strings, e.g. ["Step 1", "Step 2"] without any explanation.")
    response = await llm.agenerate([{"role": "user", "content": f"{prompt}\nTask: {state['task']}"}])
    text = response.generations[0][0].text.strip()
    try:
        plan = json.loads(text)
        if not isinstance(plan, list): raise ValueError
    except Exception as e:
        # fallback simple split by lines
        plan = [line.strip() for line in text.splitlines() if line.strip()]
    return {**state, "plan": plan, "current_step": 0, "results": []}

# 4. Execution node – run one step
async def execution(state: PlanningState) -> PlanningState:
    step = state['plan'][state['current_step']]
    prompt = f"Execute the following step and return only the result string: {step}"
    response = await llm.agenerate([{"role": "user", "content": prompt}])
    result = response.generations[0][0].text.strip()
    new_results = state['results'] + [result]
    return {**state, "current_step": state['current_step'] + 1, "results": new_results}

# 5. Condition node – decide to continue or finish
def should_continue(state: PlanningState) -> str:
    if state['current_step'] >= len(state['plan']):
        return END
    return "execute"

# 6. Build graph
workflow = StateGraph(PlanningState)
workflow.add_node("planning", planning)
workflow.add_node("execution", execution)
workflow.set_entry_point("planning")
workflow.add_conditional_edges("planning", lambda _: "execute")
workflow.add_edge("execution", "should_continue")
workflow.add_conditional_edges("should_continue", should_continue)
graph = workflow.compile()

# 7. Run demo
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python main.py '<task>'")
        sys.exit(1)
    task = sys.argv[1]
    result = graph.invoke({"task": task})
    plan = result['plan']
    results = result['results']
    print(f"\nTask: {task}\n")
    print("Plan:\n")
    for i, step in enumerate(plan, 1):
        print(f"{i}. {step}")
    print("\nResults:\n")
    for i, res in enumerate(results, 1):
        print(f"[Step {i}] {res}\n")
    final = "\nFinal answer: " + "\n".join(results)
    print(final)