"""Celery Bootstep that enables Sluice event capture and forwarding.

A Bootstep is Celery's extension mechanism — it hooks into the worker
lifecycle (start, stop, etc.) without monkey-patching. This is the same
pattern Sentry uses via ``CeleryIntegration``.

The bootstep:
1. Runs auto-configuration on worker startup (sets the 3 critical flags).
2. Starts the event buffer + HTTP transport pipeline.
3. Connects Celery signal handlers that normalize and buffer events.
4. Flushes remaining events on worker shutdown.
"""

from __future__ import annotations

import logging
import os
import threading
from typing import TYPE_CHECKING, Any

from celery import bootsteps

from sluice._config import auto_configure
from sluice._transport import _DEFAULT_BASE_URL

if TYPE_CHECKING:
    from celery.worker import WorkController

logger = logging.getLogger("sluice")

# Module-level storage. Celery instantiates bootstep classes without custom
# kwargs, so these must be stored here by install_bootstep().
_api_key: str = ""
_connection_id: str = ""
_config_lock = threading.Lock()

_BASE_URL_ENV = "SLUICE_BASE_URL"


class SluiceBootstep(bootsteps.StartStopStep):  # type: ignore[misc]
    """Sluice monitoring bootstep for Celery workers.

    This step runs during worker startup and:
    - Auto-configures the three critical monitoring flags
    - Starts the event forwarding pipeline (buffer + transport)
    - Connects Celery signal handlers for all task and worker events
    """

    label = "sluice"

    def __init__(
        self,
        parent: Any,
        **kwargs: Any,
    ) -> None:
        super().__init__(parent, **kwargs)
        with _config_lock:
            self._api_key = _api_key
            self._connection_id = _connection_id
        self._buffer: Any = None  # EventBuffer, set in start()
        self._memory_guard: Any = None  # MemoryGuard, set in start()
        self._handlers: list[Any] = []  # Signal handlers, prevent GC

    def start(self, parent: WorkController) -> None:
        """Called when the worker starts consuming tasks."""
        try:
            app = parent.app
            auto_configure(app)

            # Set up the event forwarding pipeline.
            base_url = os.environ.get(_BASE_URL_ENV, _DEFAULT_BASE_URL)

            from sluice._buffer import EventBuffer
            from sluice._safety import MemoryGuard, resolve_memory_threshold
            from sluice._transport import SluiceTransport

            # Memory guard: disable event capture if RSS grows too much.
            # Threshold is configurable via SLUICE_MEMORY_THRESHOLD_MB env.
            threshold = resolve_memory_threshold()
            self._memory_guard = MemoryGuard(threshold_bytes=threshold)
            self._memory_guard.start()

            # Shared cancel event: buffer sets it on stop(), transport
            # checks it to interrupt retry sleeps during shutdown.
            cancel = threading.Event()
            transport = SluiceTransport(
                self._api_key,
                base_url=base_url,
                connection_id=self._connection_id,
                cancel_event=cancel,
            )
            self._buffer = EventBuffer(transport, memory_guard=self._memory_guard)
            self._buffer._stopping = cancel
            self._buffer.start()

            # Connect Celery signal handlers for event capture.
            # Store handlers to prevent garbage collection (weak refs).
            self._handlers = _connect_event_handlers(self._buffer)

            logger.info("[sluice] Bootstep started. Event capture enabled.")
        except Exception:
            # Production-cannot-fail: never crash the worker.
            logger.exception(
                "[sluice] Bootstep startup failed. "
                "Monitoring may be incomplete. "
                "See: https://docs.sluice.sh/troubleshooting/bootstep"
            )

    def stop(self, parent: WorkController) -> None:
        """Called when the worker is shutting down."""
        try:
            if self._memory_guard is not None:
                self._memory_guard.stop()
                self._memory_guard = None
            if self._buffer is not None:
                self._buffer.stop()
                self._buffer = None
            logger.info("[sluice] Bootstep stopped.")
        except Exception:
            # Never crash on shutdown either.
            logger.exception("[sluice] Error during bootstep shutdown.")


def _connect_event_handlers(buffer: Any) -> list[Any]:
    """Connect Celery signal handlers that normalize and buffer events.

    Each handler wraps its body in try/except to ensure no signal
    handler can ever crash the worker. Events that fail to normalize
    are silently skipped.

    Returns the list of handler functions. The caller MUST keep a
    reference to them — Celery signals use weak references by default.

    Note: Worker lifecycle events (online/offline/heartbeat) are NOT captured
    here. The Go agent monitors worker state via Redis pub/sub, which is more
    reliable than in-process signals (survives OOM kills, segfaults, etc.).
    """
    from celery.signals import (
        task_failure,
        task_postrun,
        task_prerun,
        task_received,
        task_rejected,
        task_retry,
        task_revoked,
        task_sent,
    )

    from sluice._events import normalize_task_event

    def _on_task_sent(sender: Any = None, headers: Any = None, **kwargs: Any) -> None:
        try:
            if headers is None:
                return
            event: dict[str, Any] = {
                "type": "task-sent",
                "uuid": headers.get("id", ""),
                "name": headers.get("task", ""),
                "queue": headers.get("queue", kwargs.get("queue", "")),
            }
            normalized = normalize_task_event(event)
            if normalized is not None:
                buffer.push(normalized)
        except Exception:
            logger.debug("[sluice] Error handling task-sent.", exc_info=True)

    def _on_task_received(
        sender: Any = None, request: Any = None, **kwargs: Any
    ) -> None:
        try:
            event: dict[str, Any] = {
                "type": "task-received",
                "uuid": getattr(request, "id", "") if request else "",
                "name": (getattr(request, "name", "") if request else str(sender)),
                "hostname": getattr(request, "hostname", ""),
            }
            normalized = normalize_task_event(event)
            if normalized is not None:
                buffer.push(normalized)
        except Exception:
            logger.debug("[sluice] Error handling task-received.", exc_info=True)

    def _on_task_prerun(
        sender: Any = None, task_id: str | None = None, **kwargs: Any
    ) -> None:
        try:
            event: dict[str, Any] = {
                "type": "task-started",
                "uuid": task_id or "",
                "name": (getattr(sender, "name", str(sender)) if sender else ""),
            }
            normalized = normalize_task_event(event)
            if normalized is not None:
                buffer.push(normalized)
        except Exception:
            logger.debug("[sluice] Error handling task-prerun.", exc_info=True)

    def _on_task_postrun(
        sender: Any = None,
        task_id: str | None = None,
        state: str | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            if state != "SUCCESS":
                return
            event: dict[str, Any] = {
                "type": "task-succeeded",
                "uuid": task_id or "",
                "name": (getattr(sender, "name", str(sender)) if sender else ""),
            }
            runtime = kwargs.get("runtime")
            if runtime is not None:
                event["runtime"] = runtime
            normalized = normalize_task_event(event)
            if normalized is not None:
                buffer.push(normalized)
        except Exception:
            logger.debug("[sluice] Error handling task-postrun.", exc_info=True)

    def _on_task_failure(
        sender: Any = None,
        task_id: str | None = None,
        exception: BaseException | None = None,
        traceback: Any = None,
        **kwargs: Any,
    ) -> None:
        try:
            event: dict[str, Any] = {
                "type": "task-failed",
                "uuid": task_id or "",
                "name": (getattr(sender, "name", str(sender)) if sender else ""),
                "exception": str(exception) if exception else "",
                "traceback": str(traceback) if traceback else "",
            }
            normalized = normalize_task_event(event)
            if normalized is not None:
                buffer.push(normalized)
        except Exception:
            logger.debug("[sluice] Error handling task-failure.", exc_info=True)

    def _on_task_retry(
        sender: Any = None,
        request: Any = None,
        reason: Any = None,
        **kwargs: Any,
    ) -> None:
        try:
            event: dict[str, Any] = {
                "type": "task-retried",
                "uuid": getattr(request, "id", "") if request else "",
                "name": (getattr(sender, "name", str(sender)) if sender else ""),
                "exception": str(reason) if reason else "",
            }
            normalized = normalize_task_event(event)
            if normalized is not None:
                buffer.push(normalized)
        except Exception:
            logger.debug("[sluice] Error handling task-retry.", exc_info=True)

    def _on_task_revoked(
        sender: Any = None,
        request: Any = None,
        terminated: bool = False,
        signum: Any = None,
        expired: bool = False,
        **kwargs: Any,
    ) -> None:
        try:
            event: dict[str, Any] = {
                "type": "task-revoked",
                "uuid": getattr(request, "id", "") if request else "",
                "name": (getattr(sender, "name", str(sender)) if sender else ""),
            }
            normalized = normalize_task_event(event)
            if normalized is not None:
                buffer.push(normalized)
        except Exception:
            logger.debug("[sluice] Error handling task-revoked.", exc_info=True)

    def _on_task_rejected(
        sender: Any = None, message: Any = None, **kwargs: Any
    ) -> None:
        try:
            event: dict[str, Any] = {
                "type": "task-rejected",
                "uuid": (
                    getattr(message, "headers", {}).get("id", "") if message else ""
                ),
                "name": (
                    getattr(message, "headers", {}).get("task", "") if message else ""
                ),
            }
            normalized = normalize_task_event(event)
            if normalized is not None:
                buffer.push(normalized)
        except Exception:
            logger.debug("[sluice] Error handling task-rejected.", exc_info=True)

    handlers = [
        _on_task_sent,
        _on_task_received,
        _on_task_prerun,
        _on_task_postrun,
        _on_task_failure,
        _on_task_retry,
        _on_task_revoked,
        _on_task_rejected,
    ]

    task_sent.connect(_on_task_sent, weak=False)
    task_received.connect(_on_task_received, weak=False)
    task_prerun.connect(_on_task_prerun, weak=False)
    task_postrun.connect(_on_task_postrun, weak=False)
    task_failure.connect(_on_task_failure, weak=False)
    task_retry.connect(_on_task_retry, weak=False)
    task_revoked.connect(_on_task_revoked, weak=False)
    task_rejected.connect(_on_task_rejected, weak=False)

    return handlers


def install_bootstep(app: Any, api_key: str, connection_id: str) -> None:
    """Install the Sluice bootstep on a Celery app.

    This registers the bootstep on the ``consumer`` blueprint.
    It is idempotent — calling it multiple times is safe.
    """
    global _api_key, _connection_id  # noqa: PLW0603
    with _config_lock:
        _api_key = api_key
        _connection_id = connection_id

    # Add bootstep to the consumer blueprint.
    # The consumer blueprint controls the worker's message consumption loop.
    if SluiceBootstep not in app.steps["consumer"]:
        app.steps["consumer"].add(SluiceBootstep)
        logger.info("[sluice] Bootstep registered on Celery app.")
