"""P4.5 - the model residency of the pack (ModelPatcher) and ``ensure_flow``.

The pack wraps its three model components (the KEEP net with its flownet, the
face detector, the face parser) in ``ModelPatcher`` instances and lets ComfyUI
arbitrate the VRAM.  The tests run against the **mocked**
``comfy.model_management`` of the conftest: every residency call is recorded,
so the tests can assert what the pack asks ComfyUI to do - and that it asks
exactly once per execute.
"""

from __future__ import annotations

import hashlib
import importlib
import logging
import os

import pytest
import torch


@pytest.fixture(scope="module")
def loader_module():
    return importlib.import_module("ComfyUI_KEEP.src.keep_model_loader")


@pytest.fixture(scope="module")
def node_module():
    return importlib.import_module("ComfyUI_KEEP.nodes.keep_face_upscale_image")


@pytest.fixture(scope="module")
def configs_module():
    return importlib.import_module("ComfyUI_KEEP.src.keep_model_configs")


@pytest.fixture
def model_management():
    """The mocked ``comfy.model_management`` (the calls are reset per test)."""
    module = importlib.import_module("comfy.model_management")
    module.calls.clear()
    return module


class _StubModule:
    """A minimal stand-in for a torch module (it only has to move)."""

    def __init__(self, name: str = "module") -> None:
        self.name = name
        self.moved_to: list = []

    def to(self, device):
        self.moved_to.append(device)
        return self


class _StubKeepNet(_StubModule):
    """The KEEP net - it owns the flownet as a submodule."""

    def __init__(self) -> None:
        super().__init__("keep_net")
        self.flownet = _StubModule("flownet")


class _StubFaceHelper:
    """The facelib helper - it owns the detector and the parser."""

    def __init__(self) -> None:
        self.face_detector = _StubModule("detector")
        self.face_parse = _StubModule("parser")
        self.device = None


@pytest.fixture
def stub_models():
    return _StubKeepNet(), _StubFaceHelper()


@pytest.fixture
def pack(loader_module, stub_models):
    """A real ``KEEPModelPack`` around the stub models."""
    keep_net, face_helper = stub_models
    return loader_module.KEEPModelPack(
        keep_net=keep_net,
        face_helper=face_helper,
        bg_upscale_model=None,
        face_upscale_model=None,
        model_type_str="KEEP",
    )


# ---------------------------------------------------------------------------
# the patchers
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_the_pack_builds_three_patchers(pack):
    """Three patchers: the KEEP net, the face detector and the face parser."""
    assert len(pack.patchers) == 3, pack.patchers
    assert pack.keep_patcher.model is pack.keep_net
    assert pack.detector_patcher.model is pack.face_helper.face_detector
    assert pack.parser_patcher.model is pack.face_helper.face_parse


@pytest.mark.unit
def test_exactly_one_patcher_contains_the_flownet(pack):
    """The flownet belongs to the KEEP net - it must not be patched twice.

    A second patcher for the flownet would count its VRAM twice (the plan
    P4.5).
    """
    assert pack.flownet is pack.keep_net.flownet
    assert pack.keep_patcher.model.flownet is pack.flownet
    # the flownet is not the model of a patcher of its own
    assert all(patcher.model is not pack.flownet for patcher in pack.patchers)
    assert pack.detector_patcher.model is not pack.flownet
    assert pack.parser_patcher.model is not pack.flownet


@pytest.mark.unit
def test_the_patchers_are_shared_by_the_packs_of_the_same_models(
    loader_module, stub_models
):
    """The patchers belong to the models - they are built once per build."""
    keep_net, face_helper = stub_models
    first = loader_module.KEEPModelPack(keep_net, face_helper, None, None, "KEEP")
    second = loader_module.KEEPModelPack(
        keep_net, face_helper, None, None, "KEEP", patchers=first.patchers_by_role
    )
    assert second.keep_patcher is first.keep_patcher
    assert second.patchers == first.patchers


# ---------------------------------------------------------------------------
# the residency
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_load_device_lets_comfyui_arbitrate_the_vram(pack, model_management):
    """The pack hands all its patchers to ``load_models_gpu`` - once."""
    pack.load_device()

    load_calls = [
        call for call in model_management.calls if call[0] == "load_models_gpu"
    ]
    assert len(load_calls) == 1, model_management.calls
    _, models, memory_required = load_calls[0]
    assert list(models) == list(pack.patchers)
    assert memory_required == 0


@pytest.mark.unit
def test_offload_moves_the_models_back_and_frees_the_vram(pack, model_management):
    """The release path: the models go back to the offload device and ComfyUI
    frees the VRAM."""
    pack.offload()

    names = [call[0] for call in model_management.calls]
    assert "unload_all_models" in names, model_management.calls
    assert "soft_empty_cache" in names, model_management.calls
    # the models were moved back to the offload device
    assert pack.keep_net.moved_to[-1] == pack.offload_device
    assert pack.face_helper.face_detector.moved_to[-1] == pack.offload_device
    assert pack.face_helper.face_parse.moved_to[-1] == pack.offload_device


@pytest.mark.unit
def test_the_cpu_mode_makes_the_same_calls(
    loader_module, model_management, monkeypatch
):
    """``get_torch_device()`` -> cpu: the same residency calls are made.

    The conftest mock answers ``cpu`` by default (``conftest.py``: the
    ``get_torch_device`` / ``unet_offload_device`` of the mocked
    ``comfy.model_management``), so patching the device to ``cpu`` alone
    would exercise the default path and prove nothing.  The CPU leg is
    therefore compared with a **GPU leg**: both devices must produce the
    same residency calls - the pack must not branch on the device.  A
    device dependent path is exactly the B7 CPU crash (the pack called
    ``torch.cuda.empty_cache()`` and crashed on a CPU-only host).
    """

    def residency_calls(device):
        monkeypatch.setattr(model_management, "get_torch_device", lambda: device)
        monkeypatch.setattr(model_management, "unet_offload_device", lambda: device)
        model_management.calls.clear()

        pack = loader_module.KEEPModelPack(
            _StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
        )
        # the patched device really reaches the pack - the leg is not a no-op
        assert pack.device == device
        assert pack.offload_device == device

        pack.load_device()
        pack.offload()
        return [call[0] for call in model_management.calls]

    gpu_calls = residency_calls(torch.device("cuda"))
    cpu_calls = residency_calls(torch.device("cpu"))

    assert cpu_calls == gpu_calls, (cpu_calls, gpu_calls)
    assert cpu_calls.count("load_models_gpu") == 1, cpu_calls
    assert cpu_calls.count("unload_all_models") == 1, cpu_calls
    assert cpu_calls.count("soft_empty_cache") == 1, cpu_calls


# ---------------------------------------------------------------------------
# the device placement of the face nets (round 4, D-R1)
# ---------------------------------------------------------------------------
def _pack_on_distinct_devices(loader_module, model_management, monkeypatch):
    """A pack whose pack device is distinct from its offload device.

    The conftest mock answers ``cpu`` for both device functions, so a
    placement on the wrong device would be indistinguishable from a correct
    one.  The patch gives the pack a distinct device object - recording
    only, no CUDA op ever runs (the CPU/GPU-leg pattern above).
    """
    monkeypatch.setattr(
        model_management, "get_torch_device", lambda: torch.device("cuda")
    )
    monkeypatch.setattr(
        model_management, "unet_offload_device", lambda: torch.device("cpu")
    )
    pack = loader_module.KEEPModelPack(
        _StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
    )
    assert pack.device != pack.offload_device  # the leg is not a no-op
    return pack


class _YoloDetector(_StubModule):
    """The YOLO detector shape - it stores its own ``device`` attribute.

    The vendored YOLO detector reads ``self.device`` when it pushes its
    inputs (the device diagnosis D8b), so a moved net with a stale attribute
    would run on the wrong device.
    """

    def __init__(self) -> None:
        super().__init__("detector")
        # a sentinel distinct from every leg device - a sync shows up
        self.device = torch.device("meta")


class _FailingParser(_StubModule):
    """A parser whose move raises - the placement must degrade cleanly."""

    def __init__(self) -> None:
        super().__init__("parser")
        self.device = torch.device("meta")

    def to(self, device):
        raise RuntimeError("the VRAM is full")


@pytest.fixture
def keep_records(loader_module, caplog):
    """The records of the pack logger (the pattern of ``test_keep_download``).

    Without a ComfyUI the pack logger owns a private handler and does not
    propagate, so the records never reach the root logger of pytest: the
    capture handler is attached to the pack logger itself.
    """
    loader_module.logger.addHandler(caplog.handler)
    try:
        yield caplog
    finally:
        loader_module.logger.removeHandler(caplog.handler)


@pytest.mark.unit
def test_the_face_helper_is_constructed_on_the_pack_device(
    loader_module, model_management, monkeypatch, loader
):
    """P-D1: the helper is constructed with the pack device.

    The vendored helper stores the construction device and builds its
    detector and parser nets on it (``face_restoration_helper.py:66-93``) -
    constructed on the offload device, the nets were born on the wrong side
    of the patcher flow.  The ``_StubHelper`` of the ``loader`` fixture
    records the constructor kwargs.
    """
    monkeypatch.setattr(
        model_management, "get_torch_device", lambda: torch.device("cuda")
    )
    monkeypatch.setattr(
        model_management, "unet_offload_device", lambda: torch.device("cpu")
    )
    # the device functions are read at construction - a loader built after
    # the patch sees the distinct pack device (the fixture's loader saw cpu)
    fresh = loader_module.KEEPModelLoader()
    pack = fresh.load_keep_model_pack("KEEP", "retinaface_resnet50")

    assert pack.face_helper.kwargs["device"] == pack.device
    assert pack.face_helper.kwargs["device"] != pack.offload_device


@pytest.mark.unit
def test_load_device_places_the_face_nets_on_the_pack_device(
    loader_module, model_management, monkeypatch
):
    """P-D2: ``load_device`` places the detector and the parser explicitly.

    The patcher flow may silently skip a module under VRAM pressure - the
    explicit placement makes the residency of the two nets a fact, not a
    hope, and the helper's own ``device`` stays synced.
    """
    pack = _pack_on_distinct_devices(loader_module, model_management, monkeypatch)

    pack.load_device()

    assert pack.face_helper.face_detector.moved_to[-1] == pack.device
    assert pack.face_helper.face_parse.moved_to[-1] == pack.device
    assert pack.face_helper.device == pack.device


@pytest.mark.unit
def test_load_device_falls_back_cleanly_when_a_net_cannot_move(
    loader_module, model_management, keep_records
):
    """P-D3: a net that cannot move degrades cleanly.

    A failed move must not raise out of ``load_device``, must be logged as a
    warning, and must leave the net's own ``.device`` attribute untouched -
    the sync belongs to the success path only (a failed move that still
    re-synced the attribute would point the net at a device its weights
    never reached).  The other net is still placed.
    """
    helper = _StubFaceHelper()
    failing = _FailingParser()
    helper.face_parse = failing
    pack = loader_module.KEEPModelPack(
        _StubKeepNet(), helper, None, None, "KEEP"
    )

    pack.load_device()  # must not raise

    # the detector is still placed, the helper device still synced
    assert pack.face_helper.face_detector.moved_to[-1] == pack.device
    assert pack.face_helper.device == pack.device
    # the failed parser: no move recorded, no attribute re-sync
    assert failing.moved_to == []
    assert failing.device == torch.device("meta")
    # the failure is a logged warning, not a silent skip
    warnings = [
        record for record in keep_records.records
        if record.levelno == logging.WARNING
    ]
    assert len(warnings) == 1, [r.getMessage() for r in keep_records.records]
    assert "face parser" in warnings[0].getMessage()
    assert "could not be placed" in warnings[0].getMessage()


@pytest.mark.unit
def test_load_device_syncs_a_detector_own_device_attribute(
    loader_module, model_management, monkeypatch
):
    """P-D4: the YOLO-style ``.device`` attribute follows the net - both ways.

    Synced on ``load_device`` (after a successful move, D-R1.2) and restored
    on ``offload`` (D-R1.3) - the attribute must never lag the weights.
    """
    monkeypatch.setattr(
        model_management, "get_torch_device", lambda: torch.device("cuda")
    )
    monkeypatch.setattr(
        model_management, "unet_offload_device", lambda: torch.device("cpu")
    )
    helper = _StubFaceHelper()
    helper.face_detector = _YoloDetector()
    pack = loader_module.KEEPModelPack(
        _StubKeepNet(), helper, None, None, "KEEP"
    )
    detector = pack.face_helper.face_detector
    assert detector.device == torch.device("meta")  # the sentinel so far

    pack.load_device()
    assert detector.moved_to[-1] == pack.device
    assert detector.device == pack.device

    pack.offload()
    assert detector.moved_to[-1] == pack.offload_device
    assert detector.device == pack.offload_device


@pytest.mark.unit
def test_the_cpu_mode_places_the_nets_on_cpu(
    loader_module, model_management, monkeypatch
):
    """P-D5: the CPU leg places the nets on ``cpu`` with the same calls.

    The CPU/GPU-leg comparison of the residency (the B7 guard): the
    placement is one more residency step that must not branch on the device,
    and on the CPU leg it must land on ``cpu`` - the inherent fallback of a
    CPU-only install, where ``get_torch_device()`` answers ``cpu``.
    """

    def leg(device):
        monkeypatch.setattr(model_management, "get_torch_device", lambda: device)
        monkeypatch.setattr(model_management, "unet_offload_device", lambda: device)
        model_management.calls.clear()
        pack = loader_module.KEEPModelPack(
            _StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
        )
        assert pack.device == device  # the patched device reaches the pack
        pack.load_device()
        pack.offload()
        placements = (
            pack.face_helper.face_detector.moved_to,
            pack.face_helper.face_parse.moved_to,
        )
        return [call[0] for call in model_management.calls], placements

    gpu_calls, _ = leg(torch.device("cuda"))
    cpu_calls, cpu_placements = leg(torch.device("cpu"))

    assert cpu_calls == gpu_calls, (cpu_calls, gpu_calls)
    assert cpu_calls.count("load_models_gpu") == 1, cpu_calls
    assert cpu_calls.count("unload_all_models") == 1, cpu_calls
    assert cpu_calls.count("soft_empty_cache") == 1, cpu_calls
    # both nets were placed on cpu (the placement) and moved back (the
    # offload) - exactly the two moves, no device branch
    for moved_to in cpu_placements:
        assert moved_to == [torch.device("cpu"), torch.device("cpu")], moved_to


# ---------------------------------------------------------------------------
# the node around the pack (mock-integration)
# ---------------------------------------------------------------------------
def _patch_the_node(monkeypatch, node_module, processor_cls):
    """Stub the heavy parts of the single image node."""
    monkeypatch.setattr(node_module, "KEEPFaceProcessor", processor_cls)
    monkeypatch.setattr(node_module, "comfy_image_to_cv2", lambda image: image)
    monkeypatch.setattr(node_module, "cv2_to_comfy_image", lambda image: image)
    monkeypatch.setattr(node_module, "report_progress", lambda *args: True)


class _StubProcessor:
    """A processor that returns a stub image."""

    def __init__(self, pack) -> None:
        self.pack = pack

    def process_image(self, **kwargs):
        return _StubModule("processed")


class _FailingProcessor:
    """A processor that fails (the node must give the VRAM back)."""

    def __init__(self, pack) -> None:
        self.pack = pack

    def process_image(self, **kwargs):
        raise RuntimeError("the processing failed")


@pytest.mark.mock_integration
def test_one_execute_asks_for_the_vram_once_and_gives_it_back(
    node_module, loader_module, model_management, monkeypatch
):
    """One execute -> one ``load_models_gpu`` with all the pack patchers and
    the release afterwards."""
    pack = loader_module.KEEPModelPack(
        _StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
    )
    _patch_the_node(monkeypatch, node_module, _StubProcessor)

    node_module.KEEP_FaceUpscaleImageNode.execute(
        image=torch.zeros((1, 8, 8, 3)),
        keep_model=pack,
        final_upscale_factor=1.0,
        has_aligned_face=False,
        only_center_face=True,
        draw_bounding_box=False,
    )

    names = [call[0] for call in model_management.calls]
    assert names.count("load_models_gpu") == 1, model_management.calls
    assert names.count("unload_all_models") == 1, model_management.calls
    assert names.count("soft_empty_cache") == 1, model_management.calls
    load_call = [c for c in model_management.calls if c[0] == "load_models_gpu"][0]
    assert list(load_call[1]) == list(pack.patchers)


@pytest.mark.mock_integration
def test_the_node_gives_the_vram_back_when_the_processing_fails(
    node_module, loader_module, model_management, monkeypatch
):
    """The release happens in the ``finally`` path - also on a failure."""
    pack = loader_module.KEEPModelPack(
        _StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
    )
    _patch_the_node(monkeypatch, node_module, _FailingProcessor)

    with pytest.raises(RuntimeError):
        node_module.KEEP_FaceUpscaleImageNode.execute(
            image=torch.zeros((1, 8, 8, 3)),
            keep_model=pack,
            final_upscale_factor=1.0,
            has_aligned_face=False,
            only_center_face=True,
            draw_bounding_box=False,
        )

    names = [call[0] for call in model_management.calls]
    assert names.count("load_models_gpu") == 1, model_management.calls
    assert names.count("unload_all_models") == 1, model_management.calls
    assert names.count("soft_empty_cache") == 1, model_management.calls


# ---------------------------------------------------------------------------
# ensure_flow
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_ensure_flow_is_idempotent(loader_module, stub_models):
    """``ensure_flow`` resolves once per pack - the second call is a no-op.

    A checkpoint that carries the flownet weights (the shipped one, P0.4
    section 2) needs no resolution at all.
    """
    net, helper = stub_models
    pack = loader_module.KEEPModelPack(
        net, helper, None, None, "KEEP", flownet_from_checkpoint=True
    )
    assert pack.ensure_flow() is True, "the flownet comes with the KEEP net"
    # the flownet is gone -> the cached result is still returned
    pack.flownet = None
    assert pack.ensure_flow() is True


@pytest.mark.unit
def test_ensure_flow_reports_a_missing_flownet(loader_module):
    """A pack without a flownet reports that nothing is available (yet).

    P5.5 resolves the GMFlow checkpoint behind this hook.
    """
    net = _StubKeepNet()
    net.flownet = None
    pack = loader_module.KEEPModelPack(net, _StubFaceHelper(), None, None, "KEEP")
    assert pack.ensure_flow() is False


# ---------------------------------------------------------------------------
# P5.3 - the fail-fast weight validation
# ---------------------------------------------------------------------------
PAYLOAD = b"not a real checkpoint"
PAYLOAD_SHA256 = hashlib.sha256(PAYLOAD).hexdigest()


@pytest.fixture
def synthetic_manifest(loader_module, configs_module, monkeypatch):
    """The manifest with the hash of the test payload.

    The real hashes cannot be reproduced by a test, so the manifest the
    loader and the downloader see is replaced by entries that carry the hash
    of the payload the transport double writes.  Everything else - the
    resolution order, the fail-fast rule, the error messages - is the
    production code.
    """
    entries = tuple(
        configs_module.WeightEntry(
            name=entry.name,
            url=entry.url,
            sha256=PAYLOAD_SHA256,
            dest_folder=entry.dest_folder,
            subdir=entry.subdir,
            required=entry.required,
            available=entry.available,
            unverified=entry.unverified,
            sha256_prefix=entry.sha256_prefix,
        )
        for entry in configs_module.WEIGHT_MANIFEST
    )
    by_name = {entry.name: entry for entry in entries}
    monkeypatch.setattr(loader_module, "WEIGHT_MANIFEST_BY_NAME", by_name)
    monkeypatch.setattr(
        loader_module,
        "weight_entry_for",
        lambda model_type: by_name[
            "KEEP_Asian-4765ebe0.pth" if model_type == "Asian" else "KEEP-b76feb75.pth"
        ],
    )
    monkeypatch.setattr(
        loader_module.keep_download, "WEIGHT_MANIFEST", entries
    )
    return by_name


@pytest.fixture
def transport(monkeypatch):
    """The transport double: it records the URLs and writes the payload."""
    download_module = importlib.import_module("ComfyUI_KEEP.src.keep_download")
    calls: list[str] = []

    def _download(url, dst, hash_prefix=None, progress=True):
        calls.append(url)
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        with open(dst, "wb") as handle:
            handle.write(PAYLOAD)

    monkeypatch.setattr(download_module, "download_url_to_file", _download)
    return calls


@pytest.fixture
def loader(loader_module, monkeypatch, synthetic_manifest, transport):
    """A loader whose model build is stubbed out (the weights are real)."""

    class _StubArch:
        """A stand-in for the vendored arch that behaves like ``nn.Module``.

        The real arch **owns** the flownet as a submodule
        (``keep_arch.py:882``), so the state dict of the net always carries
        the ``flownet.*`` keys - and a checkpoint without them fails the
        **strict** load.  The stub keeps that behaviour (a no-op recorder
        would hide exactly that defect).
        """

        #: The state dict of the stub net (the flownet is part of it).
        STATE = {
            "weight": torch.zeros(1),
            "flownet.model.conv1.weight": torch.ones(1),
        }

        def __init__(self, **kwargs):
            self.kwargs = kwargs
            # the real arch owns the flownet as a submodule (keep_arch.py:882)
            self.flownet = _StubModule("flownet")
            self.state = dict(self.STATE)

        def to(self, device):
            return self

        def state_dict(self):
            return dict(self.state)

        def load_state_dict(self, state_dict, strict=True):
            missing = [key for key in self.state if key not in state_dict]
            unexpected = [key for key in state_dict if key not in self.state]
            if strict and (missing or unexpected):
                raise RuntimeError(
                    "Error(s) in loading state_dict for KEEP: missing "
                    f"{missing}, unexpected {unexpected}"
                )
            self.loaded = state_dict
            return self

        def eval(self):
            return self

    class _StubHelper:
        def __init__(self, **kwargs):
            self.kwargs = kwargs

    monkeypatch.setattr(loader_module, "ARCH_REGISTRY", {"KEEP": _StubArch})
    monkeypatch.setattr(
        loader_module.torch,
        "load",
        lambda *args, **kwargs: {"params_ema": {"weight": torch.zeros(1)}},
    )
    monkeypatch.setattr(loader_module, "_face_restore_helper", lambda: _StubHelper)
    return loader_module.KEEPModelLoader()


@pytest.mark.unit
def test_a_missing_required_weight_triggers_a_download(loader, transport):
    """The weights are resolved (and downloaded) before the net is built."""
    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    assert pack.keep_net is not None
    urls = [url.rsplit("/", 1)[-1] for url in transport]
    assert "KEEP-b76feb75.pth" in urls, urls
    # D5: only the **selected** detector and the face parsing net are
    # fetched - the other detectors of the manifest are never downloaded.
    assert "detection_Resnet50_Final.pth" in urls, urls
    assert "parsing_parsenet.pth" in urls, urls


@pytest.mark.unit
def test_a_failed_download_raises_the_actionable_message(
    loader_module, loader, monkeypatch, transport
):
    """A missing required weight names the file, the path and the URL."""
    download_module = importlib.import_module("ComfyUI_KEEP.src.keep_download")

    def _fail(url, dst, hash_prefix=None, progress=True):
        raise RuntimeError("the server is not reachable")

    monkeypatch.setattr(download_module, "download_url_to_file", _fail)

    with pytest.raises(FileNotFoundError) as excinfo:
        loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    message = str(excinfo.value)
    assert "KEEP-b76feb75.pth" in message
    assert "https://github.com/jnjaby/KEEP/releases/download/v1.0.0/KEEP-b76feb75.pth" in message
    assert "keep_models" in message and "KEEP" in message


@pytest.mark.unit
def test_an_unresolvable_optional_gmflow_weight_does_not_raise(
    loader_module, loader, monkeypatch
):
    """The GMFlow checkpoint is optional - its absence never breaks the pack.

    The flow weights are only *needed* when the checkpoint does not carry
    them; the resolution is deferred to ``ensure_flow()`` (P5.5).
    """
    monkeypatch.setattr(
        loader_module.keep_download,
        "ensure_optional",
        lambda entry, progress=True: None,
    )

    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    assert pack is not None
    assert pack.keep_net.kwargs["flownet_path"] is None


@pytest.mark.unit
def test_the_unavailable_asian_entry_does_not_break_the_default_pack(loader, transport):
    """The Asian checkpoint is not required - the default pack still builds."""
    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    assert pack is not None
    urls = [url.rsplit("/", 1)[-1] for url in transport]
    assert "KEEP_Asian-4765ebe0.pth" not in urls, urls


@pytest.mark.unit
def test_selecting_the_asian_model_raises_an_actionable_error(loader):
    """The Asian checkpoint is not published upstream (HTTP 404)."""
    with pytest.raises(FileNotFoundError) as excinfo:
        loader.load_keep_model_pack("Asian", "retinaface_resnet50")

    message = str(excinfo.value)
    assert "KEEP_Asian-4765ebe0.pth" in message
    assert "not published upstream" in message
    assert "KEEP_Asian-4765ebe0.pth" in message
    assert "keep_models" in message


# ---------------------------------------------------------------------------
# P5.5 - the flownet wiring
# ---------------------------------------------------------------------------
def _downloaded(transport) -> list[str]:
    return [url.rsplit("/", 1)[-1] for url in transport]


@pytest.mark.unit
def test_the_arch_gets_a_resolved_flownet_path_without_flownet_weights(loader, transport):
    """A checkpoint **without** flownet weights -> the GMFlow file is resolved.

    The resolved path is put into the architecture config **before** the net
    is built, so the vendored arch loads the trained weights instead of
    random ones.
    """
    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    flownet_path = pack.keep_net.kwargs["flownet_path"]
    assert flownet_path is not None, "the GMFlow checkpoint was not resolved"
    assert flownet_path.endswith("gmflow_sintel-0c07dcb3.pth"), flownet_path
    assert "gmflow_sintel-0c07dcb3.pth" in _downloaded(transport)


@pytest.mark.unit
def test_a_checkpoint_without_flownet_weights_loads_strictly(loader):
    """The fallback branch must not break the strict load.

    The real arch **always** owns the flownet (``keep_arch.py:882``), so the
    state dict of the net always carries the ``flownet.*`` keys.  A
    checkpoint without them used to make the strict load raise
    ``RuntimeError: Missing key(s) in state_dict: flownet.*`` - and the
    GMFlow weights the arch had just loaded were thrown away.  The flownet
    tensors of the built net fill the gap, so the load stays strict **and**
    the trained flow weights survive.
    """
    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    # the strict load happened (the stub arch raises on a missing key)
    assert pack.keep_net.loaded is not None
    # the flownet of the built net (the trained GMFlow weights) is in the
    # loaded state dict - it was not discarded
    assert torch.equal(
        pack.keep_net.loaded["flownet.model.conv1.weight"],
        pack.keep_net.state_dict()["flownet.model.conv1.weight"],
    )


@pytest.mark.unit
def test_the_strict_load_still_rejects_an_unexpected_key(loader_module, loader, monkeypatch):
    """The flownet merge does not weaken the strict load."""
    monkeypatch.setattr(
        loader_module.torch,
        "load",
        lambda *args, **kwargs: {
            "params_ema": {"weight": torch.zeros(1), "unexpected.key": torch.zeros(1)}
        },
    )

    with pytest.raises(RuntimeError):
        loader.load_keep_model_pack("KEEP", "retinaface_resnet50")


@pytest.mark.unit
def test_a_checkpoint_with_flownet_weights_downloads_no_gmflow(
    loader_module, loader, transport, monkeypatch
):
    """The shipped checkpoint carries the flownet -> nothing is downloaded."""
    monkeypatch.setattr(
        loader_module.torch,
        "load",
        lambda *args, **kwargs: {
            "params_ema": {
                "flownet.model.conv1.weight": torch.zeros(1),
                "weight": torch.zeros(1),
            }
        },
    )

    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    assert pack.keep_net.kwargs["flownet_path"] is None
    assert "gmflow_sintel-0c07dcb3.pth" not in _downloaded(transport)
    # the flownet is restored from the checkpoint (the strict load)
    assert "flownet.model.conv1.weight" in pack.keep_net.loaded
    assert pack.flownet_from_checkpoint is True
    # the hook is a no-op - the weights are already there
    assert pack.ensure_flow() is True
    assert "gmflow_sintel-0c07dcb3.pth" not in _downloaded(transport)


class _RecordingInner:
    """The inner module of the flownet (it records the strict load)."""

    def __init__(self) -> None:
        self.loaded: tuple | None = None

    def load_state_dict(self, weights, strict=True):
        self.loaded = (weights, strict)
        return self


class _RecordingFlownet:
    """A flownet that records the weights it receives."""

    def __init__(self) -> None:
        self.model = _RecordingInner()


@pytest.mark.unit
def test_ensure_flow_loads_the_gmflow_weights_strictly(
    loader_module, stub_models, monkeypatch, tmp_path
):
    """A net without flow weights gets them from the GMFlow file."""
    net, helper = stub_models
    net.flownet = _RecordingFlownet()
    pack = loader_module.KEEPModelPack(
        net, helper, None, None, "KEEP",
        flownet_from_checkpoint=False, flownet_path=None,
    )

    gmflow_path = tmp_path / "gmflow_sintel-0c07dcb3.pth"
    gmflow_path.write_bytes(b"weights")
    monkeypatch.setattr(
        loader_module.keep_download,
        "ensure_optional",
        lambda entry, progress=True: str(gmflow_path),
    )
    monkeypatch.setattr(
        loader_module.torch,
        "load",
        lambda *args, **kwargs: {"model": {"w": torch.zeros(1)}},
    )

    assert pack.ensure_flow() is True
    weights, strict = net.flownet.model.loaded
    assert strict is True, "the flow weights must be loaded with strict=True"
    assert "w" in weights
    # the resolution is cached - a second call does nothing
    assert pack.ensure_flow() is True


@pytest.mark.mock_integration
def test_the_image_node_does_not_call_ensure_flow(
    loader_module, node_module, monkeypatch
):
    """The image path leaves the flow resolution deferred."""
    net = _StubKeepNet()
    pack = loader_module.KEEPModelPack(
        net, _StubFaceHelper(), None, None, "KEEP"
    )
    calls: list[str] = []
    monkeypatch.setattr(pack, "ensure_flow", lambda: calls.append("ensure_flow"))
    _patch_the_node(monkeypatch, node_module, _StubProcessor)

    node_module.KEEP_FaceUpscaleImageNode.execute(
        image=torch.zeros((1, 8, 8, 3)),
        keep_model=pack,
        final_upscale_factor=1.0,
        has_aligned_face=False,
        only_center_face=True,
        draw_bounding_box=False,
    )

    assert calls == [], "the image path must not resolve the flow weights"