"""ForEachNode — iterate over a list and run a sub-pathway per item.

A lightweight batch primitive for workflows that iterate over
small-to-medium lists.  Not a full pipeline framework — for heavy ETL
use Prefect, Dagster, or plain async Python.

Usage in PathwayBuilder::

    pipeline = (
        PathwayBuilder("process_jobs")
        .input({"jobs": "list"})
        .for_each("process", items="input.jobs", body=invoice_pathway, on_error="skip")
        .output({"results": "{{process.output}}"})
        .build()
    )
"""

from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING, Any, Literal

from core.pathways.models.io import NodeInputs, NodeOutputs
from core.pathways.nodes.base import NodeBase

if TYPE_CHECKING:
    from core.pathways.context import Context

logger = logging.getLogger(__name__)


class ForEachNode(NodeBase):
    """Iterate over a list and execute a sub-pathway for each item.

    Attributes
    ----------
    items_expr
        Expression that resolves to a list from upstream inputs
        (e.g. ``"input.jobs"`` or ``"fetch.output.data.records"``).
    body_pathway_id
        ID of the sub-pathway to execute for each item.
    on_error
        ``"fail"`` aborts the batch on first error.
        ``"skip"`` skips the failed item and continues.
    max_parallel
        Maximum number of items to process concurrently.
    collect_results
        Whether to collect per-item results into the output list.
    """

    type: Literal["for_each"] = "for_each"
    items_expr: str
    body_pathway_id: str
    on_error: Literal["fail", "skip"] = "skip"
    max_parallel: int = 1
    collect_results: bool = True

    async def compute(self, inputs: NodeInputs, ctx: Context) -> NodeOutputs:
        """Iterate over items and run the body pathway for each."""
        template_data = inputs.to_dict()

        # Resolve the items list from the expression
        items = self._resolve_items(template_data)
        if not isinstance(items, list):
            return NodeOutputs(
                output={"error": f"items_expr '{self.items_expr}' did not resolve to a list"},
                error=f"items_expr resolved to {type(items).__name__}, expected list",
            )

        # Load the body pathway
        pathway_service = ctx.services.pathway_service if ctx.services else None
        if pathway_service is None:
            return NodeOutputs(
                output={"error": "No pathway service available"},
                error="PathwayService not available in context",
            )

        body_pathway = pathway_service.load(self.body_pathway_id)
        if body_pathway is None:
            return NodeOutputs(
                output={"error": f"Body pathway not found: {self.body_pathway_id}"},
                error=f"Pathway not found: {self.body_pathway_id}",
            )

        # Get the PathwayVM from context for sub-pathway execution
        # The VM is not directly on ctx, so we create a lightweight executor
        from core.pathways.vm.executor import PathwayVM

        vm = PathwayVM(ctx)

        # Execute items with concurrency control
        semaphore = asyncio.Semaphore(max(1, self.max_parallel))
        results: list[dict[str, Any]] = []
        errors: list[dict[str, Any]] = []

        async def _process_item(index: int, item: Any) -> dict[str, Any] | None:
            async with semaphore:
                try:
                    record = await vm.execute(
                        body_pathway,
                        inputs={"item": item, "index": index},
                    )
                    return {
                        "index": index,
                        "success": record.success,
                        "outputs": record.outputs,
                    }
                except Exception as e:
                    err = {
                        "index": index,
                        "success": False,
                        "error": str(e),
                    }
                    if self.on_error == "fail":
                        raise
                    return err

        if self.max_parallel <= 1:
            # Sequential execution
            for i, item in enumerate(items):
                try:
                    result = await _process_item(i, item)
                    if result is not None:
                        if result["success"]:
                            results.append(result)
                        else:
                            errors.append(result)
                except Exception as e:
                    errors.append({"index": i, "success": False, "error": str(e)})
                    if self.on_error == "fail":
                        break
        else:
            # Parallel execution
            tasks = [asyncio.create_task(_process_item(i, item)) for i, item in enumerate(items)]
            for coro in asyncio.as_completed(tasks):
                try:
                    result = await coro
                    if result is not None:
                        if result["success"]:
                            results.append(result)
                        else:
                            errors.append(result)
                except Exception as e:
                    errors.append({"success": False, "error": str(e)})
                    if self.on_error == "fail":
                        # Cancel remaining tasks
                        for t in tasks:
                            if not t.done():
                                t.cancel()
                        break

        output: dict[str, Any] = {
            "total": len(items),
            "succeeded": len(results),
            "failed": len(errors),
            "errors": errors,
        }
        if self.collect_results:
            output["results"] = results

        return NodeOutputs(output=output)

    def _resolve_items(self, data: dict[str, Any]) -> Any:
        """Resolve items_expr against template data."""
        # Try dotted path resolution first
        parts = self.items_expr.split(".")
        current: Any = data
        for part in parts:
            if isinstance(current, dict):
                current = current.get(part)
            else:
                return None
        return current


__all__ = ["ForEachNode"]
