"""
kiri/utils/progress.py

Minimal, dependency-free progress bar for training loops.
No rich, no tqdm — just stdlib + ANSI codes.
"""

import sys
import time


class ProgressBar:
    """
    Epoch-level progress bar.

    Epoch  3/10  ━━━━━━━━━━━━━━━━░░░░░░░░  62%  47/75 batches  loss: 0.3421  ETA: 8s
    """

    BAR_WIDTH = 24
    FILL      = "━"
    EMPTY     = "░"

    def __init__(self, total_epochs, total_batches, metrics=None):
        self.total_epochs  = total_epochs
        self.total_batches = total_batches
        self.metrics       = metrics or []
        self._start_time   = None
        self._epoch_start  = None

    def start_epoch(self, epoch):
        self.epoch       = epoch
        self._epoch_start = time.time()
        self._batch      = 0
        self._metric_vals = {}

    def update(self, batch, **metric_vals):
        self._batch      = batch + 1
        self._metric_vals = metric_vals
        self._render()

    def end_epoch(self, **metric_vals):
        self._batch = self.total_batches
        self._metric_vals = metric_vals
        self._render(end=True)
        print()  # newline after bar

    def _render(self, end=False):
        pct      = self._batch / max(self.total_batches, 1)
        filled   = int(self.BAR_WIDTH * pct)
        bar      = self.FILL * filled + self.EMPTY * (self.BAR_WIDTH - filled)

        elapsed  = time.time() - self._epoch_start
        if pct > 0 and not end:
            eta = elapsed / pct * (1 - pct)
            eta_str = f"ETA: {eta:.0f}s"
        else:
            eta_str = f"{elapsed:.1f}s"

        metric_str = "  ".join(
            f"{k}: {v:.4f}" for k, v in self._metric_vals.items()
        )

        line = (
            f"\r  Epoch {self.epoch:>{len(str(self.total_epochs))}}"
            f"/{self.total_epochs}  "
            f"{bar}  "
            f"{int(pct*100):>3}%  "
            f"{self._batch}/{self.total_batches} batches  "
            f"{metric_str}  {eta_str}  "
        )

        sys.stdout.write(line)
        sys.stdout.flush()


class TrainingLogger:
    """
    Manages per-epoch summary printing.
    Used inside Model.fit() when verbose=True.
    """

    def __init__(self, total_epochs):
        self.total_epochs = total_epochs
        self._history     = {"loss": []}
        print(f"\n  Training for {total_epochs} epoch(s)\n")

    def log(self, epoch, **metrics):
        for k, v in metrics.items():
            self._history.setdefault(k, []).append(v)

        parts = [f"Epoch {epoch:>{len(str(self.total_epochs))}}/{self.total_epochs}"]
        for k, v in metrics.items():
            parts.append(f"{k}: {v:.4f}")
        print("  " + "  │  ".join(parts))

    def summary(self):
        print()
        best_epoch = int(np.argmin(self._history["loss"])) + 1
        print(f"  Best loss at epoch {best_epoch}: {min(self._history['loss']):.4f}")
