from __future__ import annotations

import contextvars
import threading
from typing import Any, Callable, TYPE_CHECKING

from flownet.runtime.contexts.actor import _system_calls_scope

if TYPE_CHECKING:
    from flownet.runtime.sources.stream_sources import StreamSourceRegistry
    from flownet.runtime.topics.topic_registry import TopicRegistry


_stream_source_context: contextvars.ContextVar["StreamSourceContext" | None] = (
    contextvars.ContextVar("flownet_stream_source_context", default=None)
)


class StreamSourceContext:
    def __init__(
        self,
        source_id: str,
        registry: "StreamSourceRegistry",
        topic_registry: "TopicRegistry",
        call_manager: Any,
        local_address_fn: Callable[[], str],
        stop_event: threading.Event,
        request_id: str,
        mark_done_fn: Callable[[str], Any] | None = None,
    ):
        self._source_id = source_id
        self._registry = registry
        self._topic_registry = topic_registry
        self._call_manager = call_manager
        self._local_address_fn = local_address_fn
        self._stop_event = stop_event
        self._request_id = request_id
        self._mark_done_fn = mark_done_fn
        # Monotonic source-local sequence used for future watermark-based
        # done semantics in reducers.
        self._event_seq = 0

    @property
    def source_id(self) -> str:
        return self._source_id

    def emit(
        self,
        event_type: str,
        payload: Any = None,
        tags: dict[str, str] | None = None,
    ) -> None:
        with _system_calls_scope(True):
            if tags is None:
                tags = {}
            tags.setdefault("request_id", self._request_id)
            tags.setdefault("request_ref_id", self._request_id)
            # Explicit request-root tracking marker for topic coordinator.
            # This decouples tracker routing from request-id prefix parsing.
            tags.setdefault("__track_request_root__", "1")
            tags.setdefault("__request_tracker_origin__", self.local_address())
            is_done = (
                event_type == "__done__"
                or tags.get("__control__") == "done"
            )
            # TODO(sequence-watermark):
            # End-to-end sequence ordering is not enforced yet because frame
            # metadata is not propagated through all operator transitions.
            # We still stamp source-side sequence/final-sequence now so
            # topic/done routing and future reducers can converge on one wire
            # contract without another source API change.
            if is_done:
                tags.setdefault("__final_seq__", str(self._event_seq))
            else:
                self._event_seq += 1
                tags.setdefault("__seq__", str(self._event_seq))
            self._topic_registry.publish(
                self._source_id,
                event_type,
                payload,
                tags,
                producer_id=self._source_id,
            )
            if is_done and self._mark_done_fn is not None:
                self._mark_done_fn(self._request_id)

    def local_address(self) -> str:
        return self._local_address_fn()

    def should_stop(self) -> bool:
        return self._stop_event.is_set()

    @property
    def stop_event(self) -> threading.Event:
        return self._stop_event

    def submit(self, coro: Any, *, loop: str = "user"):
        return self._registry.submit(coro, loop=loop)


def stream_source_context() -> StreamSourceContext:
    context = _stream_source_context.get()
    if context is None:
        raise RuntimeError("Stream source context not available. Use within a source runner.")
    return context


def emit(
    event_type: str,
    payload: Any = None,
    tags: dict[str, str] | None = None,
) -> None:
    stream_source_context().emit(event_type, payload, tags)


__all__ = [
    "StreamSourceContext",
    "stream_source_context",
    "emit",
]
