# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""CPU-only unit tests for the diffusion auto-policy decision layer.

Covers the per-family footprint estimator (bf16-resident component sizes x per-scheme
factors, transient vs steady, base-repo overrides), the dense-quant candidate resolution
(with the quant selector / prequant probe monkeypatched, no torch), and the resolved
provenance record. The loader-side ordering fix is exercised through the planner: the
regression case is a GGUF whose file-size plan forces offload while the candidate's
estimate fits resident."""

from __future__ import annotations

from types import SimpleNamespace

import core.inference.diffusion_auto_policy as ap
from core.inference.diffusion_auto_policy import (
    DenseQuantEstimate,
    build_resolved_record,
    estimate_dense_quant,
    family_bf16_components_gb,
    resolve_dense_quant_candidate,
)
from core.inference.diffusion_memory import (
    OFFLOAD_NONE,
    DeviceMemory,
    plan_diffusion_memory,
)


def _fam(name = "z-image"):
    return SimpleNamespace(name = name)


# ── the per-family table ──────────────────────────────────────────────────────
def test_family_table_covers_the_dit_families():
    for name in ("flux.1", "flux.2-klein", "flux.2-dev", "qwen-image", "z-image", "krea-2"):
        comps = family_bf16_components_gb(_fam(name))
        assert comps is not None, f"{name} missing from the bf16 component table"
        transformer, text_encoders, vae = comps
        assert transformer > 1.0 and text_encoders > 0.0 and vae > 0.0


def test_family_table_unknown_family_returns_none():
    assert family_bf16_components_gb(_fam("not-a-family")) is None


def test_base_repo_override_wins_over_the_family_default():
    # flux.2-klein's family default is the 4B base; loading the 9B GGUF passes the 9B base repo,
    # whose transformer is more than twice the size.
    default = family_bf16_components_gb(_fam("flux.2-klein"))
    nine_b = family_bf16_components_gb(
        _fam("flux.2-klein"), base_repo = "black-forest-labs/FLUX.2-klein-9B"
    )
    assert nine_b is not None and default is not None
    assert nine_b[0] > 2 * default[0]


# ── the estimator ─────────────────────────────────────────────────────────────
def test_estimate_int8_steady_is_roughly_half_bf16():
    est = estimate_dense_quant(_fam("z-image"), "int8")
    assert est is not None
    bf16_mib = 12.3 * ap._MIB_PER_GB
    assert 0.5 * bf16_mib < est.steady_transformer_mib < 0.6 * bf16_mib
    # On-the-fly quantisation transiently materialises the dense bf16 transformer.
    assert est.transient_transformer_mib == int(bf16_mib)
    assert est.prequant is False


def test_estimate_prequant_transient_equals_steady():
    # A pre-quantized checkpoint loads via the meta device: dense bf16 never lands on the GPU, so the
    # build peak IS the quantised size.
    est = estimate_dense_quant(_fam("z-image"), "int8", prequant_available = True)
    assert est is not None
    assert est.transient_transformer_mib == est.steady_transformer_mib
    assert est.prequant is True


def test_estimate_nvfp4_is_smaller_than_int8():
    int8 = estimate_dense_quant(_fam("flux.1"), "int8")
    nvfp4 = estimate_dense_quant(_fam("flux.1"), "nvfp4")
    assert int8 is not None and nvfp4 is not None
    assert nvfp4.steady_transformer_mib < int8.steady_transformer_mib


def test_estimate_unknown_family_or_scheme_returns_none():
    assert estimate_dense_quant(_fam("not-a-family"), "int8") is None
    assert estimate_dense_quant(_fam("z-image"), "q4_k") is None


# ── candidate resolution (selector + prequant probe stubbed) ─────────────────
def _patch_selector(
    monkeypatch,
    *,
    supported = True,
    scheme = "int8",
    prequant = None,
):
    import core.inference.diffusion_transformer_quant as tq

    monkeypatch.setattr(tq, "dense_transformer_supported", lambda target: supported)
    monkeypatch.setattr(
        tq, "select_transformer_quant_scheme", lambda target, req, family = None: scheme
    )
    import core.inference.diffusion_prequant as pq

    monkeypatch.setattr(
        pq,
        "resolve_prequant_source",
        lambda fam, s, path_override = None, base_repo = None: prequant,
    )
    # Neutralize the cache-disk gate by default so resolution tests don't depend on the runner's free
    # space. The two disk-gate tests re-patch this to exercise it explicitly.
    monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: None)


def test_candidate_resolves_for_a_supported_request(monkeypatch):
    _patch_selector(monkeypatch, scheme = "int8")
    est = resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto")
    assert isinstance(est, DenseQuantEstimate)
    assert est.scheme == "int8"
    assert est.transient_transformer_mib > est.steady_transformer_mib


def test_candidate_none_when_request_is_off(monkeypatch):
    _patch_selector(monkeypatch)
    for off in (None, "", "none", "off"):
        assert resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = off) is None


def test_candidate_none_when_device_unsupported(monkeypatch):
    _patch_selector(monkeypatch, supported = False)
    assert resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = "auto") is None


def test_candidate_none_when_no_scheme_resolves(monkeypatch):
    _patch_selector(monkeypatch, scheme = None)
    assert resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = "auto") is None


def test_candidate_disk_gate_skips_when_cache_disk_low(monkeypatch):
    # The dense artifact may be a multi-GB download; a nearly-full model-cache disk drops the
    # candidate (the loader then keeps the GGUF build).
    import core.inference.diffusion_auto_policy as ap

    _patch_selector(monkeypatch, scheme = "int8")
    monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 1024)
    assert (
        resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto")
        is None
    )


def test_candidate_disk_gate_unprobeable_disk_passes(monkeypatch):
    # Disk probing must never sink the candidate: unprobeable (None) passes through.
    import core.inference.diffusion_auto_policy as ap

    _patch_selector(monkeypatch, scheme = "int8")
    monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: None)
    est = resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto")
    assert isinstance(est, DenseQuantEstimate)


def test_disk_gate_sizes_fp32_families_by_their_real_download(monkeypatch):
    # The size table is bf16-RESIDENT, but the disk gate is about bytes landing in the HF cache.
    # Z-Image publishes fp32 shards (measured 23,479 MiB against 11,730 MiB resident), so gating on
    # the resident figure left a window where the check passed and the download then filled the
    # disk. The gate must ask for the download size; the memory estimate must NOT change.
    import core.inference.diffusion_auto_policy as ap

    _patch_selector(monkeypatch, scheme = "int8")
    est = ap.estimate_dense_quant(_fam("z-image"), "int8")
    assert est.transient_transformer_mib == 11_730  # resident: unchanged
    assert est.download_transformer_mib == 23_460  # download: measured 23,479 MiB
    # Free space that clears the old (resident-based) bar but not the real download is refused.
    monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 11_730 + 10 * 1024 + 512)
    assert (
        resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto")
        is None
    )
    monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 23_460 + 10 * 1024 + 512)
    assert isinstance(
        resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto"),
        DenseQuantEstimate,
    )


def test_disk_gate_does_not_overcharge_a_family_published_below_bf16(monkeypatch):
    # The correction runs both ways: Ideogram 4 ships fp8 (transformer + unconditional_transformer
    # measured at 17,718 MiB) and doubles on the way to bf16, so charging the resident 35,477 MiB
    # of disk would refuse the candidate on a disk that easily holds it.
    import core.inference.diffusion_auto_policy as ap

    _patch_selector(monkeypatch, scheme = "int8")
    est = ap.estimate_dense_quant(_fam("ideogram-4"), "int8")
    assert est.transient_transformer_mib == 35_476
    assert est.download_transformer_mib == 17_738
    monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 17_738 + 10 * 1024 + 512)
    assert isinstance(
        resolve_dense_quant_candidate(fam = _fam("ideogram-4"), target = object(), requested = "auto"),
        DenseQuantEstimate,
    )


def test_disk_gate_matches_download_for_bf16_published_families(monkeypatch):
    # Families that publish bf16 download what they occupy (measured 0.99-1.07x), so the two
    # numbers stay equal and no factor entry is needed.
    import core.inference.diffusion_auto_policy as ap
    for name in ("flux.1", "flux.2-dev", "qwen-image", "krea-2", "hidream-i1"):
        est = ap.estimate_dense_quant(_fam(name), "int8")
        assert est.download_transformer_mib == est.transient_transformer_mib, name


def test_candidate_none_for_an_unlisted_family(monkeypatch):
    # No size entry means no basis to re-plan; the loader keeps today's resident-only gate.
    _patch_selector(monkeypatch)
    assert (
        resolve_dense_quant_candidate(fam = _fam("not-a-family"), target = object(), requested = "auto")
        is None
    )


def test_candidate_uses_prequant_transient_when_available(monkeypatch):
    # A hosted-repo prequant source is available without a local-path check.
    _patch_selector(monkeypatch, prequant = SimpleNamespace(kind = "repo", location = "org/int8"))
    est = resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "int8")
    assert est is not None and est.prequant is True
    assert est.transient_transformer_mib == est.steady_transformer_mib


# ── the ordering-fix regression, at the planner level ─────────────────────────
def _cuda_target():
    return SimpleNamespace(device = "cuda", supports_model_cpu_offload = True)


def test_quant_candidate_fits_resident_where_gguf_plan_offloads():
    # The ordering fix on a 32 GiB consumer card: the user picked a LARGE (BF16) GGUF, so the
    # file-size plan forces offload -- but the dense-quant candidate is far smaller, and re-planning
    # against it keeps everything resident. Before the fix the loader never attempted the fast path.
    memory = DeviceMemory("cuda", "cuda", "discrete_vram", 30000, 32768)
    z_bf16_gguf_mib = int(12.3 * ap._MIB_PER_GB * 1.05)  # BF16 GGUF resident estimate
    companions_mib = 2600  # fp8-quantised text encoders + VAE
    gguf_plan = plan_diffusion_memory(
        target = _cuda_target(),
        device_memory = memory,
        model_dense_mib = z_bf16_gguf_mib + companions_mib,
        companion_dense_mib = companions_mib,
        runtime_headroom_mib = 6963,
    )
    assert gguf_plan.offload_policy != OFFLOAD_NONE

    est = estimate_dense_quant(_fam("z-image"), "int8", prequant_available = True)
    assert est is not None
    assert est.transient_transformer_mib < z_bf16_gguf_mib / 1.8
    quant_plan = plan_diffusion_memory(
        target = _cuda_target(),
        device_memory = memory,
        model_dense_mib = est.transient_transformer_mib + companions_mib,
        companion_dense_mib = companions_mib,
        runtime_headroom_mib = 6963,
    )
    assert quant_plan.offload_policy == OFFLOAD_NONE


# ── the resolved provenance record ────────────────────────────────────────────
def test_resolved_record_marks_auto_and_explicit():
    record = build_resolved_record(
        {
            "speed_mode": (None, "default", "per-kind default"),
            "transformer_quant": ("auto", "fp8", "auto ladder"),
            "attention_backend": ("cudnn", "_native_cudnn", "requested"),
            "memory_mode": ("", "none", "planned"),
            "cpu_offload": (True, True, "legacy flag"),
        }
    )
    assert record["speed_mode"]["source"] == "auto"
    assert record["transformer_quant"]["source"] == "auto"  # "auto" delegates to backend
    assert record["attention_backend"]["source"] == "explicit"
    assert record["memory_mode"]["source"] == "auto"  # blank string delegates
    assert record["cpu_offload"]["source"] == "explicit"
    assert record["transformer_quant"]["value"] == "fp8"
    assert all("reason" in v for v in record.values())