import logging
import sys
import tqdm
from .plugin_base import PluginBase


class TqdmLoggingHandler(logging.StreamHandler):
    def emit(self, record):
        msg = self.format(record)
        tqdm.tqdm.write(msg, file=self.stream)
        self.flush()


def _is_console_handler(handler):
    return isinstance(handler, logging.StreamHandler) and handler.stream in (sys.stdout, sys.stderr)


class Plugin(PluginBase):
    def on_train_start(self, trainer, model):
        self._training_progressbar = tqdm.tqdm(total=trainer.num_epochs, position=0)
        # Replace the logging stream handlers
        self._original_handlers = logging.root.handlers
        stream_handlers = [h for h in logging.root.handlers if _is_console_handler(h)]
        if stream_handlers:
            tqdm_handler = TqdmLoggingHandler()
            tqdm_handler.setFormatter(stream_handlers[0].formatter)
            tqdm_handler.stream = stream_handlers[0].stream
            logging.root.handlers = [h for h in logging.root.handlers if not _is_console_handler(h)] + [tqdm_handler]

    def on_train_end(self, trainer, model):
        self._training_progressbar.close()
        logging.root.handlers = self._original_handlers

    def on_train_epoch_end(self, trainer, model, epoch_index):
        self._epoch_progressbar.close()
        self._training_progressbar.update(1)

    def on_train_epoch_start(self, trainer, model, epoch_index):
        self._epoch_progressbar = tqdm.tqdm(total=len(trainer.train_dataloader), position=1, leave=None)

    def on_train_batch_end(self, trainer, model, loss, batch, batch_index):
        self._epoch_progressbar.update(1)
