from __future__ import annotations

import io
import json
import pickle
import sys
import tarfile
from pathlib import Path

import numpy as np
from PIL import Image


REPO_ROOT = Path(__file__).resolve().parents[1]
STAGES = REPO_ROOT / "scripts" / "stages"
if str(STAGES) not in sys.path:
    sys.path.insert(0, str(STAGES))

from evaluate_mlperf_image_quality import (  # noqa: E402
    canonical_digest,
    load_checkpoint,
    load_cifar10_samples,
    load_vww_samples,
)


def add_bytes(archive: tarfile.TarFile, name: str, payload: bytes) -> None:
    member = tarfile.TarInfo(name)
    member.size = len(payload)
    archive.addfile(member, io.BytesIO(payload))


def test_cifar_loader_applies_pinned_indices_and_planar_to_nhwc(tmp_path: Path) -> None:
    data = np.zeros((10000, 3072), dtype=np.uint8)
    labels = [0] * 10000
    labels[17] = 3
    labels[991] = 8
    data[17, :1024] = 11
    data[17, 1024:2048] = 22
    data[17, 2048:] = 33
    data[991, :1024] = 44
    archive_path = tmp_path / "cifar.tar.gz"
    with tarfile.open(archive_path, "w:gz") as archive:
        add_bytes(
            archive,
            "cifar-10-batches-py/test_batch",
            pickle.dumps({b"data": data, b"labels": labels}),
        )
    indices_path = tmp_path / "indices.npy"
    np.save(indices_path, np.asarray([17, 991], dtype=np.int64), allow_pickle=False)
    labels_path = tmp_path / "labels.csv"
    labels_path.write_text("first.bin,10,3\nsecond.bin,10,8\n")

    samples = load_cifar10_samples(archive_path, labels_path, indices_path)

    assert [sample.label for sample in samples] == [3, 8]
    assert [sample.sample_id for sample in samples] == ["cifar10_test_00017", "cifar10_test_00991"]
    assert samples[0].image.shape == (32, 32, 3)
    assert samples[0].image[0, 0].tolist() == [11, 22, 33]
    assert samples[1].image[0, 0].tolist() == [44, 0, 0]


def test_vww_loader_matches_coco_ids_and_decodes_rgb(tmp_path: Path) -> None:
    archive_path = tmp_path / "vww.tar.gz"
    expected = {
        "000000000123": np.full((96, 96, 3), [255, 0, 0], dtype=np.uint8),
        "000000999999": np.full((96, 96, 3), [0, 255, 0], dtype=np.uint8),
    }
    with tarfile.open(archive_path, "w:gz") as archive:
        for image_id, array in expected.items():
            payload = io.BytesIO()
            Image.fromarray(array).save(payload, format="PNG")
            add_bytes(archive, f"vw_coco2014_96/class/COCO_val2014_{image_id}.png", payload.getvalue())
    labels_path = tmp_path / "labels.csv"
    labels_path.write_text("000000000123.bin,2,1\n000000999999.bin,2,0\n")

    samples = load_vww_samples(archive_path, labels_path)

    assert [sample.sample_id for sample in samples] == ["coco_000000000123", "coco_000000999999"]
    assert [sample.label for sample in samples] == [1, 0]
    assert np.array_equal(samples[0].image, expected["000000000123"])
    assert np.array_equal(samples[1].image, expected["000000999999"])


def test_checkpoint_requires_same_canonical_config(tmp_path: Path) -> None:
    digest = canonical_digest({"model_id": "VC02", "threshold": 0.85})
    checkpoint = tmp_path / "checkpoint.jsonl"
    checkpoint.write_text(
        json.dumps({"config_sha256": digest, "sample_id": "one", "label": 1}) + "\n"
    )
    assert load_checkpoint(checkpoint, digest)["one"]["label"] == 1

    try:
        load_checkpoint(checkpoint, "0" * 64)
    except ValueError as error:
        assert "config mismatch" in str(error)
    else:
        raise AssertionError("checkpoint from a different config was accepted")