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

"""Tests for the native sd.cpp diffusion backend (the no-GPU engine)."""

from __future__ import annotations

import threading
import types

import pytest
from PIL import Image

from core.inference import sd_cpp_backend as bk
from core.inference.diffusion_families import detect_family
from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles
from core.inference.sd_cpp_backend import (
    SdCppDiffusionBackend,
    _map_guidance,
    ensure_sd_cpp_binary,
)
from core.inference.sd_cpp_engine import SdCppCancelled


class _FakeEngine:
    """Stands in for SdCppEngine: writes a 1x1 PNG and records the args."""

    def __init__(
        self,
        *,
        fail = None,
        cancel_on_call = False,
    ):
        self.calls = []
        self.fail = fail
        self.cancel_on_call = cancel_on_call

    def is_available(self):
        return True

    def version(self, **_):
        return "fake sd-cli"

    def generate(
        self,
        files,
        params,
        *,
        output_path,
        cancel_event = None,
        **kw,
    ):
        self.calls.append((files, params, output_path, kw))
        if self.cancel_on_call and cancel_event is not None:
            cancel_event.set()
        if self.fail is not None:
            raise self.fail
        if cancel_event is not None and cancel_event.is_set():
            raise SdCppCancelled("cancelled")
        Image.new("RGB", (1, 1), (10, 20, 30)).save(output_path)
        from pathlib import Path

        return Path(output_path)


def _loaded_backend(fam_name = "z-image", engine = None):
    b = SdCppDiffusionBackend(engine = engine or _FakeEngine())
    fam = detect_family(fam_name)
    b._state = bk._SdState(
        repo_id = "unsloth/Z-Image-Turbo-GGUF",
        base_repo = fam.base_repo,
        family = fam,
        device = "cpu",
        files = SdCppModelFiles(
            diffusion_model = "/m/z.gguf", vae = "/m/vae.safetensors", llm = "/m/llm.safetensors"
        ),
        vae_format = fam.sd_cpp_vae_format,
        sampling_method = fam.sd_cpp_sampling_method,
        flow_shift = fam.sd_cpp_flow_shift,
        mode = "oneshot",  # this fixture injects an engine, so it exercises the one-shot path
    )
    return b


def test_loaded_repo_ids_includes_native_companions():
    # The one-shot native engine re-reads its companion VAE / text-encoder files from the HF cache on every generation, so the delete-cached guard queries loaded_repo_ids() to refuse deleting an in-use companion repo.
    # It must surface the committed family VAE + text-encoder repos (plus the main + base repos), not just the loaded GGUF, and be empty once unloaded.
    b = _loaded_backend("flux.1")
    ids = set(b.loaded_repo_ids())
    fam = detect_family("flux.1")
    assert "unsloth/Z-Image-Turbo-GGUF" in ids  # the main GGUF repo
    assert fam.base_repo in ids
    assert fam.sd_cpp_vae[0] in ids  # black-forest-labs/FLUX.1-schnell (VAE)
    for terepo, _f, _k in fam.sd_cpp_text_encoders:
        assert terepo in ids  # comfyanonymous/flux_text_encoders
    b._state = None
    assert b.loaded_repo_ids() == ()


def test_loaded_repo_ids_tracks_variant_encoder_by_gguf_filename():
    # A local *klein-9B*.gguf carries the variant keyword only in the basename, so loaded_repo_ids() must include the GGUF filename in the load identity; otherwise it falls back to the 4B default encoder and the cache-deletion guard protects the wrong repo.
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    fam = detect_family("flux.2-klein")
    b._state = bk._SdState(
        repo_id = "local/my-klein-checkpoints",  # no variant keyword; it lives in the filename
        base_repo = fam.base_repo,
        family = fam,
        device = "cpu",
        files = SdCppModelFiles(diffusion_model = "/m/FLUX.2-klein-9B-Q4_K_M.gguf"),
        vae_format = fam.sd_cpp_vae_format,
        sampling_method = fam.sd_cpp_sampling_method,
        flow_shift = fam.sd_cpp_flow_shift,
        mode = "oneshot",
        gguf_filename = "FLUX.2-klein-9B-Q4_K_M.gguf",
    )
    ids = set(b.loaded_repo_ids())
    assert "Comfy-Org/vae-text-encorder-for-flux-klein-9b" in ids  # the 8B encoder this load pulled
    assert "Comfy-Org/z_image_turbo" not in ids  # the 4B default must not be protected instead


class _FakeServer:
    """Stands in for SdCppServer: records the spawn + one img_gen per whole batch."""

    def __init__(self, binary):
        self.binary = binary
        self.started = None
        self.stopped = False
        self.payloads = []
        self.timeouts = []
        self.alive = True
        self.lora_dir = None  # set by a test to the server's --lora-model-dir scratch dir
        # Raised by the next img_gen (one shot) so a test can stage a mid-generation server death.
        self.img_gen_error = None

    def is_alive(self):
        return self.alive and not self.stopped

    def start(
        self,
        files,
        *,
        vae_format = None,
        offload = None,
        native_speed = None,
        threads = None,
        extra_args = None,
    ):
        self.started = dict(
            files = files,
            vae_format = vae_format,
            offload = offload,
            native_speed = native_speed,
            threads = threads,
            extra_args = list(extra_args or []),
        )

    def img_gen(
        self,
        payload,
        *,
        on_step = None,
        cancel_event = None,
        total_timeout = None,
    ):
        import io as _io

        self.payloads.append(payload)
        self.timeouts.append(total_timeout)
        if self.img_gen_error is not None:
            err, self.img_gen_error = self.img_gen_error, None
            self.alive = False  # a ggml abort takes the process down
            raise err
        if on_step is not None:
            steps = payload.get("sample_params", {}).get("sample_steps", 0)
            on_step(f"  {steps}/{steps}")
        n = int(payload.get("batch_count", 1))
        blobs = []
        for i in range(n):
            buf = _io.BytesIO()
            Image.new("RGB", (1, 1), (i, i, i)).save(buf, format = "PNG")
            blobs.append(buf.getvalue())
        return blobs

    def stop(self):
        self.stopped = True


# ── asset resolution ──────────────────────────────────────────────────────────


@pytest.mark.parametrize(
    "fam_name,expect_kinds",
    [
        ("flux.1", {"diffusion_model", "vae", "clip_l", "t5xxl"}),
        ("z-image", {"diffusion_model", "vae", "llm"}),
        ("qwen-image", {"diffusion_model", "vae", "qwen2vl"}),
        ("flux.2-klein", {"diffusion_model", "vae", "llm"}),
    ],
)
def test_asset_specs_cover_required_files(fam_name, expect_kinds):
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    fam = detect_family(fam_name)
    specs = b._asset_specs("unsloth/x-GGUF", "x-Q4_K_M.gguf", fam)
    kinds = {kind for _, _, kind in specs}
    assert kinds == expect_kinds
    # Every spec has a non-empty repo + filename.
    assert all(repo and fn for repo, fn, _ in specs)
    # The transformer reuses the requested GGUF, not a registry file.
    tr = [s for s in specs if s[2] == "diffusion_model"][0]
    assert tr[0] == "unsloth/x-GGUF" and tr[1] == "x-Q4_K_M.gguf"


def test_download_plan_stages_exactly_what_sd_cli_opens(monkeypatch):
    # The plan feeds the Hub download manager. Native reads the single-file VAE + text encoder and never opens the base repo sharded components, so a native-routed pick must be staged from the asset specs, not from the diffusers plan.
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    sizes = {
        ("unsloth/Z-Image-Turbo-GGUF", "z-image-turbo-Q4_K_M.gguf"): 4_000,
        ("Comfy-Org/z_image_turbo", "split_files/vae/ae.safetensors"): 300,
        ("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors"): 8_000,
    }
    monkeypatch.setattr(
        SdCppDiffusionBackend,
        "_plan_file_sizes",
        staticmethod(lambda by_repo, token: sizes),
    )

    plan = b.download_plan(
        "unsloth/Z-Image-Turbo-GGUF",
        gguf_filename = "z-image-turbo-Q4_K_M.gguf",
        model_kind = "gguf",
        # diffusers-only knobs are accepted and ignored, exactly as begin_load accepts them.
        transformer_quant = "int8",
        memory_mode = "low_vram",
    )

    fam = detect_family("z-image")
    expected = {
        (r, f)
        for r, f, _k in b._asset_specs(
            "unsloth/Z-Image-Turbo-GGUF", "z-image-turbo-Q4_K_M.gguf", fam
        )
    }
    listed = {(e["repo_id"], f) for e in plan["entries"] for f in e["files"]}
    assert listed == expected
    assert plan["total_bytes"] == 12_300
    # The transformer entry is the only one that carries the GGUF filename (the manager treats a GGUF pick differently from a plain file), and the VAE + encoder share one repo entry.
    tr = [e for e in plan["entries"] if e["gguf_filename"]]
    assert len(tr) == 1 and tr[0]["repo_id"] == "unsloth/Z-Image-Turbo-GGUF"
    assert len([e for e in plan["entries"] if e["repo_id"] == "Comfy-Org/z_image_turbo"]) == 1


def test_download_plan_skips_a_local_transformer_but_still_stages_the_assets(monkeypatch, tmp_path):
    # A local GGUF folder is already on disk; its VAE + encoder still have to come from the Hub.
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    (tmp_path / "z-image-turbo-Q4_K_M.gguf").write_bytes(b"gguf")
    monkeypatch.setattr(
        SdCppDiffusionBackend, "_plan_file_sizes", staticmethod(lambda by_repo, token: {})
    )

    plan = b.download_plan(
        str(tmp_path), gguf_filename = "z-image-turbo-Q4_K_M.gguf", model_kind = "gguf"
    )

    assert str(tmp_path) not in {e["repo_id"] for e in plan["entries"]}
    assert {e["repo_id"] for e in plan["entries"]} == {"Comfy-Org/z_image_turbo"}
    # An unreadable size understates the total; it must never fail the plan.
    assert plan["total_bytes"] == 0


def test_download_plan_refuses_a_pick_native_cannot_serve():
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    with pytest.raises(ValueError, match = "gguf_filename is required"):
        b.download_plan("unsloth/Z-Image-Turbo-GGUF", model_kind = "gguf")
    with pytest.raises(ValueError, match = "native sd.cpp asset mapping"):
        b.download_plan(
            "stabilityai/sdxl-turbo", gguf_filename = "sdxl-Q4_K_M.gguf", model_kind = "gguf"
        )


def test_asset_specs_flux2_klein_selects_encoder_by_variant():
    # FLUX.2-klein 4B pairs with Qwen3-4B, 9B with Qwen3-8B, so the encoder must be chosen from the load identity, not the family default (a mismatched encoder fails deep in sd-cli).
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    fam = detect_family("flux.2-klein")

    specs_4b = b._asset_specs("unsloth/FLUX.2-klein-4B-GGUF", "FLUX.2-klein-4B-Q4_K_M.gguf", fam)
    te_4b = [(r, f) for r, f, k in specs_4b if k == "llm"]
    assert te_4b == [("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors")]

    specs_9b = b._asset_specs("unsloth/FLUX.2-klein-9B-GGUF", "FLUX.2-klein-9B-Q4_K_M.gguf", fam)
    te_9b = [(r, f) for r, f, k in specs_9b if k == "llm"]
    assert te_9b == [
        (
            "Comfy-Org/vae-text-encorder-for-flux-klein-9b",
            "split_files/text_encoders/qwen_3_8b.safetensors",
        )
    ]


# ── guidance mapping ──────────────────────────────────────────────────────────


def test_map_guidance_flux_uses_distilled_guidance():
    cfg, g = _map_guidance(detect_family("flux.1"), 3.5)
    assert cfg is None and g == 3.5


def test_map_guidance_cfg_family_off_when_distilled():
    # qwen-image uses real CFG; a distilled 0 -> CFG off (1.0), a >1 value passes through.
    assert _map_guidance(detect_family("qwen-image"), 0.0) == (1.0, None)
    assert _map_guidance(detect_family("qwen-image"), 4.0) == (4.0, None)


# ── status ────────────────────────────────────────────────────────────────────


def test_status_unloaded_reports_sd_cpp_engine():
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    st = b.status()
    assert st["loaded"] is False and st["engine"] == "sd_cpp"


def test_status_loaded_shape():
    b = _loaded_backend()
    st = b.status()
    assert st["loaded"] is True
    assert st["engine"] == "sd_cpp"
    assert st["family"] == "z-image"
    assert st["device"] == "cpu"
    # diffusers-only fields are present (route response parity) but null.
    for k in ("transformer_quant", "attention_backend", "transformer_cache", "text_encoder_quant"):
        assert st[k] is None


# ── generate ──────────────────────────────────────────────────────────────────


def test_generate_returns_images_and_seed():
    eng = _FakeEngine()
    b = _loaded_backend(engine = eng)
    out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 123, batch_size = 2)
    assert out["seed"] == 123
    assert out["repo_id"] == "unsloth/Z-Image-Turbo-GGUF"
    assert len(out["images"]) == 2
    assert all(isinstance(im, Image.Image) for im in out["images"])
    # One sd-cli run per batch image, each a distinct seed from the base.
    assert len(eng.calls) == 2
    seeds = [params.seed for _, params, _, _ in eng.calls]
    assert seeds == [123, 124]
    # The per-image seeds are returned so the route can persist each one.
    assert out["seeds"] == [123, 124]


def test_generate_qwen_passes_sampling_args():
    eng = _FakeEngine()
    b = _loaded_backend(fam_name = "qwen-image", engine = eng)
    b.generate(prompt = "x", steps = 20, guidance = 4.0, seed = 1)
    _, params, _, kw = eng.calls[0]
    assert params.sampling_method == "euler"  # Qwen's supported sd.cpp sampler
    assert "--flow-shift" in (kw.get("extra_args") or [])


def test_generate_raises_when_not_loaded():
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    with pytest.raises(RuntimeError, match = "No diffusion model is loaded"):
        b.generate(prompt = "x")


def test_generate_passes_vae_format_for_flux2():
    eng = _FakeEngine()
    b = _loaded_backend(fam_name = "flux.2-klein", engine = eng)
    b.generate(prompt = "x", steps = 4, seed = 1)
    _, _, _, kw = eng.calls[0]
    assert kw.get("extra_args") == ["--vae-format", "flux2"]


def test_generate_cancellation_raises_cancelled_not_failure():
    # The engine cancels mid-run; the backend surfaces a cancellation, not a crash.
    eng = _FakeEngine(cancel_on_call = True)
    b = _loaded_backend(engine = eng)
    with pytest.raises(RuntimeError, match = "cancelled"):
        b.generate(prompt = "x", steps = 8, seed = 5)


def test_generate_progress_tracks_parsed_steps():
    b = _loaded_backend()
    b._gen = bk._SdGen(total_steps = 8)
    b._on_log("  sampling 4/8 done")
    p = b.generate_progress()
    assert p["active"] is True and p["step"] == 4 and p["total_steps"] == 8
    # A fraction with a different denominator must not move the bar.
    b._on_log("loaded 1/3 tensors")
    assert b.generate_progress()["step"] == 4


def test_generate_publishes_progress_before_lora_resolution(monkeypatch):
    # LoRA resolution runs during pre-generate setup while _generate_lock is held, so a progress probe in that window must read ACTIVE; _gen is published before it, mirroring the diffusers path.
    from core.inference import diffusion_lora

    eng = _FakeEngine()
    b = _loaded_backend(engine = eng)
    monkeypatch.setattr(diffusion_lora, "supports_lora", lambda **_k: True)

    seen: dict = {}

    def _resolve(
        active,
        *,
        family = None,
        hf_token = None,
        cancel_event = None,
    ):
        # Mid-setup: the in-flight generation must already be reported as active.
        seen["progress"] = b.generate_progress()
        return []

    monkeypatch.setattr(diffusion_lora, "resolve_specs", _resolve)

    out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, loras = [("some/lora", 1.0)])
    assert out["images"]
    assert seen["progress"]["active"] is True
    assert seen["progress"]["total_steps"] == 8


# ── load validation + binary install ──────────────────────────────────────────


def test_begin_load_rejects_unsupported_family(monkeypatch):
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    # A family with no native asset mapping must be rejected (router falls back).
    monkeypatch.setattr(bk, "family_sd_cpp_supported", lambda fam: False)
    with pytest.raises(ValueError, match = "no native sd.cpp asset mapping"):
        b.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z.gguf")


def test_begin_load_requires_gguf_filename():
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    with pytest.raises(ValueError, match = "gguf_filename is required"):
        b.begin_load("unsloth/Z-Image-Turbo-GGUF")


def test_begin_load_resolves_family_from_filename_only(monkeypatch):
    # A local .gguf pick whose family keyword lives only in the basename must resolve via the same filename fallback the route validated with, not dead-end with "Could not infer" on a native host.
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    monkeypatch.setattr(b, "_run_load", lambda **kwargs: None)  # skip the download thread
    b.begin_load("/models/gguf-store", gguf_filename = "Z-Image-Turbo-Q4_K_M.gguf")
    # Validation passed (no ValueError) and the family was inferred from the filename.
    assert b._loading is not None and b._loading.repo_id == "/models/gguf-store"


def test_ensure_binary_returns_found(monkeypatch):
    monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli")
    assert ensure_sd_cpp_binary() == "/usr/bin/sd-cli"


def test_ensure_binary_install_disabled_returns_none(monkeypatch):
    monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: None)
    assert ensure_sd_cpp_binary(allow_install = False) is None


def test_unload_clears_state_and_signals_cancel():
    cancel = threading.Event()
    b = _loaded_backend()
    b._active_generate_cancel = cancel
    st = b.unload()
    assert st["loaded"] is False
    assert cancel.is_set()
    assert b._cancel_event.is_set()


def test_status_reports_offload_when_flags_active():
    # status must reflect the offload flags actually passed to sd-cli, not always "none", so a balanced/low_vram (or cpu_offload) load is verifiable.
    b = _loaded_backend()
    # No flags (CPU default) -> none.
    assert b.status()["offload_policy"] == "none" and b.status()["cpu_offload"] is False
    # Flags present (off-CPU offload) -> reported active.
    s = b._state
    b._state = bk._SdState(
        repo_id = s.repo_id,
        base_repo = s.base_repo,
        family = s.family,
        device = "cuda",
        files = s.files,
        offload_flags = ("--vae-on-cpu", "--clip-on-cpu"),
    )
    st = b.status()
    assert st["cpu_offload"] is True and st["offload_policy"] == "active"


def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch):
    # A generation that started during the asset download is still running against the OLD model. _run_load must cancel it AND wait on _generate_lock before committing the new state, or a stale sd-cli run persists an image from the previous model once the new load reports ready.
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    fam = detect_family("z-image")
    monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
    monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
    monkeypatch.setattr(
        b,
        "_fetch_assets",
        lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"},
    )
    # Avoid importing torch from the worker thread (its first import deadlocks off the main thread, a test artifact); the device only needs to be CPU here.
    monkeypatch.setattr(
        bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
    )

    b._load_token = 5
    cancel = threading.Event()
    b._active_generate_cancel = cancel  # a generation is "in flight"

    committed = threading.Event()

    def _load():
        b._run_load(
            repo_id = "unsloth/Z-Image-Turbo-GGUF",
            gguf_filename = "z.gguf",
            base = fam.base_repo,
            fam = fam,
            hf_token = None,
            _load_token = 5,
        )
        committed.set()

    b._generate_lock.acquire()  # simulate the live denoise holding _generate_lock
    try:
        threading.Thread(target = _load, daemon = True).start()
        # The commit must block behind the live generation and not publish the new state, but must already have signalled the in-flight cancel.
        assert not committed.wait(0.5)
        assert b._state is None
        assert cancel.is_set()
    finally:
        b._generate_lock.release()
    assert committed.wait(5)  # only now does the commit run
    assert b._state is not None and b._state.repo_id == "unsloth/Z-Image-Turbo-GGUF"


# ── persistent sd-server mode ──────────────────────────────────────────────────


def test_resolve_backend_prefers_server(monkeypatch):
    b = SdCppDiffusionBackend()  # no injected engine
    monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
    mode, binary, engine = b._resolve_backend()
    assert mode == "server" and binary == "/x/sd-server" and engine is None


def test_resolve_backend_injected_engine_forces_oneshot():
    b = SdCppDiffusionBackend(engine = _FakeEngine())
    mode, binary, engine = b._resolve_backend()
    assert mode == "oneshot" and binary is None and engine is not None


def test_resolve_backend_falls_back_to_oneshot_without_server(monkeypatch):
    b = SdCppDiffusionBackend()
    monkeypatch.setattr(bk, "find_sd_server_binary", lambda: None)
    monkeypatch.setattr(bk, "_install_allowed", lambda: False)  # don't attempt a real install
    monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli")
    mode, binary, engine = b._resolve_backend()
    assert mode == "oneshot" and engine is not None


def test_resolve_backend_cached_fallback_engine_does_not_pin_oneshot(monkeypatch):
    # A lazily cached fallback engine (NOT an explicit injection) must not force one-shot: once a server is available again, the next load can use it.
    b = SdCppDiffusionBackend()  # no injected engine
    b._engine = _FakeEngine()  # simulate a prior lazy one-shot fallback caching the engine
    monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
    mode, binary, engine = b._resolve_backend()
    assert mode == "server" and binary == "/x/sd-server" and engine is None


def _run_server_load(
    monkeypatch,
    b,
    servers,
    fam_name = "z-image",
    device = "cpu",
):
    fam = detect_family(fam_name)
    monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
    # The fake binary path is not a real executable; skip the up-front runnability probe.
    monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True)

    def _factory(binary):
        s = _FakeServer(binary)
        servers.append(s)
        return s

    monkeypatch.setattr(bk, "SdCppServer", _factory)
    monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
    monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
    monkeypatch.setattr(
        b,
        "_fetch_assets",
        lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"},
    )
    monkeypatch.setattr(
        bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = device)
    )
    b._load_token = 1
    b._run_load(
        repo_id = "unsloth/Z-Image-Turbo-GGUF",
        gguf_filename = "z.gguf",
        base = fam.base_repo,
        fam = fam,
        hf_token = None,
        _load_token = 1,
    )


def test_server_load_spawns_once_and_status_reports_mode(monkeypatch):
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)
    assert len(servers) == 1
    assert servers[0].started is not None  # the model is loaded once, at spawn
    assert b._state is not None and b._state.mode == "server" and b._state.server is servers[0]
    assert b.status()["native_mode"] == "server"


def test_server_generate_uses_one_request_for_whole_batch(monkeypatch):
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)
    out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 7, batch_size = 3)
    assert len(out["images"]) == 3
    assert all(isinstance(im, Image.Image) for im in out["images"])
    # ONE job for the whole batch (no per-image model reload), unlike the one-shot path.
    assert len(servers[0].payloads) == 1
    assert servers[0].payloads[0]["batch_count"] == 3
    assert out["seed"] == 7 and out["seeds"] == [7, 8, 9]
    # step progress was driven from the server's stdout line.
    assert b._gen is None  # cleared after generate


_GGML_ABORT = (
    "sd-server connection lost during img_gen poll (process exited, code -6)\n"
    "Last output:\n"
    "[ERROR] ggml_extend.hpp:70   - ggml_metal_op_encode_impl: error: unsupported op 'MUL_MAT'\n"
    "1   sd-server   0x00000001044f8df4 ggml_abort + 156\n"
    "10  sd-server   0x00000001043f6ce8 StableDiffusionGGML::sample"
)


def test_server_generation_restarts_on_the_cpu_backend_after_a_ggml_abort(monkeypatch):
    # ggml checks every node against the device supports_op and calls GGML_ABORT when one is not implemented, killing sd-server mid-generation with no per-op CPU fallback. Retrying on the same backend would abort identically, so the load is restarted with --backend cpu and the generation completes (slower) instead of failing.
    # generation completes (slower) instead of failing.
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers, device = "mps")
    servers[0].img_gen_error = RuntimeError(_GGML_ABORT)

    out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 4, seed = 3)

    assert len(out["images"]) == 1  # the retry produced the image
    assert len(servers) == 2 and servers[0].stopped is True
    assert servers[1].started["extra_args"] == ["--backend", "cpu"]
    # The same checkpoint and run settings are reused; only the backend placement changed.
    assert servers[1].started["files"] is servers[0].started["files"]
    assert servers[1].started["native_speed"] == servers[0].started["native_speed"]
    # The live state points at the replacement, so the next generation does not touch the dead one.
    assert b._state is not None and b._state.server is servers[1]


def test_cpu_backend_restart_happens_once_per_load(monkeypatch):
    # The restart is a one-shot rescue: if the CPU backend aborts too, the error surfaces rather than spawning servers forever.
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers, device = "mps")
    servers[0].img_gen_error = RuntimeError(_GGML_ABORT)
    b.generate(prompt = "x", steps = 4, seed = 1)
    assert len(servers) == 2

    servers[1].img_gen_error = RuntimeError(_GGML_ABORT)
    with pytest.raises(RuntimeError, match = "unsupported op"):
        b.generate(prompt = "x", steps = 4, seed = 1)
    assert len(servers) == 2  # no third spawn


def test_server_death_without_the_abort_signature_is_not_retried(monkeypatch):
    # An OOM kill, a corrupt checkpoint or a genuine bug must not be silently retried on another backend: only the deterministic unsupported-op abort earns the CPU restart.
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers, device = "mps")
    servers[0].img_gen_error = RuntimeError(
        "sd-server connection lost during img_gen poll (process exited, code -9)"
    )
    with pytest.raises(RuntimeError, match = "code -9"):
        b.generate(prompt = "x", steps = 4, seed = 1)
    assert len(servers) == 1


def test_cpu_device_does_not_restart_on_an_abort(monkeypatch):
    # Already on CPU: the abort is not a backend-placement problem, so restarting would just repeat it. Surface the error instead.
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers, device = "cpu")
    servers[0].img_gen_error = RuntimeError(_GGML_ABORT)
    with pytest.raises(RuntimeError, match = "unsupported op"):
        b.generate(prompt = "x", steps = 4, seed = 1)
    assert len(servers) == 1


def test_server_generate_splits_batches_above_server_limit(monkeypatch):
    # A batch above the server per-job limit is chunked (the one-shot path did these image-by-image); each chunk gets a timeout proportional to its image count.
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)
    out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 100, batch_size = 10)
    assert len(out["images"]) == 10
    counts = [p["batch_count"] for p in servers[0].payloads]
    assert counts == [bk._MAX_SERVER_BATCH, 10 - bk._MAX_SERVER_BATCH]  # [8, 2]
    # Chunks share ONE request deadline rather than each getting a full budget: a batch is split only because the server caps images per job, so per-chunk budgets would let a batch outlive the window the page is still waiting on.
    # Each chunk therefore gets what is left, which is at most the ceiling and never increases.
    assert servers[0].timeouts[0] <= bk.NATIVE_GENERATION_TIMEOUT_S
    assert servers[0].timeouts[1] <= servers[0].timeouts[0]
    # A single slow image can still use the whole window (the old per-image cap was 30 minutes).
    assert servers[0].timeouts[-1] > 1800.0
    # Seeds run contiguously across chunks (chunk 2 submitted at base + 8).
    assert out["seeds"] == list(range(100, 110))
    assert servers[0].payloads[1]["seed"] == 108


def test_server_generate_masks_large_seed(monkeypatch):
    # sd.cpp image seed is signed int64, so a larger explicit seed must be masked before it reaches the server (the request model / diffusers accept up to 2**64 - 1).
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)
    out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 2**64 - 1, batch_size = 1)
    assert servers[0].payloads[0]["seed"] <= (1 << 63) - 1
    assert all(s <= (1 << 63) - 1 for s in out["seeds"])


def test_status_clears_when_server_died(monkeypatch):
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)
    assert b.status()["loaded"] is True
    servers[0].alive = False  # the resident server crashed / was OOM-killed
    st = b.status()
    assert st["loaded"] is False
    assert b._state is None  # stale state was dropped so clients reload


def test_server_generate_progress_from_stdout(monkeypatch):
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)

    seen = {}

    class _WatchServer(_FakeServer):
        def img_gen(
            self,
            payload,
            *,
            on_step = None,
            cancel_event = None,
            total_timeout = None,
        ):
            on_step("  4/8")
            seen["mid"] = b.generate_progress()
            return super().img_gen(
                payload, on_step = on_step, cancel_event = cancel_event, total_timeout = total_timeout
            )

    b._state = bk._SdState(
        repo_id = b._state.repo_id,
        base_repo = b._state.base_repo,
        family = b._state.family,
        device = b._state.device,
        files = b._state.files,
        vae_format = b._state.vae_format,
        sampling_method = b._state.sampling_method,
        flow_shift = b._state.flow_shift,
        server = _WatchServer("/x/sd-server"),
        mode = "server",
    )
    b.generate(prompt = "x", steps = 8, seed = 1)
    assert seen["mid"]["step"] == 4 and seen["mid"]["total_steps"] == 8


def test_server_unload_stops_server(monkeypatch):
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)
    st = b.unload()
    assert st["loaded"] is False
    assert servers[0].stopped is True
    assert b._state is None


def test_server_reload_stops_old_server_before_new(monkeypatch):
    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)
    # A second load must tear down the first server and start a fresh one.
    b._load_token = 2
    fam = detect_family("z-image")
    b._run_load(
        repo_id = "unsloth/Z-Image-Turbo-GGUF",
        gguf_filename = "z.gguf",
        base = fam.base_repo,
        fam = fam,
        hf_token = None,
        _load_token = 2,
    )
    assert len(servers) == 2
    assert servers[0].stopped is True  # old server stopped
    assert b._state.server is servers[1] and servers[1].stopped is False


def test_server_start_failure_falls_back_to_oneshot(monkeypatch):
    # A present-but-broken sd-server must not fail the load when sd-cli works.
    b = SdCppDiffusionBackend()
    monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
    # The probe passes; the failure exercised here is in start(), not the up-front probe.
    monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True)

    class _BadServer:
        def __init__(self, binary):
            self.stopped = False

        def start(self, *a, **k):
            raise RuntimeError("sd-server broken")

        def stop(self):
            self.stopped = True

    monkeypatch.setattr(bk, "SdCppServer", _BadServer)
    fake = _FakeEngine()
    monkeypatch.setattr(b, "_resolve_engine", lambda: fake)
    monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
    monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
    monkeypatch.setattr(
        b,
        "_fetch_assets",
        lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"},
    )
    monkeypatch.setattr(
        bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
    )
    fam = detect_family("z-image")
    b._load_token = 1
    b._run_load(
        repo_id = "unsloth/Z-Image-Turbo-GGUF",
        gguf_filename = "z.gguf",
        base = fam.base_repo,
        fam = fam,
        hf_token = None,
        _load_token = 1,
    )
    assert b._state is not None and b._state.mode == "oneshot" and b._state.server is None
    # and it can still generate via the one-shot engine
    out = b.generate(prompt = "x", steps = 4, seed = 1)
    assert len(out["images"]) == 1 and len(fake.calls) == 1


def test_run_load_redacts_paths_in_progress_error(monkeypatch):
    # A load failure surfaced via load_progress() must run through redact_native_paths, the same scrub the diffusers load path applies, so a registered native path cannot leak.
    from utils import native_path_leases as npl

    secret_root = "/managed/native/root"
    npl._remember_native_path_for_redaction(secret_root, "model dir")
    try:
        b = SdCppDiffusionBackend(engine = _FakeEngine())
        fam = detect_family("z-image")
        monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
        monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)

        def _boom(*a, **k):
            raise RuntimeError(f"failed to read {secret_root}/z.gguf")

        monkeypatch.setattr(b, "_fetch_assets", _boom)

        b._load_token = 1
        b._loading = bk._SdLoading(repo_id = "unsloth/Z-Image-Turbo-GGUF", base_repo = fam.base_repo)
        b._run_load(
            repo_id = "unsloth/Z-Image-Turbo-GGUF",
            gguf_filename = "z.gguf",
            base = fam.base_repo,
            fam = fam,
            hf_token = None,
            _load_token = 1,
        )
        err = b.load_progress()["error"]
        assert err and secret_root not in err and "<native_path>" in err
    finally:
        with npl._REDACTION_LOCK:
            if secret_root in npl._NATIVE_PATH_REDACTIONS:
                npl._NATIVE_PATH_REDACTIONS.remove(secret_root)


# ── LoRA (native engine) ────────────────────────────────────────────────────────


def _fake_materialize(resolved, dest):
    """Stand-in for diffusion_lora.materialize_native_dir: write a stub file per adapter
    into ``dest`` and return the resolved list pointing at the written paths (mirroring the
    real helper's contract without touching the Hub / real weights)."""
    from pathlib import Path as _P

    from core.inference import diffusion_lora as dl

    dest.mkdir(parents = True, exist_ok = True)
    out = []
    for r in resolved:
        p = _P(dest) / f"{r.alias}.safetensors"
        p.write_bytes(b"stub")
        out.append(dl.ResolvedLora(r.id, r.alias, str(p), r.fmt, r.weight))
    return out


def _patch_lora(
    monkeypatch,
    resolved,
    supported = True,
):
    from core.inference import diffusion_lora as dl

    monkeypatch.setattr(dl, "supports_lora", lambda **k: supported)
    monkeypatch.setattr(dl, "resolve_specs", lambda specs, **k: list(resolved))
    monkeypatch.setattr(dl, "materialize_native_dir", _fake_materialize)


def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch):
    # One-shot sd-cli LoRA: adapters materialized into a --lora-model-dir and selected with <lora:ALIAS:w> tags injected into the prompt (the real inject_prompt_tags runs here).
    from core.inference import diffusion_lora as dl

    eng = _FakeEngine()
    b = _loaded_backend(engine = eng)  # mode = "oneshot"
    _patch_lora(
        monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]
    )
    b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)])
    _, params, _, _ = eng.calls[0]
    assert params.lora_dir is not None and params.lora_apply_mode == "auto"
    assert "<lora:myalias:0.8>" in params.prompt


def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tmp_path):
    # Server-mode LoRA rides the structured `lora` request field (the sdcpp API ignores <lora:> prompt tags): adapters staged into the server --lora-model-dir, referenced by their path relative to it + the validated multiplier.
    from pathlib import Path as _P

    from core.inference import diffusion_lora as dl

    b = SdCppDiffusionBackend()
    servers: list = []
    _run_server_load(monkeypatch, b, servers)
    servers[0].lora_dir = str(tmp_path)
    _patch_lora(
        monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]
    )
    b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)])
    payload = servers[0].payloads[0]
    assert "lora" in payload and len(payload["lora"]) == 1
    assert payload["lora"][0]["multiplier"] == 0.7
    assert payload["lora"][0]["path"].endswith("myalias.safetensors")
    assert "<lora:" not in payload["prompt"]  # no prompt-tag mechanism on the server
    # The per-request stage subdir under the server's lora dir is removed after the batch.
    assert not list(_P(tmp_path).glob("gen_*"))


def test_generate_rejects_loras_on_unsupported_family(monkeypatch):
    b = _loaded_backend(engine = _FakeEngine())
    _patch_lora(monkeypatch, [], supported = False)
    with pytest.raises(ValueError, match = "LoRA is not supported"):
        b.generate(prompt = "x", steps = 4, seed = 1, loras = [("id1", 1.0)])


def test_generate_zero_weight_loras_are_noop(monkeypatch):
    # weight-0 rows are dropped BEFORE the support gate, so a request carrying only disabled adapters stays a no-op even on a family where native LoRA is unsupported.
    eng = _FakeEngine()
    b = _loaded_backend(engine = eng)
    _patch_lora(monkeypatch, [], supported = False)  # would raise if the gate were reached
    b.generate(prompt = "x", steps = 4, seed = 1, loras = [("id1", 0.0)])
    _, params, _, _ = eng.calls[0]
    assert params.lora_dir is None  # nothing applied


def test_generate_rejects_controlnet_on_native_engine():
    # ControlNet is diffusers-only. The route passes `controlnet` to whichever engine is active, so the native backend must reject it with a clean ValueError (400) rather than TypeError (500).
    b = _loaded_backend(engine = _FakeEngine())
    with pytest.raises(ValueError, match = "ControlNet is not yet supported on the native"):
        b.generate(prompt = "x", steps = 4, seed = 1, controlnet = ("id", "img", "canny", 1.0, 0.0, 1.0))


@pytest.mark.parametrize("cn_strength", [0, 0.0, None])
def test_generate_treats_zero_strength_controlnet_as_disabled(cn_strength):
    # strength 0 (or None) disables ControlNet -- the diffusers path treats it as plain txt2img -- so a strength-0 spec must succeed on the native engine too. Only a genuinely active ControlNet is rejected.
    eng = _FakeEngine()
    b = _loaded_backend(engine = eng)
    out = b.generate(
        prompt = "x",
        steps = 4,
        seed = 1,
        controlnet = ("id", "img", "canny", cn_strength, 0.0, 1.0),
    )
    assert len(out["images"]) == 1


def test_generate_rejects_image_conditioned_on_native_engine():
    # img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call with an init image on the native engine gets a clean ValueError, not a silent txt2img.
    b = _loaded_backend(engine = _FakeEngine())
    with pytest.raises(ValueError, match = "not yet supported on the native"):
        b.generate(prompt = "x", steps = 4, seed = 1, init_image = "data:image/png;base64,AAAA")


def test_status_native_reports_supports_controlnet_false():
    b = _loaded_backend()
    assert b.status()["supports_controlnet"] is False