"""
Agent implementation for the Planning Agent assignment.

This module contains all logic for building a LangGraph agent that:
1. Plans a task into discrete steps using an LLM.
2. Executes each step sequentially, collecting results.
3. Returns a final summary of all results.

The graph is built in :func:`build_agent` and returned as a compiled object.
"""

import os
from typing import TypedDict, List
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage

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

# ---------------------------------------------------------------------------
# 2. LLM configuration – BroJS provider
# ---------------------------------------------------------------------------
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 the task into steps
# ---------------------------------------------------------------------------
def planning(state: PlanningState) -> PlanningState:
    """Ask LLM to produce a numbered list of steps.

    The prompt asks for JSON output with a single key ``steps`` containing an array
    of strings.  This guarantees deterministic parsing.
    """
    task = state["task"]
    system_prompt = (
        "You are a helpful assistant that breaks down a user request into a list of concrete, actionable steps."
    )
    user_prompt = (
        f"Task: {task}\n\nReturn a JSON object with a single key 'steps' containing an array of strings.  Each string should be a concise step.  Provide between 3 and 6 steps.")

    response = llm.invoke([SystemMessage(content=system_prompt), HumanMessage(content=user_prompt)])
    # Parse JSON safely
    import json, re
    try:
        data = json.loads(response.content)
        steps = data.get("steps", [])
    except Exception:
        # Fallback: extract numbered list via regex
        pattern = r"\d+\.\s*(.+)"
        steps = [m.group(1).strip() for m in re.finditer(pattern, response.content)]
    return {
        "task": task,
        "plan": steps,
        "current_step": 0,
        "results": [],
    }

# ---------------------------------------------------------------------------
# 4. Execution node – run one step and record result
# ---------------------------------------------------------------------------
def execution(state: PlanningState) -> PlanningState:
    idx = state["current_step"]
    plan = state["plan"] or []
    if idx >= len(plan):
        return state
    step_text = plan[idx]
    # For demonstration, we simply echo the step as result.
    result = f"Result of step {idx+1}: {step_text}"
    new_results = state["results"] + [result]
    return {
        "task": state["task"],
        "plan": plan,
        "current_step": idx + 1,
        "results": new_results,
    }

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

# ---------------------------------------------------------------------------
# 6. Build graph and compile
# ---------------------------------------------------------------------------

def build_agent() -> StateGraph:
    graph = StateGraph(PlanningState)
    graph.add_node("planning", planning)
    graph.add_node("execution", execution)
    graph.add_conditional_edges(
        "planning",
        lambda _: "execute" if _["plan"] else "finish",
    )
    graph.add_edge("execution", "should_continue")
    graph.add_conditional_edges(
        "should_continue",
        should_continue,
        {"execute": "execution", "finish": END},
    )
    graph.set_entry_point("planning")
    return graph

# ---------------------------------------------------------------------------
# 7. Public helper to get compiled agent
# ---------------------------------------------------------------------------
def get_agent():
    """Return a compiled LangGraph agent ready for invocation."""
    return build_agent().compile()

# ---------------------------------------------------------------------------
# If run as script, demonstrate usage with one example task.
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    from rich.console import Console
    console = Console()
    agent = get_agent()
    task_text = "Compare Python and JavaScript in terms of performance, syntax simplicity, and ecosystem support."
    result_state = agent.invoke({"task": task_text})
    console.print("\n[bold cyan]Plan:\n[/bold cyan]")
    for i, step in enumerate(result_state.get("plan", []), 1):
        console.print(f"{i}. {step}")
    console.print("\n[blue]Execution results:\n[/blue]")
    for r in result_state.get("results", []):
        console.print(r)
    console.print("\n[bold magenta]Final summary:[/bold magenta]\n")
    # Final LLM summarization
    summary_prompt = (
        "You have executed the following steps: \n" + "\n".join(result_state.get("results", [])) + "\nProvide a concise final answer.")
    summary_response = llm.invoke([HumanMessage(content=summary_prompt)])
    console.print(summary_response.content)