"""
Main entry point for the "Planning Agent" assignment.

The project demonstrates a LangGraph agent that first plans a task into discrete steps and then executes those steps one by one.

Examples are provided in the ``__main__`` section – run the script with different tasks to see how the planner works.

Requirements:
- langgraph>=0.2.0
- langchain-openai>=0.3.0
- python-dotenv>=1.0.0
- rich>=13.0.0
"""

from __future__ import annotations

import os
from typing import TypedDict, List, Optional

from langgraph.graph import StateGraph, START, END
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from rich.console import Console
from rich.table import Table

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

# ---------------------------------------------------------------------------
# 2. LLM configuration – BroJS endpoint
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
    api_key=os.getenv("JOURNAL_MCP_PAT"),
    temperature=0.2,
)

# ---------------------------------------------------------------------------
# 3. Planning node – split task into steps
# ---------------------------------------------------------------------------
def planning(state: PlanningState) -> PlanningState:
    """Ask the LLM to produce a numbered list of steps.

    The prompt forces JSON output for reliable parsing.
    """
    system = SystemMessage(
        content="You are a helpful assistant that splits a task into clear, actionable steps. Return a JSON array of strings under the key `plan`."
    )
    user = HumanMessage(content=f"Plan the following task: {state['task']}")

    response = llm.invoke([system, user])
    text = response.content.strip()

    # Try to parse JSON; if fails, fall back to simple split by lines.
    try:
        import json
        data = json.loads(text)
        plan: List[str] = data.get("plan", [])
    except Exception:
        # Fallback – split on newlines and strip numbering
        plan = []
        for line in text.splitlines():
            line = line.strip()
            if not line:
                continue
            # Remove leading numbers like "1. " or "- "
            if line[0].isdigit() and (len(line) > 2 and line[1] in ".-"):
                line = line.split("", 1)[1]
            plan.append(line)
    return {
        **state,
        "plan": plan,
        "current_step": 0,
        "results": [],
    }

# ---------------------------------------------------------------------------
# 4. Execution node – perform one step (here we just echo the step)
# ---------------------------------------------------------------------------
def execution(state: PlanningState) -> PlanningState:
    """Execute a single step.

    In a real assignment you would replace this with calls to tools or other logic.
    For demonstration, we simply record the step text as the result.
    """
    idx = state["current_step"]
    if idx >= len(state["plan"]):
        return state
    step_text = state["plan"][idx]
    # Simulate execution – in practice you might call a tool here.
    result = f"Executed: {step_text}"
    new_results = state["results"].copy()
    new_results.append(result)
    return {
        **state,
        "current_step": idx + 1,
        "results": new_results,
    }

# ---------------------------------------------------------------------------
# 5. Decision node – should we continue?
# ---------------------------------------------------------------------------
def should_continue(state: PlanningState) -> str:
    if state["current_step"] >= len(state.get("plan", [])):
        return "finish"
    return "execute"

# ---------------------------------------------------------------------------
# 6. Build the graph
# ---------------------------------------------------------------------------
graph = StateGraph(PlanningState)
graph.add_node("planning", planning)
graph.add_node("execution", execution)
graph.add_conditional_edges(
    "planning",
    lambda _: "execute" if _.get("plan") else "finish",
)
graph.add_edge("execution", "should_continue")
graph.add_conditional_edges(
    "should_continue",
    should_continue,
    {"execute": "execution", "finish": END},
)
agent = graph.compile()

# ---------------------------------------------------------------------------
# 7. Helper to run the agent and pretty‑print results
# ---------------------------------------------------------------------------
def run_agent(task: str) -> None:
    console = Console()
    state: PlanningState = {"task": task, "plan": None, "current_step": 0, "results": []}
    config = {"configurable": {"thread_id": f"{task[:8]}"}}

    # Run the agent – we capture intermediate states via a callback.
    results: List[PlanningState] = []
    for event in agent.stream(state, config):
        if isinstance(event, dict) and "messages" in event:
            continue  # ignore final message
        results.append(event)

    # Print plan
    console.print("\n[bold underline]Task:[/]", task)
    if state.get("plan"):
        table = Table(title="Plan", show_header=False, box=None)
        for i, step in enumerate(state["plan"], 1):
            table.add_row(f"{i}. {step}")
        console.print(table)

    # Print execution results
    console.print("\n[bold underline]Execution Results:[/]")
    for r in state.get("results", []):
        console.print(r)

    console.print("\n[green]Finished.[/]\n")

# ---------------------------------------------------------------------------
# 8. Demo examples
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    examples = [
        "Compare Python and JavaScript in web development.",
        "Plan a weekend trip to the mountains.",
        "Explain how a neural network learns.",
    ]
    for ex in examples:
        run_agent(ex)