"""
Shared state schema for the Sisyphus orchestrator graph.

In LangGraph, state is the single source of truth that flows through every node.
This is analogous to Loki's per-agent session context, but unified into one typed
dictionary that the entire graph shares.

Loki Concept Mapping:
    - Loki session context          → SisyphusState (TypedDict)
    - Loki todo__init / todo__add   → SisyphusState.todos list
    - Loki agent__spawn outputs     → SisyphusState.agent_outputs dict
    - Loki intent classification    → SisyphusState.intent field
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Annotated, Literal

from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict

# ---------------------------------------------------------------------------
# Intent types — mirrors Loki's Sisyphus classification table
# ---------------------------------------------------------------------------
IntentType = Literal[
    "trivial",        # Single file, known location, typo fix → handle yourself
    "exploration",    # "Find X", "Where is Y" → spawn explore
    "implementation", # "Add feature", "Fix bug" → spawn coder
    "architecture",   # Design questions, oracle triggers → spawn oracle
    "ambiguous",      # Unclear scope → ask user
]

# ---------------------------------------------------------------------------
# Todo item — mirrors Loki's built-in todo system
# ---------------------------------------------------------------------------
@dataclass
class TodoItem:
    """A single task in the orchestrator's todo list."""
    id: int
    task: str
    done: bool = False


def _merge_todos(existing: list[TodoItem], new: list[TodoItem]) -> list[TodoItem]:
    """
    Reducer for the todos field.

    LangGraph requires a reducer for any state field that can be written by
    multiple nodes.  This merges by id: if a todo with the same id already
    exists, the incoming version wins (allows marking done).
    """
    by_id = {t.id: t for t in existing}
    for t in new:
        by_id[t.id] = t
    return list(by_id.values())


# ---------------------------------------------------------------------------
# Core graph state
# ---------------------------------------------------------------------------
class SisyphusState(TypedDict):
    """
    The shared state that flows through every node in the Sisyphus graph.

    Annotated fields use *reducers* — functions that merge concurrent writes.
    Without reducers, parallel node outputs would overwrite each other.
    """

    # Conversation history — the `add_messages` reducer appends new messages
    # instead of replacing the list.  This is critical: every node adds its
    # response here, and downstream nodes see the full history.
    #
    # Loki equivalent: each agent's chat session accumulates messages the same
    # way, but messages are scoped per-agent.  In LangGraph the shared message
    # list IS the inter-agent communication channel.
    messages: Annotated[list[BaseMessage], add_messages]

    # Classified intent for the current request
    intent: IntentType

    # Which agent the supervisor routed to last
    next_agent: str

    # Iteration counter — safety valve analogous to Loki's max_auto_continues
    iteration_count: int

    # Todo list for multi-step tracking (mirrors Loki's todo__* tools)
    todos: Annotated[list[TodoItem], _merge_todos]

    # Accumulated outputs from sub-agent nodes, keyed by agent name.
    # The supervisor reads these to decide what to do next.
    agent_outputs: dict[str, str]

    # Final synthesized answer to return to the user
    final_output: str

    # The working directory / project path (mirrors Loki's project_dir variable)
    project_dir: str