"""A lightweight, in-process asynchronous pipeline runtime."""

from __future__ import annotations

from collections import defaultdict, deque
from collections.abc import AsyncIterable, AsyncIterator, Iterable, Mapping
from dataclasses import dataclass
from enum import Enum
from types import MappingProxyType
from typing import Any

import asyncio
import inspect

from .component import (
    DEFAULT_INPUT_PORT,
    DEFAULT_OUTPUT_PORT,
    Component,
    ComponentContext,
    Emission,
    HealthReport,
    Operator,
    Sink,
    Source,
    iter_emissions,
)
from .edge_queue import (
    EdgeQueue,
    EdgeQueueClosed,
    OverflowPolicy,
    QueueStats,
)
from .envelope import Envelope


class RuntimeState(str, Enum):
    CREATED = "created"
    STARTING = "starting"
    RUNNING = "running"
    STOPPING = "stopping"
    STOPPED = "stopped"


class PipelineRuntimeError(RuntimeError):
    """Base exception raised by the pipeline runtime."""


class GraphValidationError(PipelineRuntimeError, ValueError):
    """Raised when nodes and edges do not form a valid executable DAG."""


class NodeExecutionError(PipelineRuntimeError):
    """Wrap an exception with the id of the component that raised it."""

    def __init__(self, node_id: str, cause: BaseException) -> None:
        self.node_id = node_id
        self.cause = cause
        super().__init__(f"node {node_id!r} failed: {cause}")


@dataclass(frozen=True, slots=True)
class Edge:
    """Declarative connection between one output port and one input port."""

    source_node: str
    target_node: str
    source_port: str = DEFAULT_OUTPUT_PORT
    target_port: str = DEFAULT_INPUT_PORT
    capacity: int = 1
    overflow: OverflowPolicy | str = OverflowPolicy.BLOCK
    name: str | None = None

    def __post_init__(self) -> None:
        for field_name in (
            "source_node",
            "target_node",
            "source_port",
            "target_port",
        ):
            value = getattr(self, field_name)
            if not value or not value.strip():
                raise ValueError(f"{field_name} must not be empty")
        if self.capacity < 1:
            raise ValueError("edge capacity must be at least 1")
        if self.name is not None and not self.name.strip():
            raise ValueError("edge name must not be empty")
        try:
            object.__setattr__(self, "overflow", OverflowPolicy(self.overflow))
        except ValueError as exc:
            choices = ", ".join(policy.value for policy in OverflowPolicy)
            raise ValueError(
                f"unknown overflow policy {self.overflow!r}; use {choices}"
            ) from exc

    @property
    def resolved_name(self) -> str:
        return self.name or (
            f"{self.source_node}.{self.source_port}"
            f"->{self.target_node}.{self.target_port}"
        )


@dataclass(slots=True)
class _LiveEdge:
    spec: Edge
    queue: EdgeQueue[Envelope[Any]]


class PipelineRuntime:
    """Execute a configuration-compiled DAG in one asyncio event loop.

    The runtime intentionally does not own threads, processes, transports or
    retry policy.  Those belong to components or a future worker supervisor.
    Every graph edge owns a separate bounded queue, and only emissions matching
    that edge's ``source_port`` are routed to it.

    Runtime instances are one-shot.  Create a new instance to restart a stopped
    pipeline; this avoids accidentally reusing closed queues or partially torn
    down model resources.
    """

    def __init__(
        self,
        nodes: Mapping[str, Component],
        edges: Iterable[Edge],
        *,
        pipeline_id: str = "pipeline",
        shutdown_timeout_s: float = 10.0,
        metadata: Mapping[str, Any] | None = None,
    ) -> None:
        if not pipeline_id or not pipeline_id.strip():
            raise ValueError("pipeline_id must not be empty")
        if shutdown_timeout_s < 0:
            raise ValueError("shutdown_timeout_s must not be negative")

        self._pipeline_id = pipeline_id
        self._nodes = dict(nodes)
        self._edge_specs = tuple(edges)
        self._shutdown_timeout_s = shutdown_timeout_s
        self._metadata = MappingProxyType(dict(metadata or {}))
        self._state = RuntimeState.CREATED
        self._shutdown_event = asyncio.Event()
        self._stop_lock = asyncio.Lock()
        self._tasks: dict[str, asyncio.Task[None]] = {}
        self._setup_nodes: list[str] = []
        self._stopped_nodes: set[str] = set()

        self._validate_graph()
        self._topological_order = self._compute_topological_order()
        self._live_edges = tuple(
            _LiveEdge(
                spec=edge,
                queue=EdgeQueue(
                    edge.capacity,
                    edge.overflow,
                    name=edge.resolved_name,
                ),
            )
            for edge in self._edge_specs
        )
        self._outgoing: dict[tuple[str, str], tuple[_LiveEdge, ...]] = (
            self._index_outgoing()
        )
        self._incoming: dict[str, tuple[_LiveEdge, ...]] = self._index_incoming()

    @property
    def pipeline_id(self) -> str:
        return self._pipeline_id

    @property
    def state(self) -> RuntimeState:
        return self._state

    @property
    def nodes(self) -> Mapping[str, Component]:
        return MappingProxyType(self._nodes)

    @property
    def edges(self) -> tuple[Edge, ...]:
        return self._edge_specs

    @property
    def shutdown_event(self) -> asyncio.Event:
        return self._shutdown_event

    @property
    def edge_queues(self) -> Mapping[str, EdgeQueue[Envelope[Any]]]:
        return MappingProxyType(
            {edge.spec.resolved_name: edge.queue for edge in self._live_edges}
        )

    def edge_stats(self) -> dict[str, QueueStats]:
        """Return snapshots keyed by resolved edge name."""

        return {
            edge.spec.resolved_name: edge.queue.stats() for edge in self._live_edges
        }

    async def health(self) -> dict[str, HealthReport]:
        """Collect health from all nodes without failing the whole query."""

        async def check(node_id: str, component: Component) -> tuple[str, HealthReport]:
            try:
                report = await asyncio.wait_for(component.health(), timeout=1.0)
                task = self._tasks.get(node_id)
                if task is not None and task.done():
                    if task.cancelled():
                        report = HealthReport(False, "runtime task was cancelled")
                    else:
                        failure = task.exception()
                        if failure is not None:
                            report = HealthReport(
                                False, f"runtime task failed: {failure}"
                            )
                        elif self._state is RuntimeState.RUNNING:
                            report = HealthReport(
                                False,
                                "runtime task completed while pipeline is running",
                            )
            except Exception as exc:  # health is diagnostic, not pipeline control
                report = HealthReport(False, f"health check failed: {exc}")
            return node_id, report

        reports = await asyncio.gather(
            *(check(node_id, component) for node_id, component in self._nodes.items())
        )
        return dict(reports)

    async def start(self) -> None:
        """Run lifecycle hooks and start one task per component."""

        async with self._stop_lock:
            if self._state is not RuntimeState.CREATED:
                raise PipelineRuntimeError(
                    f"cannot start pipeline in state {self._state.value!r}"
                )
            self._state = RuntimeState.STARTING

            try:
                # Prepare downstream consumers before upstream producers.
                for node_id in reversed(self._topological_order):
                    component = self._nodes[node_id]
                    context = ComponentContext(
                        pipeline_id=self._pipeline_id,
                        node_id=node_id,
                        shutdown_event=self._shutdown_event,
                        metadata=dict(self._metadata),
                    )
                    self._setup_nodes.append(node_id)
                    await component.setup(context)

                for node_id in reversed(self._topological_order):
                    component = self._nodes[node_id]
                    await component.start()
            except BaseException:
                self._shutdown_event.set()
                deadline = asyncio.get_running_loop().time() + self._shutdown_timeout_s
                await self._close_all_edges(discard=True)
                await self._stop_components(deadline=deadline)
                self._state = RuntimeState.STOPPED
                raise

            self._tasks = {
                node_id: asyncio.create_task(
                    self._run_node(node_id, self._nodes[node_id]),
                    name=f"{self._pipeline_id}:{node_id}",
                )
                for node_id in reversed(self._topological_order)
            }
            self._state = RuntimeState.RUNNING

    async def wait(self) -> None:
        """Wait for natural completion, propagating the first node failure."""

        if self._state not in (RuntimeState.RUNNING, RuntimeState.STOPPING):
            if self._state is RuntimeState.STOPPED:
                return
            raise PipelineRuntimeError(
                f"cannot wait for pipeline in state {self._state.value!r}"
            )
        try:
            await asyncio.gather(*self._tasks.values())
        except asyncio.CancelledError:
            raise
        except BaseException:
            await self.stop(graceful=False)
            raise

    async def run(self) -> None:
        """Start, wait, and always execute lifecycle teardown."""

        await self.start()
        try:
            await self.wait()
        finally:
            await self.stop()

    async def stop(
        self,
        *,
        graceful: bool = True,
        timeout_s: float | None = None,
    ) -> None:
        """Stop sources, optionally drain downstream queues, then tear down.

        Graceful shutdown cancels source tasks first.  Their output queues close
        in ``finally`` blocks, allowing operators and sinks to drain all already
        accepted messages.  If draining exceeds the timeout, remaining tasks
        are cancelled and queued messages are discarded.
        """

        if timeout_s is not None and timeout_s < 0:
            raise ValueError("timeout_s must not be negative")

        async with self._stop_lock:
            if self._state is RuntimeState.STOPPED:
                return
            if self._state is RuntimeState.CREATED:
                self._state = RuntimeState.STOPPED
                self._shutdown_event.set()
                await self._close_all_edges(discard=True)
                return

            self._state = RuntimeState.STOPPING
            self._shutdown_event.set()
            effective_timeout = (
                self._shutdown_timeout_s if timeout_s is None else timeout_s
            )
            deadline = asyncio.get_running_loop().time() + effective_timeout

            if graceful:
                source_ids = tuple(
                    node_id
                    for node_id, component in self._nodes.items()
                    if isinstance(component, Source)
                )
                self._cancel_tasks(source_ids)
                # Source hooks close external subscriptions promptly.  Operators
                # and sinks remain alive until accepted messages have drained.
                await self._stop_components(source_ids, deadline=deadline)
                drained = await self._wait_for_tasks(self._remaining(deadline))
                if not drained:
                    self._cancel_tasks()
                    await self._close_all_edges(discard=True)
                    await self._stop_components(deadline=deadline)
                await self._join_tasks(self._remaining(deadline))
            else:
                self._cancel_tasks()
                await self._close_all_edges(discard=True)
                await self._stop_components(deadline=deadline)
                await self._join_tasks(self._remaining(deadline))

            await self._stop_components(deadline=deadline)
            self._tasks.clear()
            self._state = RuntimeState.STOPPED

    async def _run_node(self, node_id: str, component: Component) -> None:
        try:
            if isinstance(component, Source):
                await self._run_source(node_id, component)
            elif isinstance(component, Operator):
                await self._run_operator(node_id, component)
            elif isinstance(component, Sink):
                await self._run_sink(node_id, component)
            else:  # guarded by graph validation
                raise TypeError("unsupported component kind")
        except asyncio.CancelledError:
            raise
        except BaseException as exc:
            raise NodeExecutionError(node_id, exc) from exc
        finally:
            await self._close_outgoing(node_id)

    async def _run_source(self, node_id: str, source: Source) -> None:
        stream = source.messages()
        if inspect.isawaitable(stream):
            # Be tolerant of an ``async def`` factory that returns an iterator.
            stream = await stream
        if not isinstance(stream, AsyncIterable):
            raise TypeError("Source.messages() must return an async iterable")

        async for item in stream:
            async for emission in iter_emissions(item):
                await self._route(node_id, emission)

    async def _run_operator(self, node_id: str, operator: Operator) -> None:
        async for live_edge, envelope in self._receive(node_id):
            result = await operator.process(
                envelope,
                input_port=live_edge.spec.target_port,
            )
            async for emission in iter_emissions(result):
                await self._route(node_id, emission)

    async def _run_sink(self, node_id: str, sink: Sink) -> None:
        async for live_edge, envelope in self._receive(node_id):
            await sink.consume(
                envelope,
                input_port=live_edge.spec.target_port,
            )

    async def _route(self, source_node: str, emission: Emission) -> None:
        edges = self._outgoing.get((source_node, emission.port), ())
        if not edges:
            return

        results = await asyncio.gather(
            *(edge.queue.put(emission.envelope) for edge in edges),
            return_exceptions=True,
        )
        for result in results:
            if isinstance(result, EdgeQueueClosed) and self._shutdown_event.is_set():
                continue
            if isinstance(result, BaseException):
                raise result

    async def _receive(
        self,
        node_id: str,
    ) -> AsyncIterator[tuple[_LiveEdge, Envelope[Any]]]:
        incoming = self._incoming.get(node_id, ())
        pending: dict[asyncio.Task[Envelope[Any]], _LiveEdge] = {
            asyncio.create_task(
                edge.queue.get(),
                name=f"{self._pipeline_id}:{node_id}:read:{edge.spec.resolved_name}",
            ): edge
            for edge in incoming
        }

        try:
            while pending:
                done, _ = await asyncio.wait(
                    pending,
                    return_when=asyncio.FIRST_COMPLETED,
                )
                for task in done:
                    edge = pending.pop(task)
                    try:
                        envelope = task.result()
                    except EdgeQueueClosed:
                        continue

                    yield edge, envelope

                    if not (edge.queue.closed and edge.queue.empty()):
                        next_read = asyncio.create_task(
                            edge.queue.get(),
                            name=(
                                f"{self._pipeline_id}:{node_id}:read:"
                                f"{edge.spec.resolved_name}"
                            ),
                        )
                        pending[next_read] = edge
        finally:
            for task in pending:
                task.cancel()
            if pending:
                await asyncio.gather(*pending, return_exceptions=True)

    async def _close_outgoing(self, node_id: str) -> None:
        queues = [
            live_edge.queue
            for (source, _), edges in self._outgoing.items()
            if source == node_id
            for live_edge in edges
        ]
        if queues:
            await asyncio.gather(*(queue.close() for queue in queues))

    async def _close_all_edges(self, *, discard: bool) -> None:
        if self._live_edges:
            await asyncio.gather(
                *(edge.queue.close(discard=discard) for edge in self._live_edges)
            )

    async def _wait_for_tasks(self, timeout_s: float) -> bool:
        if not self._tasks:
            return True
        _, pending = await asyncio.wait(
            self._tasks.values(),
            timeout=timeout_s,
            return_when=asyncio.ALL_COMPLETED,
        )
        return not pending

    def _cancel_tasks(self, node_ids: Iterable[str] | None = None) -> None:
        selected = self._tasks.keys() if node_ids is None else node_ids
        for node_id in selected:
            task = self._tasks.get(node_id)
            if task is None:
                continue
            if not task.done():
                task.cancel()

    async def _join_tasks(self, timeout_s: float) -> bool:
        tasks = tuple(self._tasks.values())
        if not tasks:
            return True
        _, pending = await asyncio.wait(
            tasks,
            timeout=timeout_s,
            return_when=asyncio.ALL_COMPLETED,
        )
        for task in pending:
            task.cancel()
        return not pending

    async def _stop_components(
        self,
        node_ids: Iterable[str] | None = None,
        *,
        deadline: float | None = None,
    ) -> None:
        # setup may allocate resources even when a later start hook fails.
        selected = None if node_ids is None else set(node_ids)
        setup_order = list(dict.fromkeys(self._setup_nodes))
        for node_id in reversed(setup_order):
            if selected is not None and node_id not in selected:
                continue
            if node_id in self._stopped_nodes:
                continue
            try:
                if deadline is None:
                    await self._nodes[node_id].stop()
                else:
                    remaining = self._remaining(deadline)
                    if remaining <= 0:
                        continue
                    await asyncio.wait_for(
                        self._nodes[node_id].stop(), timeout=remaining
                    )
            except Exception:
                # Teardown is best effort so one broken connector cannot leak all
                # remaining resources.  Supervisors should surface node health/logs.
                pass
            finally:
                self._stopped_nodes.add(node_id)
        if selected is None:
            self._setup_nodes.clear()

    @staticmethod
    def _remaining(deadline: float) -> float:
        return max(0.0, deadline - asyncio.get_running_loop().time())

    def _index_outgoing(self) -> dict[tuple[str, str], tuple[_LiveEdge, ...]]:
        mutable: dict[tuple[str, str], list[_LiveEdge]] = defaultdict(list)
        for edge in self._live_edges:
            mutable[(edge.spec.source_node, edge.spec.source_port)].append(edge)
        return {key: tuple(value) for key, value in mutable.items()}

    def _index_incoming(self) -> dict[str, tuple[_LiveEdge, ...]]:
        mutable: dict[str, list[_LiveEdge]] = defaultdict(list)
        for edge in self._live_edges:
            mutable[edge.spec.target_node].append(edge)
        return {key: tuple(value) for key, value in mutable.items()}

    def _validate_graph(self) -> None:
        if not self._nodes:
            raise GraphValidationError("pipeline must contain at least one node")

        for node_id, component in self._nodes.items():
            if not node_id or not node_id.strip():
                raise GraphValidationError("node id must not be empty")
            kinds = sum(
                isinstance(component, kind) for kind in (Source, Operator, Sink)
            )
            if kinds != 1:
                raise GraphValidationError(
                    f"node {node_id!r} must implement exactly one of "
                    "Source, Operator, or Sink"
                )

        edge_names: set[str] = set()
        edge_keys: set[tuple[str, str, str, str]] = set()
        incoming_count: dict[str, int] = defaultdict(int)
        outgoing_count: dict[str, int] = defaultdict(int)
        adjacency: dict[str, set[str]] = defaultdict(set)
        indegree = {node_id: 0 for node_id in self._nodes}

        for edge in self._edge_specs:
            if edge.source_node not in self._nodes:
                raise GraphValidationError(
                    f"edge {edge.resolved_name!r} references unknown source node "
                    f"{edge.source_node!r}"
                )
            if edge.target_node not in self._nodes:
                raise GraphValidationError(
                    f"edge {edge.resolved_name!r} references unknown target node "
                    f"{edge.target_node!r}"
                )
            if edge.resolved_name in edge_names:
                raise GraphValidationError(
                    f"duplicate edge name {edge.resolved_name!r}"
                )
            edge_names.add(edge.resolved_name)

            edge_key = (
                edge.source_node,
                edge.source_port,
                edge.target_node,
                edge.target_port,
            )
            if edge_key in edge_keys:
                raise GraphValidationError(f"duplicate edge {edge.resolved_name!r}")
            edge_keys.add(edge_key)
            incoming_count[edge.target_node] += 1
            outgoing_count[edge.source_node] += 1

            if edge.target_node not in adjacency[edge.source_node]:
                adjacency[edge.source_node].add(edge.target_node)
                indegree[edge.target_node] += 1

        for node_id, component in self._nodes.items():
            if isinstance(component, Source) and incoming_count[node_id]:
                raise GraphValidationError(
                    f"source node {node_id!r} cannot have incoming edges"
                )
            if isinstance(component, (Operator, Sink)) and not incoming_count[node_id]:
                raise GraphValidationError(
                    f"{type(component).__name__.lower()} node {node_id!r} "
                    "must have at least one incoming edge"
                )
            if isinstance(component, Sink) and outgoing_count[node_id]:
                raise GraphValidationError(
                    f"sink node {node_id!r} cannot have outgoing edges"
                )

        ready = deque(node_id for node_id, degree in indegree.items() if degree == 0)
        visited = 0
        while ready:
            node_id = ready.popleft()
            visited += 1
            for target in adjacency[node_id]:
                indegree[target] -= 1
                if indegree[target] == 0:
                    ready.append(target)
        if visited != len(self._nodes):
            raise GraphValidationError("pipeline graph must be acyclic")

    def _compute_topological_order(self) -> tuple[str, ...]:
        adjacency: dict[str, set[str]] = defaultdict(set)
        indegree = {node_id: 0 for node_id in self._nodes}
        for edge in self._edge_specs:
            if edge.target_node not in adjacency[edge.source_node]:
                adjacency[edge.source_node].add(edge.target_node)
                indegree[edge.target_node] += 1

        ready = deque(node_id for node_id, degree in indegree.items() if degree == 0)
        ordered: list[str] = []
        while ready:
            node_id = ready.popleft()
            ordered.append(node_id)
            for target in adjacency[node_id]:
                indegree[target] -= 1
                if indegree[target] == 0:
                    ready.append(target)
        return tuple(ordered)


__all__ = [
    "Edge",
    "GraphValidationError",
    "NodeExecutionError",
    "PipelineRuntime",
    "PipelineRuntimeError",
    "RuntimeState",
]
