"""Tests for adaptive lossless v2 compression."""

from __future__ import annotations

import numpy as np
import pytest

from model_osmosis.format.adaptive_compression import (
    PayloadIntegrityError,
    inspect_osmo_model_lossless_v2,
    pack_osmo_model_lossless_v2,
    unpack_osmo_model_lossless_v2,
)
from model_osmosis.format.converters import UniversalConverter


pytest.importorskip("zstandard")


def _make_osmo_with_tensors():
    rng = np.random.default_rng(123)
    base = {
        "c_float32": rng.standard_normal((16, 8), dtype=np.float32),
        "f_float16": np.asfortranarray(rng.standard_normal((9, 7), dtype=np.float32)).astype(np.float16),
        "int8": rng.integers(-128, 127, size=(64,), dtype=np.int8),
    }
    # Build a minimal valid OsmoModel from a tiny torch-like equivalent.
    import torch
    import torch.nn as nn

    model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 4))
    osmo = UniversalConverter().to_osmo(model, name="adaptive-test")

    # Replace a few parameters with known tensors to assert exact bytes.
    for i, key in enumerate(list(osmo.parameters.keys())[:3]):
        osmo.parameters[key] = list(base.values())[i]
    return osmo


def test_v2_pack_unpack_bit_exact():
    osmo = _make_osmo_with_tensors()
    original = {k: v.copy() for k, v in osmo.parameters.items()}

    payload, stats = pack_osmo_model_lossless_v2(osmo)
    loaded, loaded_stats = unpack_osmo_model_lossless_v2(payload)

    assert stats["tensor_count"] >= 1
    assert loaded_stats["tensor_count"] >= 1
    for name, arr in original.items():
        out = loaded.parameters[name]
        assert out.dtype == arr.dtype
        assert out.shape == arr.shape
        assert out.tobytes() == arr.tobytes()


def test_v2_inspect_without_tensor_decode():
    osmo = _make_osmo_with_tensors()
    payload, _ = pack_osmo_model_lossless_v2(osmo)

    meta_only, stats = inspect_osmo_model_lossless_v2(payload)
    assert meta_only.manifest.name == osmo.manifest.name
    assert stats["tensor_count"] >= 1
    assert meta_only.parameters == {}


def test_v2_corruption_raises_clear_error():
    osmo = _make_osmo_with_tensors()
    payload, _ = pack_osmo_model_lossless_v2(osmo)

    import io
    import zipfile

    in_buf = io.BytesIO(payload)
    out_buf = io.BytesIO()
    with zipfile.ZipFile(in_buf, "r") as zin, zipfile.ZipFile(out_buf, "w", compression=zipfile.ZIP_STORED) as zout:
        for name in zin.namelist():
            data = zin.read(name)
            if name.startswith("parameters/") and data:
                broken = bytearray(data)
                broken[0] ^= 0x01
                data = bytes(broken)
            zout.writestr(name, data)
    corrupted = out_buf.getvalue()

    with pytest.raises(PayloadIntegrityError):
        unpack_osmo_model_lossless_v2(corrupted)
