"""
LangGraph Agent that plans and executes a task step by step.

Usage:
    python agent.py "Compare Python and JavaScript"

Dependencies:
    - langgraph
    - langchain-openai
    - langchain-ollama (optional)
"""

from __future__ import annotations

import json
import os
import re
import sys
import traceback
from typing import TypedDict

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI

# ----- State definition (from snippet) -----
class PlanningState(TypedDict):
    task: str
    plan: list[str] | None
    current_step: int
    results: list[str]

# ----- LLM configuration -----
# Prefer OpenAI if API key is set; otherwise fallback to Ollama
if os.getenv("OPENAI_API_KEY"):
    llm = ChatOpenAI(
        model="gpt-4o-mini",
        temperature=0,
    )
else:
    # Fall back to local Ollama model if available
    llm = ChatOpenAI(
        model=os.getenv("CHAT_MODEL", "llama3"),
        base_url=os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1"),
        api_key="ollama",
        temperature=0,
    )

# ----- Planning node -----

def planning(state: PlanningState) -> PlanningState:
    """Prompt LLM to break the task into 3–6 numbered steps."""
    try:
        prompt = (
            f"Given the task '{state['task']}', break it into 3-6 numbered steps. "
            "Return only the numbered list or a JSON array of steps."
        )
        raw = llm.invoke(prompt)
        text = raw if isinstance(raw, str) else raw.content
    except Exception as e:
        raise RuntimeError(f"LLM failed in planning node: {e}")

    # Parse plan – first try JSON, then regex
    plan: list[str] | None = None
    try:
        plan = json.loads(text)
        if not isinstance(plan, list):
            plan = None
    except Exception:
        plan = None

    if plan is None:
        pattern = r"^\s*\d+\.\s+(.*)$"
        plan = [m.group(1).strip() for m in re.finditer(pattern, text, re.MULTILINE)]

    if not plan or not (3 <= len(plan) <= 6):
        raise ValueError(
            f"Planning node returned invalid plan: {plan}. Expected 3-6 steps."
        )

    state["plan"] = plan
    state["current_step"] = 0
    state["results"] = []
    return state

# ----- Execution node -----

def execution(state: PlanningState) -> PlanningState:
    """Execute a single step from the plan and record the result."""
    step_idx = state["current_step"]
    if state["plan"] is None or step_idx >= len(state["plan"]):
        return state
    step = state["plan"][step_idx]
    try:
        prompt = f"Execute step: {step}\nProvide a concise result."
        raw = llm.invoke(prompt)
        result = raw if isinstance(raw, str) else raw.content
    except Exception as e:
        result = f"Error executing step: {e}"
    state["results"].append(result.strip())
    state["current_step"] = step_idx + 1
    return state

# ----- Graph construction -----
# Graph diagram (from snippet):
# START → planning → execution → should_continue
#                      ↑____________| (execute)
#                      finish → END

builder = StateGraph(PlanningState)
builder.add_node("planning", planning)
builder.add_node("execution", execution)

# Conditional edges to loop execution until all steps processed
builder.add_conditional_edges(
    "planning",
    lambda state: "execution" if state.get("plan") else "END"
)
builder.add_conditional_edges(
    "execution",
    lambda state: "execution"
    if state.get("current_step", 0) < len(state.get("plan", []))
    else "END"
)

builder.set_entry_point("planning")

# Compile graph with in-memory checkpointing
graph = builder.compile(checkpointer=InMemorySaver())

# ----- Demo -----

def run_demo(task: str) -> None:
    initial_state: PlanningState = {
        "task": task,
        "plan": None,
        "current_step": 0,
        "results": [],
    }
    try:
        result = graph.invoke(initial_state)
    except Exception:
        traceback.print_exc()
        sys.exit(1)

    plan = result.get("plan", [])
    print("\n===== PLAN =====")
    for idx, step in enumerate(plan, 1):
        print(f"{idx}. {step}")
    print("\n===== RESULTS =====")
    for idx, res in enumerate(result.get("results", []), 1):
        print(f"[Step {idx}] {res}\n")
    print("===== SUMMARY =====")
    try:
        summary_prompt = f"Given the collected results: {result.get('results', [])}, produce a concise summary of the task outcome."
        summary = llm.invoke(summary_prompt)
        print(summary if isinstance(summary, str) else summary.content)
    except Exception as e:
        print(f"Error generating summary: {e}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python agent.py '<task>'")
        sys.exit(1)
    task = " ".join(sys.argv[1:])
    run_demo(task)