from concurrent.futures import ThreadPoolExecutor

from lightning.pytorch import callbacks as pl_callbacks
from typing_extensions import override

from src.utils import logger


class ModelCheckpointParallel(pl_callbacks.ModelCheckpoint):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.threads = []
        self.thread_pool = ThreadPoolExecutor(1, thread_name_prefix="ModelCheckpointParallel")

    @override
    def on_train_batch_end(self, *args, **kwargs):
        trainer = args[0]
        if self._should_skip_saving_checkpoint(trainer):
            return
        self.threads.append(self.thread_pool.submit(super().on_train_batch_end, *args, **kwargs))

    @override
    def on_train_epoch_end(self, *args, **kwargs):
        trainer = args[0]
        if not self._should_skip_saving_checkpoint(trainer) and self._should_save_on_train_epoch_end(trainer):
            self.threads.append(self.thread_pool.submit(super().on_train_epoch_end, *args, **kwargs))

    @override
    def on_validation_end(self, *args, **kwargs):
        trainer = args[0]
        if not self._should_skip_saving_checkpoint(trainer) and not self._should_save_on_train_epoch_end(trainer):
            self.threads.append(self.thread_pool.submit(super().on_validation_end, *args, **kwargs))

    def wait(self):
        for thread in self.threads:
            try:
                thread.result()
            except Exception as e:
                logger.print_error(f"Exception during checkpoint saving in thread: {e}")
        self.thread_pool.shutdown(wait=True)
        self.thread_pool = ThreadPoolExecutor(1, thread_name_prefix="ModelCheckpointParallel")
        self.threads = []

    @override
    def on_train_end(self, *args, **kwargs):
        self.wait()

    @override
    def on_test_start(self, *args, **kwargs):
        self.wait()