"""Rich progress bar callback with epoch-level train/val losses."""

from __future__ import annotations

from typing import Any

import torch
from lightning.pytorch.callbacks import RichProgressBar


class LossLoggingCallback(RichProgressBar):
    """Show epoch-level mean train/val losses directly in Rich progress bar."""

    def __init__(self) -> None:
        """Initialize callback state used to aggregate losses per epoch."""
        super().__init__()
        self._train_loss_sum = 0.0
        self._train_loss_count = 0
        self._val_loss_sum = 0.0
        self._val_loss_count = 0
        self._latest_train_loss: float | None = None
        self._latest_val_loss: float | None = None

    @staticmethod
    def _as_float(value: Any) -> float | None:
        """Convert tensor/scalar values to float for logging."""
        if value is None:
            return None
        if torch.is_tensor(value):
            return float(value.detach().cpu().item())
        if isinstance(value, (int, float)):
            return float(value)
        return None

    @staticmethod
    def _extract_loss_from_outputs(outputs: Any) -> float | None:
        """Extract loss scalar from Lightning step outputs."""
        if outputs is None:
            return None
        if isinstance(outputs, dict):
            for key in ("loss", "train_loss", "val_loss"):
                loss = LossLoggingCallback._as_float(outputs.get(key))
                if loss is not None:
                    return loss
        loss = LossLoggingCallback._as_float(outputs)
        if loss is not None:
            return loss
        if hasattr(outputs, "loss"):
            return LossLoggingCallback._as_float(outputs.loss)
        return None

    @staticmethod
    def _extract_train_loss_fallback(metrics: dict[str, Any]) -> float | None:
        """Fallback: get train loss from callback metrics."""
        for key in ("train_loss_epoch", "train_loss", "train/loss"):
            loss = LossLoggingCallback._as_float(metrics.get(key))
            if loss is not None:
                return loss
        return None

    @staticmethod
    def _extract_val_loss_fallback(metrics: dict[str, Any]) -> float | None:
        """Fallback: get validation loss from callback metrics."""
        for key in ("val_loss_epoch", "val_loss", "val/loss"):
            loss = LossLoggingCallback._as_float(metrics.get(key))
            if loss is not None:
                return loss
        return None

    def on_train_epoch_start(self, trainer, pl_module) -> None:
        """Reset train accumulators at epoch start."""
        super().on_train_epoch_start(trainer, pl_module)
        self._train_loss_sum = 0.0
        self._train_loss_count = 0

    def on_validation_epoch_start(self, trainer, pl_module) -> None:
        """Reset validation accumulators at epoch start."""
        super().on_validation_epoch_start(trainer, pl_module)
        self._val_loss_sum = 0.0
        self._val_loss_count = 0

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx) -> None:
        """Accumulate per-batch train loss."""
        super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)
        loss = self._extract_loss_from_outputs(outputs)
        if loss is None:
            return
        self._train_loss_sum += loss
        self._train_loss_count += 1

    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0) -> None:
        """Accumulate per-batch validation loss."""
        super().on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)
        if trainer.sanity_checking:
            return
        loss = self._extract_loss_from_outputs(outputs)
        if loss is None:
            return
        self._val_loss_sum += loss
        self._val_loss_count += 1

    def on_train_epoch_end(self, trainer, pl_module) -> None:
        """Store mean train loss at train epoch end."""
        super().on_train_epoch_end(trainer, pl_module)
        train_loss = (
            self._train_loss_sum / self._train_loss_count
            if self._train_loss_count > 0
            else None
        )
        if train_loss is None:
            metrics = dict(trainer.callback_metrics)
            train_loss = self._extract_train_loss_fallback(metrics)
        self._latest_train_loss = train_loss

    def on_validation_epoch_end(self, trainer, pl_module) -> None:
        """Store mean val loss at validation epoch end."""
        super().on_validation_epoch_end(trainer, pl_module)
        if trainer.sanity_checking:
            return

        val_loss = (
            self._val_loss_sum / self._val_loss_count
            if self._val_loss_count > 0
            else None
        )
        metrics = dict(trainer.callback_metrics)
        if val_loss is None:
            val_loss = self._extract_val_loss_fallback(metrics)
        self._latest_val_loss = val_loss

    def get_metrics(self, trainer, pl_module) -> dict[str, float]:
        """Append epoch-level losses to Rich progress bar metrics."""
        metrics = super().get_metrics(trainer, pl_module)
        if self._latest_train_loss is not None:
            metrics["train_loss"] = round(float(self._latest_train_loss), 6)
        if self._latest_val_loss is not None:
            metrics["val_loss"] = round(float(self._latest_val_loss), 6)
        return metrics