"""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.
"""

import asyncio
import time
from collections.abc import AsyncIterator
from typing import Any

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


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,
            )
