"""
Inference Probe Callback
========================

After each checkpoint is saved, runs BADAS inference on a fixed probe video
and saves a visualization of the collision risk timeline.

Useful for monitoring model quality progression during training without
manually evaluating checkpoints.

Implementation note
-------------------
PyTorch Lightning 2.x places all ModelCheckpoint callbacks **after** every
other callback in the execution order (see callback_connector.py line ~239):

    return tuner_callbacks + other_callbacks + checkpoint_callbacks

This means on_validation_end in a regular Callback fires *before* the
checkpoint is saved. The workaround is to set a _pending flag inside
on_validation_end and do the actual probe work inside on_train_batch_start,
which fires after the checkpoint has been written to disk.
"""

import glob
import os

import pytorch_lightning as pl


class InferenceProbeCallback(pl.Callback):
    """
    Runs inference on a fixed probe video after each checkpoint is saved
    and saves a risk-timeline visualization PNG alongside the checkpoint.

    Args:
        video_path: Absolute path to the probe video.
        output_dir: Directory to save PNG visualizations.
        stride: Frame stride for inference (4 = fast, 1 = every frame).
    """

    def __init__(self, video_path: str, output_dir: str, stride: int = 1):
        self.video_path = video_path
        self.output_dir = output_dir
        self.stride = stride
        self._last_step = -1
        self._log_path = os.path.join(output_dir, "probe.log")

        # Deferred probe state (set by on_validation_end, consumed by on_train_batch_start)
        self._pending = False
        self._pending_step = -1
        self._pending_metrics: dict = {}

        if not os.path.exists(video_path):
            print(f"[InferenceProbe] WARNING: probe video not found: {video_path}", flush=True)
            self._enabled = False
        else:
            self._enabled = True

    def _log(self, msg: str):
        """Write to stdout (flushed) and to a sidecar log file."""
        print(msg, flush=True)
        try:
            os.makedirs(self.output_dir, exist_ok=True)
            with open(self._log_path, "a") as f:
                f.write(msg + "\n")
        except Exception:
            pass

    def on_validation_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
        """
        Do NOT look for checkpoints here — ModelCheckpoint hasn't saved yet.
        Just record that we want to run a probe and defer to on_train_batch_start.
        """
        rank = getattr(trainer, "global_rank", "?")
        is_zero = trainer.is_global_zero
        self._log(
            f"[InferenceProbe] on_validation_end rank={rank} is_global_zero={is_zero}"
            f" step={trainer.global_step} (deferring probe to next batch start)"
        )

        if not is_zero:
            return
        if not self._enabled:
            return
        if trainer.global_step == self._last_step:
            return

        self._pending = True
        self._pending_step = trainer.global_step
        self._pending_metrics = dict(trainer.callback_metrics)

    def on_train_batch_start(
        self,
        trainer: pl.Trainer,
        pl_module: pl.LightningModule,
        batch,
        batch_idx: int,
    ):
        """
        Fires after ModelCheckpoint has saved its files, so last.ckpt is available.
        Consume the pending probe request exactly once.
        """
        if not trainer.is_global_zero:
            return
        if not self._pending:
            return
        if not self._enabled:
            return

        self._pending = False
        step = self._pending_step
        metrics = self._pending_metrics

        ckpt_cb = trainer.checkpoint_callback
        if ckpt_cb is None:
            self._log("[InferenceProbe] No checkpoint callback found, skipping.")
            return

        dirpath = getattr(ckpt_cb, "dirpath", None)
        if not dirpath:
            self._log("[InferenceProbe] checkpoint_callback has no dirpath, skipping.")
            return

        ckpt_path = None
        candidate = os.path.join(dirpath, "last.ckpt")
        if os.path.exists(candidate):
            ckpt_path = candidate
        else:
            ckpts = sorted(glob.glob(os.path.join(dirpath, "*.ckpt")), key=os.path.getmtime)
            if ckpts:
                ckpt_path = ckpts[-1]

        self._log(f"[InferenceProbe] step={step}  ckpt={ckpt_path!r}")

        if not ckpt_path:
            self._log("[InferenceProbe] No checkpoint found on disk yet, skipping.")
            return

        self._last_step = step
        self._run_probe(ckpt_path, step, metrics)

    def _run_probe(self, ckpt_path: str, step: int, metrics: dict):
        try:
            import torch
            from badas.inference import BADAS, BADASConfig, SmoothingConfig
            from badas.visualization import visualize_predictions

            val_ap = float(metrics.get("val_ap", 0.0))
            self._log(
                f"[InferenceProbe] Running inference  val_ap={val_ap:.4f}"
                f"  ckpt={os.path.basename(ckpt_path)}"
            )

            config = BADASConfig(
                smoothing=SmoothingConfig(enabled=True),
                use_compile=False,
                quiet=True,
            )
            predictor = BADAS(ckpt_path, config=config)
            predictions = predictor.predict_video(self.video_path, stride=self.stride, quiet=True)

            ckpt_stem = os.path.splitext(os.path.basename(ckpt_path))[0]
            save_path = os.path.join(
                self.output_dir, f"probe_step{step:06d}_ap{val_ap:.4f}_{ckpt_stem}.png"
            )
            os.makedirs(self.output_dir, exist_ok=True)

            visualize_predictions(
                predictions,
                title=f"Collision Risk — step {step:,}  val_ap={val_ap:.4f}",
                show=False,
                save_path=save_path,
            )

            self._log(f"[InferenceProbe] Saved → {save_path}")

            del predictor
            torch.cuda.empty_cache()

        except Exception as e:
            import traceback

            self._log(f"[InferenceProbe] Failed: {e}\n{traceback.format_exc()}")
