from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import cast

from celery import Celery, signals

from core.config.settings import Settings
from core.logging.banner import log_service_banner
from core.logging.config import configure_logging
from core.logging.context import bind_context, clear_context


@dataclass(frozen=True)
class CelerySignalHandlers:
    on_setup_logging: Callable[..., None]
    on_after_setup_task_logger: Callable[..., None]
    on_task_prerun: Callable[..., None]
    on_task_postrun: Callable[..., None]


def build_celery_signal_handlers(
    settings: Settings | None = None,
) -> CelerySignalHandlers:
    active_settings = settings or Settings()

    def on_setup_logging(*_args: object, **_kwargs: object) -> None:
        configure_logging(settings)
        log_service_banner(
            service_name=active_settings.runtime.service_name,
            environment=active_settings.runtime.environment,
        )

    def on_after_setup_task_logger(*_args: object, **_kwargs: object) -> None:
        configure_logging(settings)

    def on_task_prerun(*_args: object, **kwargs: object) -> None:
        task_id = cast(str | None, kwargs.get("task_id"))
        task = kwargs.get("task")
        task_name = getattr(task, "name", None)
        bind_context(task_id=task_id, task_name=task_name)

    def on_task_postrun(*_args: object, **_kwargs: object) -> None:
        clear_context()

    return CelerySignalHandlers(
        on_setup_logging=on_setup_logging,
        on_after_setup_task_logger=on_after_setup_task_logger,
        on_task_prerun=on_task_prerun,
        on_task_postrun=on_task_postrun,
    )


def configure_celery_app(app: Celery, settings: Settings | None = None) -> None:
    app.conf.worker_hijack_root_logger = False

    handlers = build_celery_signal_handlers(settings)
    signals.setup_logging.connect(handlers.on_setup_logging, weak=False)
    signals.after_setup_task_logger.connect(
        handlers.on_after_setup_task_logger, weak=False
    )
    signals.task_prerun.connect(handlers.on_task_prerun, weak=False)
    signals.task_postrun.connect(handlers.on_task_postrun, weak=False)
