"""StageNode — orchestration stage for stateful pathways."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Literal

from pydantic import Field, model_validator

from core.pathways.expressions import safe_eval
from core.pathways.models.io import NodeInputs, NodeOutputs
from core.pathways.nodes.base import NodeBase
from core.pathways.nodes.models import StageOutput
from core.tools import get_tool_registry

if TYPE_CHECKING:
    from core.common import JsonObject
    from core.pathways.context import Context


class StageNode(NodeBase):
    """Orchestration stage — a state in a stateful pathway.

    Represents a distinct phase in a multi-turn pathway with
    an associated agent and conditional transitions.

    Attributes
    ----------
    agent_id : str | None
        Agent to run for this stage.
    goal : str
        What this stage should accomplish.
    transitions : list[tuple[str, str]]
        List of (condition, target_stage) pairs.
    on_enter : list[str]
        Tool calls to run when entering the stage.
    on_exit : list[str]
        Tool calls to run when leaving the stage.
    max_iterations : int
        Safety limit to prevent infinite loops.
    """

    type: Literal["stage"] = "stage"
    agent_id: str | None = None
    goal: str = ""
    transitions: list[tuple[str, str]] = Field(default_factory=list)

    @model_validator(mode="before")
    @classmethod
    def _normalize_transitions(cls, data: Any) -> Any:
        """Convert list-of-lists (from YAML) to list-of-tuples."""
        if not isinstance(data, dict):
            return data
        raw = data.get("transitions")
        if isinstance(raw, list):
            data["transitions"] = [
                (str(t[0]), str(t[1])) for t in raw if isinstance(t, (list, tuple)) and len(t) == 2
            ]
        return data

    on_enter: list[str] = Field(default_factory=list)
    on_exit: list[str] = Field(default_factory=list)
    max_iterations: int = 10

    async def compute(self, inputs: NodeInputs, ctx: Context) -> NodeOutputs:
        """Execute the stage — run agent and determine next transition."""
        template_data = inputs.to_dict()
        formatted_goal = self._format_template(self.goal, template_data)

        agent_result: JsonObject = {}
        if self.agent_id:
            registry = get_tool_registry()
            tool_def = registry.get("agent.turn")
            turn_handler = tool_def.handler if tool_def else None
            if turn_handler:
                agent_result = await turn_handler(
                    agent_id=self.agent_id,
                    message=formatted_goal,
                    context=ctx,
                )

        eval_context: JsonObject = {**template_data, "stage_result": agent_result, **agent_result}

        next_stage: str | None = None
        transition_reason: str | None = None
        for condition, target in self.transitions:
            formatted_condition = self._format_template(condition, eval_context)
            if safe_eval(formatted_condition, env=eval_context):
                next_stage = target
                transition_reason = condition
                break

        stage_output = StageOutput(
            stage_id=self.id,
            goal=formatted_goal,
            agent_result=agent_result,
            next_stage=next_stage,
            transition_reason=transition_reason,
            success=agent_result.get("success", True) if isinstance(agent_result, dict) else True,
        )
        return NodeOutputs(
            output=stage_output.model_dump(),
            selected_route=next_stage,
        )


__all__ = ["StageNode"]
