"""Binary-safe auxiliary serialization tests."""

from pathlib import Path

import torch
import torch.nn as nn

from model_osmosis.format.image import OsmoImage
from model_osmosis.format.serialization import save_osmo, load_osmo


def _make_model(seed: int = 0) -> nn.Module:
    torch.manual_seed(seed)
    model = nn.Sequential(
        nn.Linear(4, 8),
        nn.ReLU(),
        nn.Linear(8, 2),
    )
    model.eval()
    return model


def test_auxiliary_loader_is_binary_safe(tmp_path):
    image = OsmoImage.build(_make_model(0), name="aux-binary-safe")
    osmo = image.osmo_model
    osmo.auxiliary_files["tokenizer/spiece.model"] = b"\x00\xff\x10BINARY\x81"
    osmo.auxiliary_files["tokenizer/vocab.txt"] = "a\nb\nc\n"
    osmo.auxiliary_files["tokenizer/tokenizer_config.json"] = {"model_max_length": 128}

    out = Path(tmp_path) / "aux-binary-safe.osmo"
    save_osmo(osmo, out, compress=True, obfuscate=False)
    loaded = load_osmo(out)

    assert isinstance(loaded.auxiliary_files["tokenizer/spiece.model"], bytes)
    assert loaded.auxiliary_files["tokenizer/spiece.model"] == b"\x00\xff\x10BINARY\x81"
    assert isinstance(loaded.auxiliary_files["tokenizer/vocab.txt"], str)
    assert loaded.auxiliary_files["tokenizer/vocab.txt"] == "a\nb\nc\n"
    assert isinstance(loaded.auxiliary_files["tokenizer/tokenizer_config.json"], dict)
