import logging
from collections import abc
from collections.abc import Callable, Iterable
from typing import Any, TypeAlias, cast, final

from lightning.pytorch import Callback, LightningModule
from lightning.pytorch.callbacks import LambdaCallback
from typing_extensions import override

from ...util.typing_utils import mixin_base_type

log = logging.getLogger(__name__)

CallbackFn: TypeAlias = Callable[[], Callback | Iterable[Callback] | None]


class CallbackRegistrarModuleMixin:
    @override
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self._ll_callbacks: list[CallbackFn] = []

    def register_callback(
        self,
        callback: Callback | Iterable[Callback] | CallbackFn | None = None,
        *,
        setup: Callable | None = None,
        teardown: Callable | None = None,
        on_fit_start: Callable | None = None,
        on_fit_end: Callable | None = None,
        on_sanity_check_start: Callable | None = None,
        on_sanity_check_end: Callable | None = None,
        on_train_batch_start: Callable | None = None,
        on_train_batch_end: Callable | None = None,
        on_train_epoch_start: Callable | None = None,
        on_train_epoch_end: Callable | None = None,
        on_validation_epoch_start: Callable | None = None,
        on_validation_epoch_end: Callable | None = None,
        on_test_epoch_start: Callable | None = None,
        on_test_epoch_end: Callable | None = None,
        on_validation_batch_start: Callable | None = None,
        on_validation_batch_end: Callable | None = None,
        on_test_batch_start: Callable | None = None,
        on_test_batch_end: Callable | None = None,
        on_train_start: Callable | None = None,
        on_train_end: Callable | None = None,
        on_validation_start: Callable | None = None,
        on_validation_end: Callable | None = None,
        on_test_start: Callable | None = None,
        on_test_end: Callable | None = None,
        on_exception: Callable | None = None,
        on_save_checkpoint: Callable | None = None,
        on_load_checkpoint: Callable | None = None,
        on_before_backward: Callable | None = None,
        on_after_backward: Callable | None = None,
        on_before_optimizer_step: Callable | None = None,
        on_before_zero_grad: Callable | None = None,
        on_predict_start: Callable | None = None,
        on_predict_end: Callable | None = None,
        on_predict_batch_start: Callable | None = None,
        on_predict_batch_end: Callable | None = None,
        on_predict_epoch_start: Callable | None = None,
        on_predict_epoch_end: Callable | None = None,
    ):
        if callback is None:
            callback = LambdaCallback(
                setup=setup,
                teardown=teardown,
                on_fit_start=on_fit_start,
                on_fit_end=on_fit_end,
                on_sanity_check_start=on_sanity_check_start,
                on_sanity_check_end=on_sanity_check_end,
                on_train_batch_start=on_train_batch_start,
                on_train_batch_end=on_train_batch_end,
                on_train_epoch_start=on_train_epoch_start,
                on_train_epoch_end=on_train_epoch_end,
                on_validation_epoch_start=on_validation_epoch_start,
                on_validation_epoch_end=on_validation_epoch_end,
                on_test_epoch_start=on_test_epoch_start,
                on_test_epoch_end=on_test_epoch_end,
                on_validation_batch_start=on_validation_batch_start,
                on_validation_batch_end=on_validation_batch_end,
                on_test_batch_start=on_test_batch_start,
                on_test_batch_end=on_test_batch_end,
                on_train_start=on_train_start,
                on_train_end=on_train_end,
                on_validation_start=on_validation_start,
                on_validation_end=on_validation_end,
                on_test_start=on_test_start,
                on_test_end=on_test_end,
                on_exception=on_exception,
                on_save_checkpoint=on_save_checkpoint,
                on_load_checkpoint=on_load_checkpoint,
                on_before_backward=on_before_backward,
                on_after_backward=on_after_backward,
                on_before_optimizer_step=on_before_optimizer_step,
                on_before_zero_grad=on_before_zero_grad,
                on_predict_start=on_predict_start,
                on_predict_end=on_predict_end,
                on_predict_batch_start=on_predict_batch_start,
                on_predict_batch_end=on_predict_batch_end,
                on_predict_epoch_start=on_predict_epoch_start,
                on_predict_epoch_end=on_predict_epoch_end,
            )

        if not callable(callback):
            callback_ = cast(CallbackFn, lambda: callback)
        else:
            callback_ = callback

        self._ll_callbacks.append(callback_)


class CallbackModuleMixin(
    CallbackRegistrarModuleMixin,
    mixin_base_type(LightningModule),
):
    def _gather_all_callbacks(self):
        modules: list[Any] = []
        if isinstance(self, CallbackRegistrarModuleMixin):
            modules.append(self)
        if (
            datamodule := getattr(self.trainer, "datamodule", None)
        ) is not None and isinstance(datamodule, CallbackRegistrarModuleMixin):
            modules.append(datamodule)
        modules.extend(
            module
            for module in self.children()
            if isinstance(module, CallbackRegistrarModuleMixin)
        )
        for module in modules:
            yield from module._ll_callbacks

    @final
    @override
    def configure_callbacks(self):
        callbacks = super().configure_callbacks()
        if not isinstance(callbacks, abc.Sequence):
            callbacks = [callbacks]

        callbacks = list(callbacks)
        for callback_fn in self._gather_all_callbacks():
            callback_result = callback_fn()
            if callback_result is None:
                continue

            if not isinstance(callback_result, abc.Iterable):
                callback_result = [callback_result]

            for callback in callback_result:
                log.info(
                    f"Registering {callback.__class__.__qualname__} callback {callback}"
                )
                callbacks.append(callback)

        return callbacks
