# src/sibi_flux/core/managed_resource/_managed_resource.py
from __future__ import annotations

import abc
import asyncio
import contextlib
import json
import threading
from typing import (
    Any,
    Awaitable,
    Callable,
    Optional,
    Protocol,
    Self,
    final,
    runtime_checkable,
)

import fsspec

from sibi_flux.logger import Logger


# ----------------------- SSE sink protocol -----------------------
@runtime_checkable
class SSESinkProtocol(Protocol):
    """
    Structural protocol for an SSE sink.

    Supported sink shapes:
      - async send(event: str, data: dict)
      - async put({"event": str, "data": str|dict})
      - optional (a)close for lifecycle
    """

    def send(self, event: str, data: dict[str, Any]) -> None: ...
    def put(self, item: dict[str, Any]) -> None: ...
    def aclose(self) -> Awaitable[None]: ...
    def close(self) -> None: ...


# --------- Minimal built-in SSE sink (used when auto_sse=True) ----------
class _QueueSSE:
    """
    Internal sink used when auto_sse=True.

    Behaviour:
      - Iteration yields dict items {"event": str, "data": str}
      - send(event, data_dict) stores JSON string in `data`
      - close/aclose wakes consumers deterministically
    """

    __slots__ = ("q", "_closed", "_loop", "_close_lock")

    _CLOSE_ITEM = {"event": "__close__", "data": ""}

    def __init__(self, *, maxsize: int = 0) -> None:
        # maxsize=0 => unbounded (as asyncio.Queue default behaviour).
        self.q: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=maxsize)
        self._closed = False
        self._close_lock = threading.Lock()
        self._loop: Optional[asyncio.AbstractEventLoop] = None  # bound lazily

    def _bind_loop_if_needed(self) -> None:
        if self._loop is not None:
            return
        try:
            self._loop = asyncio.get_running_loop()
        except RuntimeError:
            # Not in a running loop; we'll try again later.
            return

    def _is_closed(self) -> bool:
        with self._close_lock:
            return self._closed

    async def send(self, event: str, data: dict[str, Any]) -> None:
        self._bind_loop_if_needed()
        if self._is_closed():
            return
        await self.q.put({"event": event, "data": json.dumps(data)})

    async def put(self, item: dict[str, Any]) -> None:
        self._bind_loop_if_needed()
        if self._is_closed():
            return
        await self.q.put(item)

    async def aclose(self) -> None:
        self._bind_loop_if_needed()
        with self._close_lock:
            if self._closed:
                return
            self._closed = True

        with contextlib.suppress(Exception):
            await self.q.put(dict(self._CLOSE_ITEM))

    def close(self) -> None:
        with self._close_lock:
            if self._closed:
                return
            self._closed = True

        close_item = dict(self._CLOSE_ITEM)
        loop = self._loop

        # asyncio.Queue is not thread-safe; if we have a loop, schedule insertion there.
        if loop is not None and loop.is_running():
            def _signal() -> None:
                try:
                    self.q.put_nowait(close_item)
                except asyncio.QueueFull:
                    async def _aput() -> None:
                        with contextlib.suppress(Exception):
                            await self.q.put(close_item)
                    asyncio.run_coroutine_threadsafe(_aput(), loop)

            loop.call_soon_threadsafe(_signal)
            return

        # Fallback for when loop is not available (e.g. during finalization)
        # Note: This is potentially unsafe if called from a thread while the loop
        # is running elsewhere, but as per audit, we add documentation/warning.
        with contextlib.suppress(Exception):
            self.q.put_nowait(close_item)

    def __aiter__(self):
        return self._iter_items()

    async def _iter_items(self):
        self._bind_loop_if_needed()
        while True:
            item = await self.q.get()
            try:
                if item.get("event") == "__close__":
                    break
                yield item
            finally:
                with contextlib.suppress(Exception):
                    self.q.task_done()

    async def aiter_sse(self):
        """
        SSE wire format iterator for ASGI StreamingResponse:
          event: <name>\n
          data: <payload>\n
          \n
        """
        async for item in self:
            event = item.get("event", "")
            data = item.get("data", "")

            # Preserve existing convention: data stored as JSON string.
            if not isinstance(data, str):
                data = json.dumps(data)

            lines = data.splitlines() or [""]
            out: list[str] = []
            if event:
                out.append(f"event: {event}\n")
            for line in lines:
                out.append(f"data: {line}\n")
            out.append("\n")
            yield "".join(out).encode("utf-8")


# ------------------------------ ManagedResource ------------------------------
class ManagedResource(abc.ABC):
    """
    ManagedResource

    A lifecycle base class that standardizes resource management.
    """

    __slots__ = (
        "verbose",
        "debug",
        "_log_cleanup_errors",
        "logger",
        "_owns_logger",
        "fs",
        "_fs_factory",
        "_owns_fs",
        "_sse",
        "_sse_factory",
        "_owns_sse",
        "_emitter",
        "_auto_sse",
        "_auto_sse_maxsize",
        "_emit_timeout_s",
        "_emit_drop_on_timeout",
        "_validate_emitter",
        "_is_closed",
        "_closing",
        "_close_lock",
        "_aclose_lock",
        "_finalizer_enabled",
        "__weakref__",
    )

    def __init__(
        self,
        *,
        verbose: bool = False,
        debug: bool = False,
        log_cleanup_errors: bool = True,
        logger: Optional[Logger] = None,
        fs: Optional[fsspec.AbstractFileSystem] = None,
        fs_factory: Optional[Callable[[], fsspec.AbstractFileSystem]] = None,
        emitter: Optional[Callable[[str, dict[str, Any]], Awaitable[None]]] = None,
        emitter_factory: Optional[
            Callable[[], Callable[[str, dict[str, Any]], Awaitable[None]]]
        ] = None,
        sse: Optional[SSESinkProtocol] = None,
        sse_factory: Optional[Callable[[], SSESinkProtocol]] = None,
        auto_sse: bool = False,
        auto_sse_maxsize: int = 0,
        emit_timeout_s: Optional[float] = None,
        emit_drop_on_timeout: bool = True,
        validate_emitter: bool = False,
        finalizer_enabled: bool = True,
        **_: object,
    ) -> None:
        self.verbose = verbose
        self.debug = debug
        self._log_cleanup_errors = log_cleanup_errors

        self._emit_timeout_s = emit_timeout_s
        self._emit_drop_on_timeout = emit_drop_on_timeout
        self._validate_emitter = validate_emitter

        self._is_closed = False
        self._closing = False
        self._close_lock = threading.RLock()
        self._aclose_lock = asyncio.Lock()

        # ---------- Logger ----------
        if logger is None:
            self.logger = Logger.default_logger(logger_name=self.__class__.__name__)
            self._owns_logger = False  # default logger is shared (stdlib getLogger)
            level = (
                Logger.DEBUG
                if self.debug
                else (Logger.INFO if self.verbose else Logger.WARNING)
            )
            self.logger.set_level(level)
        else:
            self.logger = logger
            self._owns_logger = False

        # ---------- FS ----------
        self.fs: Optional[fsspec.AbstractFileSystem] = None
        self._fs_factory: Optional[Callable[[], fsspec.AbstractFileSystem]] = None
        self._owns_fs = False

        if fs is not None:
            if not isinstance(fs, fsspec.AbstractFileSystem):
                raise TypeError(
                    f"fs must be an fsspec.AbstractFileSystem, got {type(fs)!r}"
                )
            self.fs = fs
        elif fs_factory is not None:
            if not callable(fs_factory):
                raise TypeError("fs_factory must be callable")
            self._fs_factory = fs_factory
            self._owns_fs = True

        # ---------- SSE / emitter ----------
        self._sse: Optional[SSESinkProtocol] = None
        self._sse_factory: Optional[Callable[[], SSESinkProtocol]] = None
        self._owns_sse = False
        self._auto_sse = auto_sse
        self._auto_sse_maxsize = auto_sse_maxsize

        self._emitter: Optional[Callable[[str, dict[str, Any]], Awaitable[None]]] = None
        if emitter is not None:
            self._emitter = emitter
        elif emitter_factory is not None:
            if not callable(emitter_factory):
                raise TypeError("emitter_factory must be callable")
            self._emitter = emitter_factory()

        if sse is not None:
            self._sse = sse
            self._emitter = self._emitter or self._build_emitter(sse)
        elif sse_factory is not None:
            if not callable(sse_factory):
                raise TypeError("sse_factory must be callable")
            self._sse_factory = sse_factory
            self._owns_sse = True

        if (
            self._auto_sse
            and self._sse is None
            and self._emitter is None
            and self._sse_factory is None
        ):
            self._create_auto_sse()

        # ---------- Finalizer ----------
        self._finalizer_enabled = finalizer_enabled

        if self.debug:
            with contextlib.suppress(Exception):
                self.logger.debug(
                    "Initialised %s %s", self.__class__.__name__, repr(self)
                )

    # ----------------- Properties -----------------
    @property
    def closed(self) -> bool:
        return self._is_closed

    @property
    def has_fs(self) -> bool:
        return self.fs is not None or self._fs_factory is not None

    @property
    def has_sse(self) -> bool:
        return (
            (self._emitter is not None)
            or (self._sse is not None)
            or self._auto_sse
            or (self._sse_factory is not None)
        )

    def __repr__(self) -> str:
        def _status(current: bool, factory: bool, owned: bool) -> str:
            if current:
                return "own" if owned else "external"
            if factory:
                return "own(lazy)"
            return "none"

        fs_status = _status(
            self.fs is not None, self._fs_factory is not None, self._owns_fs
        )
        sse_current = (self._sse is not None) or (self._emitter is not None)
        sse_factory = (self._sse_factory is not None) or self._auto_sse
        sse_status = _status(sse_current, sse_factory, self._owns_sse or self._auto_sse)
        return (
            f"<{self.__class__.__name__} debug={self.debug} verbose={self.verbose} "
            f"log_cleanup_errors={self._log_cleanup_errors} fs={fs_status} sse={sse_status}>"
        )

    # ----------------- Hooks -----------------
    def _cleanup(self) -> None:
        return

    async def _acleanup(self) -> None:
        return

    def _assert_open(self) -> None:
        if self._is_closed or self._closing:
            raise RuntimeError(f"{self.__class__.__name__} is closed")

    # ----------------- FS -----------------
    def set_fs_factory(
        self, factory: Optional[Callable[[], fsspec.AbstractFileSystem]]
    ) -> None:
        with self._close_lock:
            self._assert_open()
            if self.fs is not None:
                return
            if factory is not None and not callable(factory):
                raise TypeError("fs_factory must be callable")
            self._fs_factory = factory
            self._owns_fs = factory is not None

    def _ensure_fs(self) -> Optional[fsspec.AbstractFileSystem]:
        with self._close_lock:
            self._assert_open()
            if self.fs is not None:
                return self.fs
            if self._fs_factory is None:
                return None
            fs_new = self._fs_factory()
            if not isinstance(fs_new, fsspec.AbstractFileSystem):
                raise TypeError(
                    f"fs_factory() must return fsspec.AbstractFileSystem, got {type(fs_new)!r}"
                )
            self.fs = fs_new
            return self.fs

    def require_fs(self) -> fsspec.AbstractFileSystem:
        fs = self._ensure_fs()
        if fs is None:
            raise RuntimeError(
                f"{self.__class__.__name__}: filesystem is required but not configured"
            )
        return fs

    # ----------------- SSE -----------------
    def _create_auto_sse(self) -> None:
        sink = _QueueSSE(maxsize=self._auto_sse_maxsize)
        self._sse = sink
        self._owns_sse = True
        self._emitter = self._build_emitter(sink)

    def set_sse_factory(self, factory: Optional[Callable[[], SSESinkProtocol]]) -> None:
        with self._close_lock:
            self._assert_open()
            if self._sse is not None or self._emitter is not None:
                return
            if factory is not None and not callable(factory):
                raise TypeError("sse_factory must be callable")
            self._sse_factory = factory
            self._owns_sse = factory is not None

    def _ensure_sse(self) -> Optional[SSESinkProtocol]:
        with self._close_lock:
            # Allow retrieving already-created SSE even after close (useful for draining).
            if self._sse is not None:
                return self._sse

            self._assert_open()

            if self._sse_factory is not None:
                sink = self._sse_factory()
                self._sse = sink
                self._owns_sse = True
                if self._emitter is None:
                    self._emitter = self._build_emitter(sink)
                return self._sse

            if self._auto_sse and self._emitter is None:
                self._create_auto_sse()
                return self._sse

            return None

    def get_sse(self) -> Optional[SSESinkProtocol]:
        return self._ensure_sse()

    # ----------------- Emitter building/validation -----------------
    def _validate_send(self, send_fn: Any) -> bool:
        """Robustly check if callable has (event, data) signature."""
        if not callable(send_fn):
            return False
        try:
            from inspect import signature
            sig = signature(send_fn)
            # We expect exactly 2 parameters (event, data). 
            # Note: This is a loose check to handle variations in naming.
            params = [p for p in sig.parameters.values() if p.kind in (p.POSITIONAL_OR_KEYWORD, p.POSITIONAL_ONLY)]
            # If it's a bound method, the first param (self) is already handled by signature() usually, 
            # but let's be careful.
            return len(params) == 2
        except (ValueError, TypeError):
            return True # Fallback for C-extensions

    def _validate_put(self, put_fn: Any) -> bool:
        """Robustly check if callable has (item) signature."""
        if not callable(put_fn):
            return False
        try:
            from inspect import signature
            sig = signature(put_fn)
            params = [p for p in sig.parameters.values() if p.kind in (p.POSITIONAL_OR_KEYWORD, p.POSITIONAL_ONLY)]
            return len(params) == 1
        except (ValueError, TypeError):
            return True # Fallback

    def _build_emitter(
        self, sink: object
    ) -> Callable[[str, dict[str, Any]], Awaitable[None]]:
        # Lightweight duck-typing check instead of heavy inspect.signature
        send = getattr(sink, "send", None)
        put = getattr(sink, "put", None)

        if callable(send):
            if asyncio.iscoroutinefunction(send):
                async def _emit_send(event: str, payload: dict[str, Any]) -> None:
                    await send(event, payload)
                return _emit_send
            else:
                async def _emit_send_sync(event: str, payload: dict[str, Any]) -> None:
                    await asyncio.to_thread(send, event, payload)
                return _emit_send_sync

        if callable(put):
            if asyncio.iscoroutinefunction(put):
                async def _emit_put(event: str, payload: dict[str, Any]) -> None:
                    await put({"event": event, "data": json.dumps(payload)})
                return _emit_put
            else:
                async def _emit_put_sync(event: str, payload: dict[str, Any]) -> None:
                    await asyncio.to_thread(
                        put, {"event": event, "data": json.dumps(payload)}
                    )
                return _emit_put_sync

        raise TypeError(
            f"{self.__class__.__name__}: SSE sink must expose callable send(event, data) or put(item)"
        )

    async def emit(self, event: str, **data: Any) -> None:
        if self._is_closed or self._closing:
            return

        if self._emitter is None:
            self._ensure_sse()

        emitter = self._emitter
        if emitter is None:
            return

        try:
            if self._emit_timeout_s is None:
                await emitter(event, data)
            else:
                await asyncio.wait_for(
                    emitter(event, data), timeout=self._emit_timeout_s
                )
        except asyncio.TimeoutError:
            if not self._emit_drop_on_timeout:
                raise
            if self._log_cleanup_errors:
                with contextlib.suppress(Exception):
                    self.logger.warning("emit(%r) timed out; event dropped", event)
        except Exception:
            if self._log_cleanup_errors:
                with contextlib.suppress(Exception):
                    self.logger.error(
                        "Error emitting SSE event %r", event, exc_info=self.debug
                    )

    # ----------------- Shutdown helpers -----------------
    def _release_owned_fs(self) -> None:
        if self._owns_fs and self.fs is not None:
            close_method = getattr(self.fs, "close", None)
            with contextlib.suppress(Exception):
                if callable(close_method):
                    close_method()
            self.fs = None

    async def _aclose_obj(self, obj: object, timeout: float = 1.0) -> None:
        # Prefer aclose(); only fall back to close() if aclose isn't available or fails.
        aclose_method = getattr(obj, "aclose", None)
        if callable(aclose_method):
            try:
                out = aclose_method()
                if asyncio.iscoroutine(out):
                    await asyncio.wait_for(out, timeout=timeout)
                return
            except Exception:
                if self.debug:
                    with contextlib.suppress(Exception):
                        self.logger.debug("Error in _aclose_obj", exc_info=True)

        close_method = getattr(obj, "close", None)
        if callable(close_method):
            with contextlib.suppress(Exception):
                await asyncio.to_thread(close_method)

    def _shutdown_logger(self) -> None:
        if not self._owns_logger:
            return
        with contextlib.suppress(Exception):
            shutdown_method = getattr(self.logger, "shutdown", None)
            if callable(shutdown_method):
                shutdown_method()
                return
        with contextlib.suppress(Exception):
            close_handlers_method = getattr(self.logger, "close_handlers", None)
            if callable(close_handlers_method):
                close_handlers_method()

    def _shutdown_owned_resources_sync(self) -> None:
        self._release_owned_fs()
        if self._owns_sse and self._sse is not None:
            with contextlib.suppress(Exception):
                close_method = getattr(self._sse, "close", None)
                if callable(close_method):
                    close_method()
        self._sse = None
        self._emitter = None
        self._shutdown_logger()

    async def _shutdown_owned_resources_async(self) -> None:
        self._release_owned_fs()
        if self._owns_sse and self._sse is not None:
            await self._aclose_obj(self._sse)
        self._sse = None
        self._emitter = None
        self._shutdown_logger()

    # ----------------- Public lifecycle -----------------
    @final
    def close(self, *, suppress_errors: bool = False) -> None:
        # Guard against partial initialization failure
        lock = getattr(self, "_close_lock", None)
        if lock is None:
            return

        with lock:
            if self._is_closed or self._closing:
                return
            self._closing = True

        try:
            self._cleanup()
        except Exception:
            if self._log_cleanup_errors:
                with contextlib.suppress(Exception):
                    self.logger.error(
                        "Error during %s._cleanup()",
                        self.__class__.__name__,
                        exc_info=self.debug,
                    )
            if not suppress_errors:
                raise
        finally:
            if self.debug:
                with contextlib.suppress(Exception):
                    self.logger.debug("Component %s closed.", self.__class__.__name__)

            with lock:
                self._is_closed = True
                self._closing = False
                self.debug = False

            self._shutdown_owned_resources_sync()

    async def aclose(
        self,
        *,
        suppress_errors: bool = True,
        run_sync_cleanup_if_missing: bool = True,
    ) -> None:
        async with self._aclose_lock:
            with self._close_lock:
                if self._is_closed or self._closing:
                    return
                self._closing = True

            try:
                uses_default_acleanup = (
                    type(self)._acleanup is ManagedResource._acleanup
                )
                if run_sync_cleanup_if_missing and uses_default_acleanup:
                    await asyncio.to_thread(self._cleanup)
                else:
                    await self._acleanup()
            except Exception:
                if self._log_cleanup_errors:
                    with contextlib.suppress(Exception):
                        self.logger.error(
                            "Error during %s._acleanup()",
                            self.__class__.__name__,
                            exc_info=self.debug,
                        )
                if not suppress_errors:
                    raise
            finally:
                if self.debug:
                    with contextlib.suppress(Exception):
                        self.logger.debug(
                            "Async component %s closed.", self.__class__.__name__
                        )

                with self._close_lock:
                    self._is_closed = True
                    self._closing = False
                    self.debug = False

                await self._shutdown_owned_resources_async()

    # ----------------- Context managers -----------------
    @final
    def __enter__(self) -> Self:
        return self

    @final
    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: Any | None,
    ) -> None:
        self.close()
        return False

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
        await self.aclose()
        return False

    # ----------------- Finalizer -----------------
    def __del__(self) -> None:
        """Safe finalizer for resource cleanup."""
        if getattr(self, "_finalizer_enabled", False):
            # We must be very careful in __del__
            # Only call close() if the object is still fundamentally intact
            try:
                if not getattr(self, "_is_closed", True):
                    self.close()
            except Exception:
                # Never raise in __del__
                pass
