from __future__ import annotations

import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from torch.utils.tensorboard import SummaryWriter


class TensorBoardHandler(logging.Handler):
    """Convert each log record into a TensorBoard text summary."""

    def __init__(self, writer: SummaryWriter, tag: str = "logs"):
        super().__init__(level=logging.INFO)
        self.writer = writer
        self.tag = tag
        self._step = 0  # will be incremented every emit

    def emit(self, record: logging.LogRecord) -> None:
        self.writer.add_text(self.tag, self.format(record), global_step=self._step)
        self._step += 1