"""Tests for weight resolution and HuggingFace Hub auto-download system."""

from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest


class TestWeightRegistries:
    """Test that all weight registries are consistent."""

    def test_all_model_files_have_sizes(self):
        from sltk.weights.hub import MODEL_FILES, MODEL_SIZES_MB

        for model in MODEL_FILES:
            assert model in MODEL_SIZES_MB, f"{model} missing from MODEL_SIZES_MB"

    def test_all_model_files_have_env_vars(self):
        from sltk.weights.hub import MODEL_ENV_VARS, MODEL_FILES

        for model in MODEL_FILES:
            assert model in MODEL_ENV_VARS, f"{model} missing from MODEL_ENV_VARS"

    def test_weight_to_model_maps_to_valid_models(self):
        from sltk.weights.hub import MODEL_FILES, WEIGHT_TO_MODEL

        for weight_name, model in WEIGHT_TO_MODEL.items():
            assert model in MODEL_FILES, (
                f"WEIGHT_TO_MODEL['{weight_name}'] = '{model}' " f"but '{model}' not in MODEL_FILES"
            )

    def test_weight_paths_relative_consistent_with_hub(self):
        """Every entry in WEIGHT_PATHS_RELATIVE that is in WEIGHT_TO_MODEL
        should have its relative path present in the corresponding MODEL_FILES list."""
        from sltk.extraction.weights import WEIGHT_PATHS_RELATIVE
        from sltk.weights.hub import MODEL_FILES, WEIGHT_TO_MODEL

        for name, rel_path in WEIGHT_PATHS_RELATIVE.items():
            model = WEIGHT_TO_MODEL.get(name)
            if model is None:
                continue  # Not all weights are auto-downloadable
            assert rel_path in MODEL_FILES[model], (
                f"WEIGHT_PATHS_RELATIVE['{name}'] = '{rel_path}' " f"not in MODEL_FILES['{model}']"
            )

    def test_teaser_in_all_registries(self):
        """Verify teaser is fully registered across all weight system dicts."""
        from sltk.extraction.weights import ENV_OVERRIDES, WEIGHT_PATHS_RELATIVE
        from sltk.weights.hub import (
            MODEL_ENV_VARS,
            MODEL_FILES,
            MODEL_SIZES_MB,
            WEIGHT_TO_MODEL,
        )

        # Hub registries
        assert "teaser" in MODEL_FILES
        assert "Teaser.pt" in MODEL_FILES["teaser"]
        assert "teaser" in MODEL_SIZES_MB
        assert MODEL_SIZES_MB["teaser"] == 360.0
        assert "teaser" in MODEL_ENV_VARS
        assert MODEL_ENV_VARS["teaser"] == "SLTK_TEASER_CHECKPOINT"
        assert "teaser_checkpoint" in WEIGHT_TO_MODEL
        assert WEIGHT_TO_MODEL["teaser_checkpoint"] == "teaser"

        # weights.py registries
        assert "teaser_checkpoint" in WEIGHT_PATHS_RELATIVE
        assert WEIGHT_PATHS_RELATIVE["teaser_checkpoint"] == "Teaser.pt"
        assert "teaser_checkpoint" in ENV_OVERRIDES
        assert ENV_OVERRIDES["teaser_checkpoint"] == "SLTK_TEASER_CHECKPOINT"


class TestGetWeightPath:
    """Tests for the get_weight_path resolution chain."""

    def test_custom_path_takes_priority(self, tmp_path):
        from sltk.extraction.weights import get_weight_path

        custom = tmp_path / "custom_model.pt"
        custom.touch()

        result = get_weight_path("teaser_checkpoint", custom_path=str(custom))
        assert result == custom

    def test_custom_path_nonexistent_falls_through(self, tmp_path):
        from sltk.extraction.weights import get_weight_path

        result = get_weight_path(
            "teaser_checkpoint",
            custom_path=str(tmp_path / "nonexistent.pt"),
        )
        # Should not return the nonexistent custom path
        assert result is None or result != tmp_path / "nonexistent.pt"

    def test_env_var_override(self, tmp_path, monkeypatch):
        from sltk.extraction.weights import get_weight_path

        env_model = tmp_path / "env_teaser.pt"
        env_model.touch()
        monkeypatch.setenv("SLTK_TEASER_CHECKPOINT", str(env_model))

        result = get_weight_path("teaser_checkpoint")
        assert result == env_model

    def test_env_var_nonexistent_falls_through(self, monkeypatch):
        from sltk.extraction.weights import get_weight_path

        monkeypatch.setenv("SLTK_TEASER_CHECKPOINT", "/nonexistent/Teaser.pt")

        # Should not return the nonexistent env path — falls through
        result = get_weight_path("teaser_checkpoint")
        if result is not None:
            assert result.exists()

    def test_cached_path_found(self, tmp_path, monkeypatch):
        """If weight exists in HF cache dir, return it without downloading."""
        from sltk.extraction.weights import get_weight_path

        # Point cache dir to tmp_path and create the file
        monkeypatch.setattr("sltk.weights.hub.get_cache_dir", lambda: tmp_path)
        cached = tmp_path / "Teaser.pt"
        cached.touch()

        # Clear env var so it doesn't interfere
        monkeypatch.delenv("SLTK_TEASER_CHECKPOINT", raising=False)

        result = get_weight_path("teaser_checkpoint")
        assert result is not None
        assert result == cached


class TestEnsureWeights:
    """Tests for the ensure_weights download function."""

    def test_unknown_model_raises(self):
        from sltk.weights.hub import ensure_weights

        with pytest.raises(ValueError, match="Unknown model"):
            ensure_weights("nonexistent_model")

    def test_all_cached_returns_immediately(self, tmp_path, monkeypatch):
        from sltk.weights.hub import ensure_weights

        monkeypatch.setattr("sltk.weights.hub.get_cache_dir", lambda: tmp_path)

        # Create the cached file
        cached = tmp_path / "Teaser.pt"
        cached.write_bytes(b"fake weights")

        paths = ensure_weights("teaser")
        assert "Teaser.pt" in paths
        assert paths["Teaser.pt"] == cached

    @patch("sltk.weights.hub._prompt_download", return_value=False)
    def test_declined_download_raises(self, mock_prompt, tmp_path, monkeypatch):
        from sltk.weights.hub import ensure_weights

        monkeypatch.setattr("sltk.weights.hub.get_cache_dir", lambda: tmp_path)
        monkeypatch.delenv("SLTK_AUTO_DOWNLOAD", raising=False)

        with pytest.raises(RuntimeError, match="Download declined"):
            ensure_weights("teaser", confirm=True)

    @patch("sltk.weights.hub._should_auto_download", return_value=True)
    def test_auto_download_skips_prompt(self, mock_auto, tmp_path, monkeypatch):
        from sltk.weights.hub import ensure_weights

        monkeypatch.setattr("sltk.weights.hub.get_cache_dir", lambda: tmp_path)

        mock_hf_download = MagicMock(return_value=str(tmp_path / "Teaser.pt"))
        with patch("sltk.weights.hub.hf_hub_download", mock_hf_download, create=True):
            # Patch the import inside ensure_weights
            with patch.dict(
                "sys.modules",
                {
                    "huggingface_hub": MagicMock(
                        hf_hub_download=mock_hf_download,
                    )
                },
            ):
                ensure_weights("teaser", confirm=True)

        mock_hf_download.assert_called_once()
        call_kwargs = mock_hf_download.call_args
        assert (
            call_kwargs.kwargs["filename"] == "Teaser.pt"
            or call_kwargs[1].get("filename") == "Teaser.pt"
        )

    def test_confirm_false_skips_prompt(self, tmp_path, monkeypatch):
        from sltk.weights.hub import ensure_weights

        monkeypatch.setattr("sltk.weights.hub.get_cache_dir", lambda: tmp_path)
        monkeypatch.delenv("SLTK_AUTO_DOWNLOAD", raising=False)

        mock_hf_download = MagicMock(return_value=str(tmp_path / "Teaser.pt"))

        with patch.dict(
            "sys.modules",
            {
                "huggingface_hub": MagicMock(
                    hf_hub_download=mock_hf_download,
                )
            },
        ):
            ensure_weights("teaser", confirm=False)

        mock_hf_download.assert_called_once()

    def test_prompt_download_shows_teaser_info(self, monkeypatch):
        """Verify _prompt_download displays correct teaser info."""
        from io import StringIO

        from sltk.weights.hub import _prompt_download

        stderr = StringIO()
        monkeypatch.setattr("sys.stderr", stderr)
        # Simulate non-interactive (no TTY)
        monkeypatch.setattr("sys.stdin", MagicMock(isatty=lambda: False))

        result = _prompt_download("teaser")

        assert result is False  # Non-interactive declines
        output = stderr.getvalue()
        assert "TEASER" in output
        assert "360 MB" in output
        assert "SLTK_TEASER_CHECKPOINT" in output
        assert "fiskenai/vltk" in output


class TestEnsureWeightByName:
    """Tests for ensure_weight_by_name integration."""

    def test_unknown_name_returns_none(self):
        from sltk.weights.hub import ensure_weight_by_name

        result = ensure_weight_by_name("totally_unknown_weight")
        assert result is None

    def test_teaser_checkpoint_resolves(self, tmp_path, monkeypatch):
        from sltk.weights.hub import ensure_weight_by_name

        monkeypatch.setattr("sltk.weights.hub.get_cache_dir", lambda: tmp_path)
        cached = tmp_path / "Teaser.pt"
        cached.write_bytes(b"fake")

        result = ensure_weight_by_name("teaser_checkpoint")
        assert result == cached
