"""Guard the single-repository download and mounted-checkpoint loading paths."""
import sys
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from flux3_engine import Engine, download_weights
from model_config import GGUF_FILENAME, Settings


@pytest.mark.parametrize('mode', ['q8_0', 'bf16'])
def test_download_only_uses_selected_repository(monkeypatch, mode):
    hub = MagicMock()
    monkeypatch.setitem(sys.modules, 'huggingface_hub', hub)
    settings = Settings('example/complete-copy', 'pinned-commit', mode)
    download_weights(settings)
    hub.snapshot_download.assert_called_once()
    call = hub.snapshot_download.call_args
    assert call.args == (settings.model_id,)
    assert call.kwargs['revision'] == settings.revision
    assert (GGUF_FILENAME in call.kwargs['allow_patterns']) == (mode == 'q8_0')
    assert ('transformer/*.safetensors' in call.kwargs['allow_patterns']) == (mode == 'bf16')
    hub.hf_hub_download.assert_not_called()


@pytest.mark.parametrize('missing_gguf', [False, True])
def test_mounted_q8_loads_locally_without_hub_cache(monkeypatch, tmp_path, missing_gguf):
    checkpoint = tmp_path / GGUF_FILENAME
    if not missing_gguf:
        checkpoint.parent.mkdir()
        checkpoint.write_bytes(b'test-checkpoint')
    torch, diffusers, transformers, hub = (MagicMock() for _ in range(4))
    for name, module in [('torch', torch), ('diffusers', diffusers),
                         ('transformers', transformers), ('huggingface_hub', hub)]:
        monkeypatch.setitem(sys.modules, name, module)
    pipe = MagicMock()
    blocks = diffusers.ModularPipeline.from_pretrained.return_value.blocks
    # Match the pinned Diffusers API: init_pipeline accepts no Hub kwargs.
    blocks.init_pipeline.side_effect = lambda path: pipe

    def update_components(**components):
        for name, component in components.items():
            setattr(pipe, name, component)
            component.parameters.return_value = [SimpleNamespace(device='cuda:0')]

    pipe.update_components.side_effect = update_components
    with patch('flux3_engine.download_weights') as download:
        engine = Engine(settings=Settings(quantization='q8_0'), model_path=tmp_path)
        if missing_gguf:
            with pytest.raises(FileNotFoundError, match='complete Socialtrait checkpoint'):
                engine.load()
            diffusers.ModularPipeline.from_pretrained.assert_not_called()
        else:
            engine.load()
            args, kwargs = diffusers.Flux3Transformer2DModel.from_single_file.call_args
            assert args == (str(checkpoint),)
            assert kwargs['config'] == str(tmp_path)
            assert kwargs['local_files_only'] is True
            diffusers.ModularPipeline.from_pretrained.assert_called_once_with(
                str(tmp_path), local_files_only=True)
            assert engine.pipe is pipe
        download.assert_not_called()
        assert not hub.mock_calls