"""Ad-hoc check: the shared transfer engine used by fetch_image_model.py.

Proves the four guarantees the image downloader relies on, against a real
(public, ungated) Hub repo so nothing is mocked:

  1. first run fetches every file, byte-exact against the published size,
  2. a second run skips everything (download only what is missing),
  3. an interrupted run resumes from its ledger instead of starting over,
  4. a corrupt file of the CORRECT length is rejected by the sha256 gate and
     deleted, so no invalid file is ever promoted into place.

    python scripts/_selftest_transfer.py
"""
import hashlib
import json
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_i2v_model as engine  # noqa: E402

# Ungated and small, but a real ~90 MB safetensors weight file, so an
# interruption genuinely leaves a partial transfer behind.
REPO = "sentence-transformers/all-MiniLM-L6-v2"
WEIGHT = "model.safetensors"

ok = True


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


work = Path(tempfile.mkdtemp())
try:
    manifest = engine.fetch_manifest(REPO, allow=[WEIGHT], ignore=[])
    size = manifest["files"][WEIGHT]
    expected_sha = manifest["oids"][WEIGHT]
    dest = work / WEIGHT
    print(f"target: {REPO}/{WEIGHT}  ({size / 1e6:.1f} MB)")

    print("1. first run fetches the file, byte-exact")
    check("no errors", engine.download_snapshot(manifest, work) == [])
    check("size matches the published size", dest.stat().st_size == size)
    good = dest.read_bytes()
    check("sha256 matches the Hub's published hash",
          hashlib.sha256(good).hexdigest() == expected_sha)

    print("2. second run skips it (only missing files are downloaded)")
    before = engine._PROGRESS[0]
    check("no errors", engine.download_snapshot(manifest, work) == [])
    check("zero bytes re-transferred", engine._PROGRESS[0] == before)

    print("3. an interrupted run resumes instead of restarting")
    half = size // 2
    dest.unlink()
    part = dest.with_name(dest.name + ".part")
    with part.open("wb") as handle:
        handle.truncate(size)          # pre-allocated to the full size
    with part.open("r+b") as handle:  # first half is genuinely good
        handle.write(good[:half])
    part.with_name(dest.name + ".parts.json").write_text(json.dumps({
        "repo": manifest["repo"], "revision": manifest["revision"],
        "size": size, "done": [[0, half]],
    }))
    baseline = engine._PROGRESS[0]
    engine._download_file(
        manifest["repo"], manifest["revision"], WEIGHT, dest, size,
        expect_sha256=expected_sha,
    )
    re_fetched = engine._PROGRESS[0] - baseline
    check(f"only the missing tail was fetched ({re_fetched / 1e6:.1f} MB of {size / 1e6:.1f} MB)",
          0 < re_fetched < size * 0.75)
    check("resumed file is byte-exact",
          hashlib.sha256(dest.read_bytes()).hexdigest() == expected_sha)
    check("scratch file and ledger cleaned up",
          not part.exists() and not part.with_name(dest.name + ".parts.json").exists())

    print("4. a file the ledger calls complete but whose bytes are wrong is rejected")
    dest.unlink()
    with part.open("wb") as handle:
        handle.truncate(size)          # right length...
    with part.open("r+b") as handle:  # ...wrong bytes
        handle.seek(size // 2)
        handle.write(b"CORRUPTION")
    # The ledger claims every byte already landed, so the downloader has
    # nothing left to fetch and the sha256 gate is the ONLY thing that can
    # catch this. This is the case that once discarded ~2.6 GB of perfectly
    # good weights when the hash was read from the CDN `ETag` instead.
    part.with_name(dest.name + ".parts.json").write_text(json.dumps({
        "repo": manifest["repo"], "revision": manifest["revision"],
        "size": size, "done": [[0, size]],
    }))
    check("size check alone cannot catch it", part.stat().st_size == size)
    baseline = engine._PROGRESS[0]
    try:
        engine._download_file(
            manifest["repo"], manifest["revision"], WEIGHT, dest, size,
            expect_sha256=expected_sha,
        )
        check("corrupt file rejected by the sha256 gate", False)
    except RuntimeError as exc:
        check("corrupt file rejected by the sha256 gate",
              "sha256 mismatch" in str(exc))
    check("nothing was re-downloaded to paper over it",
          engine._PROGRESS[0] == baseline)
    check("corrupt scratch and ledger deleted",
          not part.exists()
          and not part.with_name(dest.name + ".parts.json").exists())
    check("no invalid file promoted into place", not dest.exists())
    check("verify_snapshot still reports it missing",
          engine.verify_snapshot(work, manifest, [WEIGHT], []) != [])
finally:
    shutil.rmtree(work, ignore_errors=True)

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