"""Ad-hoc check: the image downloader's own logic (paths, filters, pre-flight).

Separate from `_selftest_transfer.py` (which exercises the shared engine)
because these are the image-specific decisions: where the snapshot lands,
which files are wanted, and whether this host can run the result.

    python scripts/_selftest_fetch_image.py
"""
import contextlib
import io
import shutil
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import fetch_image_model as f  # noqa: E402

ok = True


def check(label, condition):
    global ok
    ok = ok and bool(condition)
    print(f"  [{'PASS' if condition else 'FAIL'}] {label}")


print("1. the snapshot lands at the expected absolute path")
target = f.resolve_target()
expected = Path(__file__).resolve().parent.parent / "models" / "flux2-klein-4b"
check(f"resolve_target() == {expected}", target == expected)
check("target is absolute", target.is_absolute())
check("repo is the FLUX.2 Klein snapshot", f.resolve_repo() == f.DEFAULT_REPO)
check(
    "repo is black-forest-labs/FLUX.2-klein-4B",
    f.DEFAULT_REPO == "black-forest-labs/FLUX.2-klein-4B",
)

print("2. the Diffusers components are wanted, the extras are not")
for rel in ("README.md",
            "LICENSE.md",
            "realism.jpg",
            "editing.jpg",
            "others.jpg",
            "flux-2-klein-4b.safetensors",
            "flux-2-klein-4b.workflow.json"):
    check(f"{rel} is ignored", not f.transfer._wanted(rel, f.ALLOW_PATTERNS, f.IGNORE_PATTERNS))
for rel in ("model_index.json",
            "scheduler/scheduler_config.json",
            "tokenizer/tokenizer.json",
            "tokenizer/chat_template.jinja",
            "text_encoder/model-00001-of-00002.safetensors",
            "text_encoder/model-00002-of-00002.safetensors",
            "text_encoder/model.safetensors.index.json",
            "transformer/diffusion_pytorch_model.safetensors",
            "vae/diffusion_pytorch_model.safetensors"):
    check(f"{rel} is still wanted",
          f.transfer._wanted(rel, f.ALLOW_PATTERNS, f.IGNORE_PATTERNS))

print("3. nothing REQUIRED_FILES needs is filtered out")
for rel in f.REQUIRED_FILES:
    check(f"required {rel}",
          f.transfer._wanted(rel, f.ALLOW_PATTERNS, f.IGNORE_PATTERNS))

print("4. reclaim_unusable_partials only drops un-resumable state")
work = Path(tempfile.mkdtemp())
try:
    orphan = work / "unet" / "diffusion_pytorch_model.fp16.safetensors.part"
    orphan.parent.mkdir(parents=True)
    orphan.write_bytes(b"x" * 1000)  # wrong size, no ledger -> dead weight

    resumable = work / "vae" / "diffusion_pytorch_model.fp16.safetensors.part"
    resumable.parent.mkdir(parents=True)
    resumable.write_bytes(b"y" * 500)
    ledger = resumable.with_name(resumable.name[: -len(".part")] + ".parts.json")
    ledger.write_text('{"repo":"x","revision":"y","size":500,"done":[[0,500]]}')

    manifest = {"files": {
        "unet/diffusion_pytorch_model.fp16.safetensors": 9999,
        "vae/diffusion_pytorch_model.fp16.safetensors": 500,
    }}
    freed = f.reclaim_unusable_partials(work, manifest)
    check("ledger-less orphan .part reclaimed", not orphan.exists())
    check(f"its bytes were freed ({freed})", freed == 1000)
    check("resumable .part kept", resumable.exists())
    check("its resume ledger kept", ledger.exists())
finally:
    shutil.rmtree(work, ignore_errors=True)

print("5. pre-flight reports this host honestly (no download is started)")
print(f"  HF token found: {'yes' if f._hub_token() else 'no'}")
problems = f.check_hardware()
print(f"  hardware problems on this machine: {len(problems)}")
for p in problems:
    print(f"    - {p.splitlines()[0]}")

print("6. an incomplete-snapshot report names each missing file once")
labelled = f.dedupe([
    "missing: model_index.json",
    "missing: vae/config.json",
    "missing: model_index.json",
    "empty (0 bytes): vae/config.json",
])
check("duplicates dropped, order kept", labelled == [
    "missing: model_index.json",
    "missing: vae/config.json",
    "empty (0 bytes): vae/config.json",
])

report_buffer = io.StringIO()
with contextlib.redirect_stdout(report_buffer):
    code = f.reject(Path("models/flux2-klein-4b"), ["missing: model_index.json"])
report = report_buffer.getvalue()
check("reject() fails the run", code == 1)
check("the report names the folder", "flux2-klein-4b" in report)
check("the report names the missing file", "missing: model_index.json" in report)
check("the report names the fix", "fetch_image_model.py" in report)

print("\nRESULT:", "ALL CHECKS PASSED" if ok else "FAILURES PRESENT")
raise SystemExit(0 if ok else 1)