from __future__ import annotations

from dataclasses import dataclass
import inspect
import threading
import uuid
from typing import Any, Callable

from flownet.core.stream_event import StreamEvent
from flownet.runtime.contexts.actor import _user_execution_scope
from flownet.runtime.contexts.source import (
    StreamSourceContext,
    stream_source_context,
    emit,
    _stream_source_context,
)
from flownet.runtime.loops import submit_runtime_coro
from flownet.runtime.runner_lifecycle import start_runner_lifecycle
from flownet.runtime.topics.topic_registry import TopicRegistry


 


@dataclass
class StreamSourceHandle:
    source_id: str
    source: Any
    status: str = "running"
    restart_policy: str = "none"
    last_error: Exception | None = None
    runner: "StreamSourceRunner" | None = None


class StreamSourceRunner:
    def __init__(
        self,
        source_id: str,
        stop_event: threading.Event,
        thread: threading.Thread | None = None,
        future: Any | None = None,
        close_fn: Callable[[], None] | None = None,
    ):
        self._source_id = source_id
        self._stop_event = stop_event
        self._thread = thread
        self._future = future
        self._close_fn = close_fn

    def stop(self) -> None:
        self._stop_event.set()
        if self._close_fn is not None:
            try:
                self._close_fn()
            except Exception as exc:  # pragma: no cover - best effort cleanup
                print(f"[StreamSourceRunner] Failed to close {self._source_id}: {exc}")
        if self._future is not None:
            try:
                self._future.cancel()
            except Exception as exc:  # pragma: no cover - best effort
                print(f"[StreamSourceRunner] Failed to cancel {self._source_id}: {exc}")
        if self._thread is not None and self._thread.is_alive():
            self._thread.join(timeout=1.0)


def _emit_item(context: StreamSourceContext, item: Any) -> None:
    if item is None:
        return
    if isinstance(item, StreamEvent):
        context.emit(item.event_type, item.payload, dict(item.tags))
        return
    if isinstance(item, tuple):
        if len(item) == 2:
            event_type, payload = item
            tags = None
        elif len(item) == 3:
            event_type, payload, tags = item
        else:
            raise ValueError(f"Invalid stream event tuple: {item}")
        context.emit(event_type, payload, tags)
        return
    context.emit("event", item, None)


def _emit_done(context: StreamSourceContext) -> None:
    context.emit("__done__", None, {"__control__": "done"})


def _call_source(
    source_spec: Any,
    *args,
    **kwargs,
) -> tuple[Any, Any]:
    if inspect.isclass(source_spec):
        source_obj = source_spec(*args, **kwargs)
        result = source_obj()
        return source_obj, result
    if inspect.isgenerator(source_spec) or inspect.isasyncgen(source_spec):
        return source_spec, source_spec
    if not callable(source_spec):
        raise TypeError(f"Stream source spec is not callable: {type(source_spec)}")
    result = source_spec(*args, **kwargs)
    return source_spec, result


class StreamSourceRegistry:
    def __init__(
        self,
        call_manager: Any,
        topic_registry: TopicRegistry,
        local_address_fn: Callable[[], str],
        submit_coro: Callable[[Any], Any] = submit_runtime_coro,
        submit_user_coro: Callable[[Any], Any] | None = None,
        request_id_factory: Callable[[str], str] | None = None,
        mark_done_fn: Callable[[str], Any] | None = None,
    ):
        self._call_manager = call_manager
        self._topic_registry = topic_registry
        self._local_address_fn = local_address_fn
        self._submit_coro = submit_coro
        self._submit_user_coro = submit_user_coro or submit_coro
        self._request_id_factory = request_id_factory
        self._mark_done_fn = mark_done_fn
        self._sources: dict[str, StreamSourceHandle] = {}
        self._lock = threading.Lock()

    def register_source(
        self,
        source_id: str,
        source_spec: Any,
        *args,
        restart_policy: str = "none",
        **kwargs,
    ) -> StreamSourceHandle:
        with self._lock:
            if source_id in self._sources:
                raise ValueError(f"Stream source already running: {source_id}")

        stop_event = threading.Event()
        request_id = self._build_request_id(source_id)
        context = StreamSourceContext(
            source_id=source_id,
            registry=self,
            topic_registry=self._topic_registry,
            call_manager=self._call_manager,
            local_address_fn=self._local_address_fn,
            stop_event=stop_event,
            request_id=request_id,
            mark_done_fn=self._mark_done_fn,
        )

        source_obj, result = _call_source(source_spec, *args, **kwargs)

        handle = StreamSourceHandle(
            source_id=source_id,
            source=source_obj,
            restart_policy=restart_policy,
        )
        with self._lock:
            self._sources[source_id] = handle

        try:
            runner = self._start_runner(source_id, context, stop_event, result, handle)
            handle.runner = runner
        except Exception:
            with self._lock:
                self._sources.pop(source_id, None)
            raise
        return handle

    def start_source(
        self,
        source_id: str,
        source_spec: Any,
        *args,
        **kwargs,
    ) -> StreamSourceHandle:
        return self.register_source(source_id, source_spec, *args, **kwargs)

    def _start_runner(
        self,
        source_id: str,
        context: StreamSourceContext,
        stop_event: threading.Event,
        result: Any,
        handle: StreamSourceHandle,
    ) -> StreamSourceRunner | None:
        def _on_stopped(runner_kind: str) -> None:
            if runner_kind in {"asyncgen", "generator"}:
                if not stop_event.is_set():
                    _emit_done(context)
            else:
                _emit_done(context)
            self._update_status(source_id, "stopped", None)

        def _on_cancelled() -> None:
            self._update_status(source_id, "stopped", None)

        def _on_error(exc: Exception) -> None:
            self._update_status(source_id, "error", exc)
            print(f"[StreamSourceRegistry] Source {source_id} failed: {exc}")

        started = start_runner_lifecycle(
            name=source_id,
            result=result,
            stop_event=stop_event,
            submit_user_coro=self._submit_user_coro,
            set_context=lambda: _stream_source_context.set(context),
            reset_context=_stream_source_context.reset,
            execution_scope_factory=lambda: _user_execution_scope(None),
            emit_item=lambda item: _emit_item(context, item),
            on_stopped=_on_stopped,
            on_error=_on_error,
            on_cancelled=_on_cancelled,
            thread_name=f"source-{source_id}",
        )
        if started is None:
            return None
        return StreamSourceRunner(
            source_id=source_id,
            stop_event=stop_event,
            thread=started.thread,
            future=started.future,
            close_fn=started.close_fn,
        )

    def submit(self, coro: Any, *, loop: str = "runtime"):
        if loop == "user" and self._submit_user_coro is not None:
            return self._submit_user_coro(coro)
        return self._submit_coro(coro)

    def _update_status(self, source_id: str, status: str, error: Exception | None) -> None:
        with self._lock:
            handle = self._sources.get(source_id)
            if handle is None:
                return
            handle.status = status
            handle.last_error = error

    def stop_source(self, source_id: str) -> bool:
        with self._lock:
            handle = self._sources.get(source_id)
        if handle is None:
            return False
        if handle.runner is not None:
            handle.runner.stop()
        stop_method = getattr(handle.source, "stop", None)
        if callable(stop_method):
            stop_result = stop_method()
            if inspect.isawaitable(stop_result):
                self._submit_user_coro(stop_result)
        with self._lock:
            self._sources.pop(source_id, None)
        return True

    def list_sources(self) -> list[dict[str, Any]]:
        with self._lock:
            handles = list(self._sources.values())
        return [
            {
                "source_id": handle.source_id,
                "status": handle.status,
                "restart_policy": handle.restart_policy,
                "last_error": str(handle.last_error) if handle.last_error else None,
            }
            for handle in handles
        ]

    def get_source(self, source_id: str) -> StreamSourceHandle | None:
        with self._lock:
            return self._sources.get(source_id)

    def _build_request_id(self, source_id: str) -> str:
        if callable(self._request_id_factory):
            try:
                request_id = self._request_id_factory(source_id)
            except TypeError:
                request_id = self._request_id_factory()  # type: ignore[misc]
            if isinstance(request_id, str) and request_id:
                return request_id
        return uuid.uuid4().hex


__all__ = [
    "StreamSourceContext",
    "StreamSourceHandle",
    "StreamSourceRegistry",
    "StreamSourceRunner",
    "stream_source_context",
    "emit",
]
