"""LoopNode — iteration until condition.

Executes body nodes inline within ``compute()``, so the DAG scheduler
does not need cycle support.  The body nodes are executed sequentially
in each iteration, with the previous iteration's output fed back as
input to the condition evaluation.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Literal

from pydantic import Field

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 LoopOutput

if TYPE_CHECKING:
    from core.pathways.context import Context

logger = logging.getLogger(__name__)


class LoopNode(NodeBase):
    """Loop iteration — executes body nodes until condition is false.

    The body nodes are executed inline within ``compute()`` so that
    the DAG scheduler (which rejects cycles) can handle loops.

    Attributes
    ----------
    condition
        Expression evaluated before each iteration.  Loop continues
        while this is truthy.  Supports ``{{var}}`` templates.
    body
        List of node IDs to execute in each iteration.
    max_iterations
        Safety limit on the number of iterations.
    collect_results
        Whether to collect per-iteration results into the output.
    """

    type: Literal["loop"] = "loop"
    condition: str
    body: list[str] = Field(default_factory=list)
    max_iterations: int = 10
    collect_results: bool = True

    async def compute(self, inputs: NodeInputs, ctx: Context) -> NodeOutputs:
        """Execute the loop: evaluate condition, run body, repeat."""
        template_data = inputs.to_dict()
        iteration_results: list[dict[str, Any]] = []

        # Get the pathway and its nodes from context
        from core.pathways.service import PathwayService

        pathway_service: PathwayService | None = (
            ctx.services.pathway_service if ctx.services else None
        )
        current_pathway = None
        if pathway_service and ctx.pathway_id:
            current_pathway = pathway_service.load(ctx.pathway_id)

        for iteration in range(self.max_iterations):
            # Evaluate condition
            template_data["_iteration"] = iteration
            template_data["_iteration_results"] = iteration_results
            formatted_condition = self._format_template(self.condition, template_data)

            try:
                should_continue = bool(safe_eval(formatted_condition, env=template_data))
            except Exception as e:
                logger.warning(
                    "[LoopNode %s] Condition evaluation failed at iteration %d: %s",
                    self.id,
                    iteration,
                    e,
                )
                should_continue = False

            if not should_continue:
                break

            # Execute body nodes inline
            iter_output: dict[str, Any] = {"iteration": iteration}
            if current_pathway and self.body:
                for body_node_id in self.body:
                    body_node = current_pathway.nodes.get(body_node_id)
                    if body_node is None:
                        logger.warning(
                            "[LoopNode %s] Body node '%s' not found in pathway",
                            self.id,
                            body_node_id,
                        )
                        continue

                    try:
                        body_inputs = NodeInputs.from_dict({**template_data, **iter_output})
                        body_output = await body_node.compute(body_inputs, ctx)
                        body_data = body_output.to_dict()
                        iter_output[body_node_id] = body_data
                        # Feed body output into template data for next node / next iteration
                        template_data[body_node_id] = body_data
                    except Exception as e:
                        logger.warning(
                            "[LoopNode %s] Body node '%s' failed at iteration %d: %s",
                            self.id,
                            body_node_id,
                            iteration,
                            e,
                        )
                        iter_output[body_node_id] = {"error": str(e)}

            if self.collect_results:
                iteration_results.append(iter_output)

        loop_output = LoopOutput(
            should_continue=False,
            iteration=len(iteration_results),
            body_nodes=self.body,
            reason="completed"
            if len(iteration_results) < self.max_iterations
            else "max_iterations_reached",
        )

        output: dict[str, Any] = loop_output.model_dump()
        output["iterations_completed"] = len(iteration_results)
        if self.collect_results:
            output["results"] = iteration_results

        return NodeOutputs(
            output=output,
            continue_loop=False,
        )


__all__ = ["LoopNode"]
