"""Streaming Scheduler - Event-driven execution with streaming nodes.

This scheduler handles pathways with EventSourceNode/IntrospectionNode
that yield events over time, executing downstream nodes for each event.
"""

from __future__ import annotations

import asyncio
import time
from typing import TYPE_CHECKING, Any

from core.common import JsonObject
from core.pathways.models.io import NodeOutputs
from core.pathways.streaming import EventSourceNode, IntrospectionNode
from core.pathways.vm.graph import reachable_from, topological_sort
from core.pathways.vm.schedulers.dag import DAGScheduler
from core.pathways.vm.utils import gather_inputs, should_execute

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from core.pathways.context import Context
    from core.pathways.pathway import Pathway
    from core.pathways.vm.observability import VMEventEmitter


class StreamingScheduler:
    """Scheduler for streaming execution."""

    def __init__(
        self,
        ctx: Context,
        max_parallel: int,
        emitter: VMEventEmitter,
        node_semaphore: asyncio.Semaphore,
    ) -> None:
        """Initialize streaming scheduler.

        Args:
            ctx: Execution context
            max_parallel: Maximum concurrent node executions
            emitter: Event emitter for observability
            node_semaphore: Semaphore for node execution concurrency
        """
        self.ctx = ctx
        self.max_parallel = max_parallel
        self.emitter = emitter
        self._node_semaphore = node_semaphore

    async def execute(
        self,
        pathway: Pathway,
        inputs: JsonObject,
        max_events: int | None = None,
    ) -> AsyncIterator[JsonObject]:
        """Execute a streaming pathway.

        This is **true streaming execution**:
        - Each streaming node yields events over time
        - For each event, the VM executes the downstream reachable subgraph

        Args:
            pathway: The pathway to execute
            inputs: Input data
            max_events: Maximum number of events to process

        Yields
        ------
            Dicts with type, source_node_id, event, and outputs
        """
        # Establish an execution_id for observability parity with execute().
        execution_id = f"exec_{int(time.time() * 1000)}"
        self.ctx.pathway_id = pathway.id
        self.ctx.execution_id = execution_id
        self.emitter.emit(
            "execution_started",
            pathway_id=pathway.id,
            execution_id=execution_id,
            inputs=inputs,
            mode="stream",
        )

        # Identify streaming source nodes (outer + inner).
        streaming_nodes: list[tuple[str, EventSourceNode | IntrospectionNode]] = [
            (node_id, node)
            for node_id, node in pathway.nodes.items()
            if isinstance(node, (EventSourceNode, IntrospectionNode))
        ]
        if not streaming_nodes:
            # No streaming nodes; just run once and yield a single completion.
            # Delegate to DAG scheduler for non-streaming pathways
            dag_scheduler = DAGScheduler(
                self.ctx, self.max_parallel, self.emitter, self._node_semaphore
            )
            record = await dag_scheduler.execute(pathway, inputs, execution_id)
            yield {
                "type": "complete",
                "outputs": record.outputs,
                "success": record.success,
                "error": record.error,
            }
            return

        # Build adjacency for reachability.
        outgoing: dict[str, set[str]] = {}
        for c in pathway.connections:
            outgoing.setdefault(c.from_node, set()).add(c.to_node)

        # Build reachability map for each streaming node
        topo_order = topological_sort(pathway)
        reach_map: dict[str, set[str]] = {
            nid: reachable_from(nid, list(pathway.connections), pathway.nodes)
            for nid, _ in streaming_nodes
        }

        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        tasks: list[asyncio.Task[None]] = []

        async def _pump(node_id: str, node: EventSourceNode | IntrospectionNode) -> None:
            try:
                async for item in node.stream(inputs, self.ctx):
                    # Normalize to an event payload for downstream consumers.
                    ev = item.get("event") if isinstance(item, dict) else item
                    await queue.put({"source_node_id": node_id, "event": ev, "raw": item})
            except asyncio.CancelledError:
                raise
            except Exception as e:
                await queue.put({"source_node_id": node_id, "error": str(e)})

        for node_id, node in streaming_nodes:
            tasks.append(asyncio.create_task(_pump(node_id, node)))

        produced = 0
        try:
            while True:
                if max_events is not None and produced >= max_events:
                    return

                item = await queue.get()
                source_node_id = item.get("source_node_id")
                if item.get("error"):
                    yield {
                        "type": "stream_error",
                        "source_node_id": source_node_id,
                        "error": item.get("error"),
                    }
                    produced += 1
                    continue

                event_payload = item.get("event")

                # Emit synthetic node events for the streaming source (it doesn't run via compute())
                self.emitter.emit(
                    "node_started",
                    node_id=str(source_node_id),
                    pathway_id=pathway.id,
                    execution_id=execution_id,
                    tick_index=produced,
                    source_node_id=source_node_id,
                )
                self.emitter.emit(
                    "node_completed",
                    node_id=str(source_node_id),
                    pathway_id=pathway.id,
                    execution_id=execution_id,
                    success=True,
                    duration_ms=0.0,
                    tick_index=produced,
                    source_node_id=source_node_id,
                )

                # Seed results with the source node output so normal connections work.
                results: dict[str, NodeOutputs] = {
                    str(source_node_id): NodeOutputs(
                        output=event_payload,
                        event=event_payload,
                        source_node_id=source_node_id,
                    )
                }

                # Execute reachable downstream nodes in topo order.
                reachable = reach_map.get(str(source_node_id), set())
                for node_id in topo_order:
                    if node_id == source_node_id:
                        continue
                    if node_id not in reachable:
                        continue

                    node = pathway.nodes.get(node_id)
                    if node is None:
                        continue

                    if not should_execute(pathway, node_id, results):
                        continue

                    node_inputs = gather_inputs(pathway, node_id, results, inputs)
                    self.emitter.emit(
                        "node_started",
                        node_id=node_id,
                        pathway_id=pathway.id,
                        execution_id=execution_id,
                        tick_index=produced,
                        source_node_id=source_node_id,
                    )
                    started_at = time.time()
                    try:
                        self.ctx.run._current_node_id = node_id
                        out = await node.compute(node_inputs, self.ctx)
                        duration_ms = (time.time() - started_at) * 1000
                        results[node_id] = out
                        self.emitter.emit(
                            "node_completed",
                            node_id=node_id,
                            pathway_id=pathway.id,
                            execution_id=execution_id,
                            success=True,
                            duration_ms=duration_ms,
                            tick_index=produced,
                            source_node_id=source_node_id,
                        )
                    except Exception as e:
                        results[node_id] = NodeOutputs(error=str(e))
                        self.emitter.emit(
                            "node_failed",
                            node_id=node_id,
                            pathway_id=pathway.id,
                            error=str(e),
                            execution_id=execution_id,
                            tick_index=produced,
                            source_node_id=source_node_id,
                        )
                        break
                    finally:
                        if self.ctx.run._current_node_id == node_id:
                            self.ctx.run._current_node_id = None

                yield {
                    "type": "tick",
                    "source_node_id": source_node_id,
                    "event": event_payload,
                    "outputs": {k: v.to_dict() for k, v in results.items()},
                }
                produced += 1
        except asyncio.CancelledError:
            self.emitter.emit(
                "execution_completed",
                pathway_id=pathway.id,
                execution_id=execution_id,
                status="cancelled",
                produced=produced,
            )
            raise
        finally:
            for t in tasks:
                t.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
            self.emitter.emit(
                "execution_completed",
                pathway_id=pathway.id,
                execution_id=execution_id,
                status="completed",
                produced=produced,
            )
