# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------

"""Backward-compatibility tests for loading past-version RF-DETR checkpoints.

Each test loads a checkpoint generated by a specific released version of
``rfdetr`` (produced by ``tests/legacy/generate_checkpoint.py``) using the
*current* codebase.  Checkpoint files are expected at::

    tests/legacy/checkpoints/checkpoint_v<version>.pth

The ``ci-legacy-checkpoints.yml`` workflow generates and downloads them before
running this suite.  Locally, without pre-generated files, the fixture that
loads them fails fast with a clear "checkpoint missing" message.

The list of versions under test is authoritative in
``tests/legacy/checkpoint_versions.txt``.
"""

from __future__ import annotations

import os
import pickle
from pathlib import Path

import numpy as np
import pytest
import supervision as sv
import torch
from PIL import Image

from tests.legacy.generate_checkpoint import _get_reference_image_path, _top_detection


def _skip_or_fail_missing_reference(version: str) -> None:
    """Skip locally, but hard-fail in CI, when a checkpoint has no reference_prediction.

    ``ci-legacy-checkpoints.yml`` always generates checkpoints with ``--use-pretrained``,
    so a missing ``reference_prediction`` under CI signals a real generation-pipeline
    regression, not the expected local-dev shortcut (skipping ``--use-pretrained`` to avoid
    a slow pretrained-weights download). A plain ``pytest.skip`` here would let that
    regression disappear silently if ``--use-pretrained`` were ever accidentally dropped
    from the workflow — this is the *only* check that catches load-path weight corruption.

    Args:
        version: rfdetr version string, used in the skip/fail message.

    Examples:
        Under CI (``GITHUB_ACTIONS=true``) the call hard-fails; locally it skips instead — both
        outcomes raise a pytest exception so neither can be demonstrated purely by value. The
        ``TestSkipOrFailMissingReference`` class tests each branch directly.

        >>> callable(_skip_or_fail_missing_reference)
        True
    """
    message = f"v{version} checkpoint has no reference_prediction (generated without --use-pretrained)"
    if os.environ.get("GITHUB_ACTIONS") == "true":
        pytest.fail(message)
    pytest.skip(message)


class TestSkipOrFailMissingReference:
    """Unit tests for _skip_or_fail_missing_reference — the CI-vs-local skip/fail split.

    Pins the one safety property this helper exists to guarantee: under CI it hard-fails
    (catching a real generation-pipeline regression), while local dev runs soften to a skip.
    An accidental inversion of the GITHUB_ACTIONS check must fail these tests, not pass silently.
    """

    def test_fails_under_ci(self, monkeypatch: pytest.MonkeyPatch) -> None:
        """Under CI (GITHUB_ACTIONS=true), a missing reference_prediction must hard-fail."""
        monkeypatch.setenv("GITHUB_ACTIONS", "true")

        with pytest.raises(pytest.fail.Exception):
            _skip_or_fail_missing_reference("1.5.0")

    def test_skips_locally(self, monkeypatch: pytest.MonkeyPatch) -> None:
        """Locally (GITHUB_ACTIONS unset), a missing reference_prediction must skip, not fail."""
        monkeypatch.delenv("GITHUB_ACTIONS", raising=False)

        with pytest.raises(pytest.skip.Exception):
            _skip_or_fail_missing_reference("1.5.0")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_VERSIONS_FILE = Path(__file__).parent / "checkpoint_versions.txt"
_CHECKPOINTS_DIR = Path(__file__).parent / "checkpoints"
_LEGACY_VISUALIZATIONS_DIR = Path(__file__).parent / "legacy_predictions"
_FROZEN_DEPENDENCIES_DIR = Path(__file__).parent / "frozen_dependencies"


def _read_versions(versions_file: Path | None = None) -> list[str]:
    """Parse a versions file and return a list of version strings.

    Args:
        versions_file: Path to the versions file.  Defaults to
            ``_VERSIONS_FILE`` (``tests/legacy/checkpoint_versions.txt``).

    Returns:
        Ordered list of version strings (comment lines and blank lines skipped).

    Examples:
        >>> import tempfile
        >>> from pathlib import Path
        >>> with tempfile.TemporaryDirectory() as d:
        ...     p = Path(d) / "versions.txt"
        ...     _ = p.write_text("# comment\\n1.5.0\\n\\n1.6.0\\n")
        ...     _read_versions(p)
        ['1.5.0', '1.6.0']
        >>> _read_versions(Path("/nonexistent/versions.txt"))
        []
    """
    path = versions_file if versions_file is not None else _VERSIONS_FILE
    if not path.is_file():
        return []
    versions: list[str] = []
    for line in path.read_text().splitlines():
        stripped = line.strip()
        if stripped and not stripped.startswith("#"):
            versions.append(stripped)
    return versions


def _checkpoint_path(version: str) -> Path:
    """Return the expected checkpoint path for *version*.

    Args:
        version: rfdetr version string, e.g. ``"1.5.0"``.

    Returns:
        Path object (may not yet exist — caller must check).

    Examples:
        >>> p = _checkpoint_path("1.5.0")
        >>> p.name
        'checkpoint_v1.5.0.pth'
    """
    return _CHECKPOINTS_DIR / f"checkpoint_v{version}.pth"


_LEGACY_VERSIONS: list[str] = _read_versions()


def _prediction_visualization_path(version: str) -> Path:
    """Return the artifact path for a legacy checkpoint prediction image.

    Args:
        version: Legacy RF-DETR version encoded in the checkpoint filename.

    Returns:
        Ignored test-artifact path whose filename includes the legacy version.

    Examples:
        >>> p = _prediction_visualization_path("1.5.0")
        >>> p.name
        'prediction_legacy_v1.5.0.png'
    """
    return _LEGACY_VISUALIZATIONS_DIR / f"prediction_legacy_v{version}.png"


def _save_prediction_visualization(detections: sv.Detections, image_path: Path, output_path: Path) -> None:
    """Save the reference image annotated with all reloaded-model detections.

    Args:
        detections: Predictions returned while loading a legacy checkpoint.
        image_path: Fixed reference image used for compatibility inference.
        output_path: Destination for the annotated PNG artifact.

    Examples:
        Requires a real image file and ``supervision`` annotators — not runnable standalone without
        a full detections object. The ``TestCheckpointBackwardCompat`` class tests the full pipeline.

        >>> callable(_save_prediction_visualization)
        True
    """
    with Image.open(image_path) as image:
        scene = np.asarray(image.convert("RGB")).copy()

    labels = [
        f"class_id={class_id} confidence={confidence:.3f}"
        for class_id, confidence in zip(detections.class_id, detections.confidence)
    ]
    scene = sv.BoxAnnotator().annotate(scene=scene, detections=detections)
    scene = sv.LabelAnnotator().annotate(scene=scene, detections=detections, labels=labels)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(scene).save(output_path)


def test_prediction_visualization_path_includes_legacy_version() -> None:
    """Prediction visualization filenames identify the checkpoint version."""
    assert _prediction_visualization_path("1.5.2") == (_LEGACY_VISUALIZATIONS_DIR / "prediction_legacy_v1.5.2.png")


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module", params=_LEGACY_VERSIONS, ids=[f"v{version}" for version in _LEGACY_VERSIONS])
def loaded_checkpoint(request: pytest.FixtureRequest) -> tuple[str, Path, dict]:
    """Load one legacy checkpoint file, shared across all compat assertions in this module.

    Args:
        request: Pytest fixture request; ``request.param`` is the rfdetr version string.

    Returns:
        Tuple of (version, checkpoint path, loaded checkpoint dict).
    """
    version: str = request.param
    ckpt_path = _checkpoint_path(version)
    assert ckpt_path.is_file(), f"Checkpoint missing: {ckpt_path}"
    ckpt: dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    return version, ckpt_path, ckpt


def test_torch_load_on_corrupt_checkpoint_raises_clear_error(tmp_path: Path) -> None:
    """A present-but-corrupt .pth file must fail with torch's own load error, not hang or crash silently.

    ``loaded_checkpoint`` only asserts ``ckpt_path.is_file()`` (the missing-file case) before
    calling ``torch.load(...)`` directly with no error wrapping — a present-but-corrupt checkpoint
    would otherwise raise mid-fixture with no assertion pinning what is raised. This test does not
    change production code (there is no wrapping to change); it pins the natural failure mode of
    ``torch.load`` on a truncated/corrupt file so a future change to the loading code cannot
    silently swallow the failure or turn it into an unrelated, harder-to-diagnose traceback.

    Args:
        tmp_path: Pytest-provided temporary directory.
    """
    corrupt_path = tmp_path / "checkpoint_vcorrupt.pth"
    corrupt_path.write_bytes(b"not a valid torch checkpoint - truncated garbage")

    with pytest.raises(pickle.UnpicklingError, match=r".+"):
        torch.load(corrupt_path, map_location="cpu", weights_only=False)


class TestCheckpointBackwardCompat:
    """Load past-version checkpoints with the current rfdetr codebase.

    Each test parametrized over version (via the ``loaded_checkpoint`` fixture)
    exercises the full L1 facade loading path (``RFDETRSmall(...,
    pretrain_weights=path)``), which internally calls
    :func:`rfdetr.models.weights.load_pretrain_weights`.  A successful load
    without exception is the acceptance criterion; we also verify that the
    returned model has the expected number of classes to confirm head alignment.
    """

    def test_load_checkpoint_with_current_code(self, loaded_checkpoint: tuple[str, Path, dict]) -> None:
        """Loading a past-version checkpoint must not raise any exception.

        The test uses ``RFDETRSmall`` (the lightest non-nano detection model)
        with ``num_classes`` auto-inferred from the checkpoint, which exercises
        the same loading path used in production fine-tuning workflows.

        Args:
            loaded_checkpoint: Shared (version, path, raw checkpoint dict) fixture.
        """
        from rfdetr import RFDETRSmall

        version, ckpt_path, raw = loaded_checkpoint

        # Read the stored class count so we can pass an explicit num_classes and avoid the
        # constructor's 90-class COCO default triggering a head mismatch with the checkpoint's
        # actual class count — 2 for a locally-generated (non-pretrained) checkpoint, or 90 for
        # the CI --use-pretrained path (see generate_checkpoint.py); either way, RFDETRSmall's
        # implicit default must not silently diverge from what this checkpoint was saved with.
        stored_num_classes: int | None = None
        if "model" in raw and "class_embed.bias" in raw["model"]:
            stored_num_classes = int(raw["model"]["class_embed.bias"].shape[0]) - 1

        kwargs: dict = dict(pretrain_weights=str(ckpt_path), device="cpu")
        if stored_num_classes is not None:
            kwargs["num_classes"] = stored_num_classes

        # This must not raise — any exception indicates a loading regression.
        loaded = RFDETRSmall(**kwargs)

        # Sanity-check: model is on CPU and has at least one parameter.
        params = list(loaded.model.model.parameters())
        assert params, f"Loaded model for v{version} has no parameters."
        first_device = next(iter(params)).device
        assert str(first_device) == "cpu", f"Expected model on CPU, got {first_device}."

    def test_checkpoint_contains_expected_keys(self, loaded_checkpoint: tuple[str, Path, dict]) -> None:
        """Checkpoint file must contain 'model' and 'args' keys.

        Validates that ``generate_checkpoint.py`` produced a structurally
        valid checkpoint file before attempting a full model load.

        Args:
            loaded_checkpoint: Shared (version, path, raw checkpoint dict) fixture.
        """
        version, _, ckpt = loaded_checkpoint

        assert "model" in ckpt, f"v{version} checkpoint missing 'model' key; keys present: {list(ckpt.keys())}"
        assert isinstance(ckpt["model"], dict), (
            f"v{version} checkpoint['model'] must be a dict, got {type(ckpt['model']).__name__}"
        )
        assert ckpt["model"], f"v{version} checkpoint['model'] state dict is empty."

        assert "args" in ckpt, f"v{version} checkpoint missing 'args' key; keys present: {list(ckpt.keys())}"

    def test_checkpoint_class_embed_shape_is_valid(self, loaded_checkpoint: tuple[str, Path, dict]) -> None:
        """'class_embed.bias' must be present and have shape [N] with N >= 2.

        ``load_pretrain_weights`` reads this tensor to determine the checkpoint
        class count; a missing or malformed tensor causes a ``KeyError`` or
        index error during loading.

        Args:
            loaded_checkpoint: Shared (version, path, raw checkpoint dict) fixture.
        """
        version, _, ckpt = loaded_checkpoint
        model_sd: dict[str, torch.Tensor] = ckpt.get("model", {})

        assert "class_embed.bias" in model_sd, (
            f"v{version} checkpoint missing 'class_embed.bias'; "
            f"keys starting with 'class': "
            f"{[k for k in model_sd if k.startswith('class')]}"
        )

        bias = model_sd["class_embed.bias"]
        assert bias.ndim == 1 and bias.shape[0] >= 2, (
            f"v{version} 'class_embed.bias' must be 1-D with ≥ 2 entries, got shape {tuple(bias.shape)}"
        )

    def test_prediction_matches_reference(self, loaded_checkpoint: tuple[str, Path, dict]) -> None:
        """Reloading a real-pretrained-weights checkpoint must predict the same thing.

        The structural checks above cannot catch a load-path regression that
        silently permutes or corrupts weights — the model would still "load"
        without raising. This exercises the real inference path instead: the
        checkpoint stores the top detection captured at generation time on a
        fixed reference image (``generate_checkpoint.py --use-pretrained``);
        reloading with current code and re-running the same image must
        reproduce it. Skipped locally for checkpoints generated without
        ``--use-pretrained`` (no ``reference_prediction`` key); hard-fails
        instead under CI, where ``--use-pretrained`` is always passed.

        Args:
            loaded_checkpoint: Shared (version, path, raw checkpoint dict) fixture.
        """
        from rfdetr import RFDETRSmall

        version, ckpt_path, ckpt = loaded_checkpoint
        reference = ckpt.get("reference_prediction")
        if reference is None:
            _skip_or_fail_missing_reference(version)

        loaded = RFDETRSmall(pretrain_weights=str(ckpt_path), device="cpu")
        reference_image_path = _get_reference_image_path()
        detections = loaded.predict(str(reference_image_path), threshold=0.5)
        visualization_path = _prediction_visualization_path(version)
        _save_prediction_visualization(detections, reference_image_path, visualization_path)
        assert visualization_path.is_file(), f"v{version}: prediction visualization was not saved: {visualization_path}"
        actual = _top_detection(detections)

        # class_id (not class_name) is the identity check: class-name mapping
        # logic has changed across rfdetr releases, while the raw category id
        # is what a load-path regression would actually corrupt.
        assert actual["class_id"] == reference["class_id"], (
            f"v{version}: top detection class_id changed after reload (expected "
            f"{reference['class_id']} {reference['class_name']!r}, got {actual['class_id']} {actual['class_name']!r})"
        )

        confidence_delta = abs(actual["confidence"] - reference["confidence"])
        assert confidence_delta < 0.05, (
            f"v{version}: top detection confidence drifted by {confidence_delta:.4f} after reload "
            f"(expected {reference['confidence']:.4f}, got {actual['confidence']:.4f})"
        )

        iou = float(sv.box_iou_batch(np.array([reference["xyxy"]]), np.array([actual["xyxy"]]))[0, 0])
        assert iou > 0.9, f"v{version}: top detection box IoU dropped to {iou:.4f} after reload (expected > 0.9)"

    def test_load_checkpoint_with_custom_resolution_interpolates_position_embeddings(
        self, loaded_checkpoint: tuple[str, Path, dict]
    ) -> None:
        """Loading with an explicit non-default resolution must not raise and must keep predictions stable.

        ``RFDETRSmall`` defaults to ``resolution=512`` (a 32x32 DINOv2 patch grid,
        ``positional_encoding_size=32``). Passing an explicit ``resolution=1280`` (an 80x80 grid)
        forces ``n_source != n_target`` in ``interpolate_position_embeddings``
        (``src/rfdetr/models/weights.py:205-262``, invoked from ``load_pretrain_weights`` at
        line 554), which must bicubic-interpolate the checkpoint's positional embeddings rather
        than raising a shape-mismatch ``RuntimeError`` from ``load_state_dict``. Skipped for
        checkpoints generated without ``--use-pretrained`` (no ``reference_prediction`` key).

        Args:
            loaded_checkpoint: Shared (version, path, raw checkpoint dict) fixture.
        """
        from rfdetr import RFDETRSmall

        version, ckpt_path, ckpt = loaded_checkpoint
        reference = ckpt.get("reference_prediction")
        if reference is None:
            pytest.skip(f"v{version} checkpoint has no reference_prediction (generated without --use-pretrained)")

        # This must not raise — interpolate_position_embeddings must handle n_source != n_target cleanly.
        loaded = RFDETRSmall(pretrain_weights=str(ckpt_path), resolution=1280, device="cpu")

        # Confirm the resolution override actually drove a PE-grid mismatch (32x32 -> 80x80),
        # i.e. that the interpolation code path was genuinely exercised, not skipped as a no-op.
        assert loaded.model_config.positional_encoding_size == 80, (
            f"v{version}: expected positional_encoding_size=80 (1280 // patch_size=16) after "
            f"resolution=1280 override, got {loaded.model_config.positional_encoding_size}"
        )

        params = list(loaded.model.model.parameters())
        assert params, f"Loaded model for v{version} (resolution=1280) has no parameters."

        reference_image_path = _get_reference_image_path()
        detections = loaded.predict(str(reference_image_path), threshold=0.5)
        actual = _top_detection(detections)

        # Interpolating positional embeddings to a much larger grid shifts confidence/box precision
        # more than a native-resolution reload, so this checks identity stability of the top
        # detection (the load path did not corrupt or permute weights), not exact reproduction.
        assert actual["class_id"] == reference["class_id"], (
            f"v{version}: top detection class_id changed after reload at resolution=1280 (expected "
            f"{reference['class_id']} {reference['class_name']!r}, got {actual['class_id']} {actual['class_name']!r})"
        )


# ---------------------------------------------------------------------------
# Sanity check for checkpoint_versions.txt parsing
# ---------------------------------------------------------------------------


def test_legacy_versions_are_parsed() -> None:
    """checkpoint_versions.txt must parse to at least one version."""
    assert _LEGACY_VERSIONS, "failed parsing legacy versions"


@pytest.mark.parametrize(
    ("content", "expected"),
    [
        pytest.param("1.5.0\n1.6.0\n", ["1.5.0", "1.6.0"], id="plain-versions"),
        pytest.param("# header comment\n1.5.0\n# trailing comment\n", ["1.5.0"], id="comment-lines-skipped"),
        pytest.param("\n\n1.5.0\n\n\n1.6.0\n\n", ["1.5.0", "1.6.0"], id="blank-lines-skipped"),
        pytest.param("  1.5.0  \n\t1.6.0\t\n", ["1.5.0", "1.6.0"], id="whitespace-stripped"),
        pytest.param("# only comments\n# nothing else\n", [], id="comments-only-yields-empty"),
        pytest.param("", [], id="empty-file-yields-empty"),
    ],
)
def test_read_versions_parsing_edge_cases(tmp_path: Path, content: str, expected: list[str]) -> None:
    """_read_versions must skip blank lines and comments, and strip whitespace, for each edge case."""
    versions_file = tmp_path / "checkpoint_versions.txt"
    versions_file.write_text(content)

    assert _read_versions(versions_file) == expected


def test_read_versions_missing_file_yields_empty() -> None:
    """_read_versions must return an empty list (not raise) when the file does not exist."""
    assert _read_versions(Path("/nonexistent/checkpoint_versions.txt")) == []


def test_every_version_has_matching_frozen_requirements_file() -> None:
    """Every version listed in checkpoint_versions.txt must have a frozen_dependencies/req-<version>.txt.

    The CI ``generate`` matrix job (``ci-legacy-checkpoints.yml``) installs
    ``tests/legacy/frozen_dependencies/req-<version>.txt`` for each version in
    ``checkpoint_versions.txt``; a version listed without a matching file
    would fail the CI matrix at install time with no local signal beforehand.
    """
    missing = [
        version for version in _LEGACY_VERSIONS if not (_FROZEN_DEPENDENCIES_DIR / f"req-{version}.txt").is_file()
    ]
    assert not missing, (
        f"checkpoint_versions.txt lists version(s) with no matching frozen requirements file: {missing} "
        f"(expected tests/legacy/frozen_dependencies/req-<version>.txt for each)"
    )
