"""Tests for trusted .osmoimg signing, verification, and reproducibility."""

from __future__ import annotations

import struct
import tempfile
from pathlib import Path

import pytest
import torch
import torch.nn as nn

from model_osmosis.format.image import OsmoImage


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


def test_signed_artifact_is_trusted(monkeypatch):
    with tempfile.TemporaryDirectory() as tmpdir:
        monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(Path(tmpdir) / "trust"))

        image = OsmoImage.build(_make_model(0), name="trusted-default")
        out = Path(tmpdir) / "trusted.osmoimg"
        image.save(out)

        verify = OsmoImage.verify(out)
        assert verify["trust_status"] == "trusted"
        assert verify["header_version"] == 3
        assert verify["signature_alg"] == "ed25519"

        info = OsmoImage.inspect(out)
        assert info["runtime"]["trust_status"] == "trusted"
        assert info["runtime"]["signer_id"] == "default"
        assert info["runtime"]["payload_sha256"]
        assert info["runtime"]["artifact_id"]
        assert info["runtime"]["manifest_sha256"]
        assert info["runtime"]["tensor_index_sha256"]


def test_unsigned_artifact_reports_unsigned(monkeypatch):
    with tempfile.TemporaryDirectory() as tmpdir:
        monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(Path(tmpdir) / "trust"))

        image = OsmoImage.build(_make_model(1), name="unsigned", trust_sign=False)
        out = Path(tmpdir) / "unsigned.osmoimg"
        image.save(out, trust_sign=False)

        verify = OsmoImage.verify(out)
        assert verify["trust_status"] == "unsigned"


def test_tampered_payload_reports_tampered(monkeypatch):
    with tempfile.TemporaryDirectory() as tmpdir:
        monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(Path(tmpdir) / "trust"))

        image = OsmoImage.build(_make_model(2), name="tamper-test")
        out = Path(tmpdir) / "tamper.osmoimg"
        image.save(out)

        with open(out, "r+b") as fh:
            # Move to payload start after magic/version/header.
            fh.seek(8)  # magic
            fh.read(4)  # version
            header_len = struct.unpack("<I", fh.read(4))[0]
            payload_offset = 8 + 4 + 4 + header_len
            fh.seek(payload_offset + 16)
            b = fh.read(1)
            fh.seek(payload_offset + 16)
            fh.write(bytes([((b[0] if b else 0) ^ 0xFF)]))

        verify = OsmoImage.verify(out)
        assert verify["trust_status"] == "tampered"


def test_strict_repro_mode_is_byte_reproducible(monkeypatch):
    with tempfile.TemporaryDirectory() as tmpdir:
        monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(Path(tmpdir) / "trust"))
        model_a = _make_model(3)
        model_b = _make_model(3)

        image_a = OsmoImage.build(
            model_a,
            name="repro-a",
            repro_mode="strict",
            trust_sign=True,
            compression_profile="lossless-max",
        )
        image_b = OsmoImage.build(
            model_b,
            name="repro-a",
            repro_mode="strict",
            trust_sign=True,
            compression_profile="lossless-max",
        )

        out_a = Path(tmpdir) / "a.osmoimg"
        out_b = Path(tmpdir) / "b.osmoimg"
        image_a.save(out_a, repro_mode="strict", trust_sign=True, compression_profile="lossless-max")
        image_b.save(out_b, repro_mode="strict", trust_sign=True, compression_profile="lossless-max")

        assert out_a.read_bytes() == out_b.read_bytes()


def test_provenance_modes_recorded(monkeypatch):
    with tempfile.TemporaryDirectory() as tmpdir:
        monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(Path(tmpdir) / "trust"))

        private_image = OsmoImage.build(_make_model(4), name="prov-private", provenance_mode="private")
        private_path = Path(tmpdir) / "private.osmoimg"
        private_image.save(private_path, provenance_mode="private")
        private_loaded = OsmoImage.load(private_path)
        assert private_loaded.runtime.provenance_mode == "private"
        assert "lineage_hash" in private_loaded.runtime.provenance

        full_image = OsmoImage.build(_make_model(5), name="prov-full", provenance_mode="full")
        full_path = Path(tmpdir) / "full.osmoimg"
        full_image.save(full_path, provenance_mode="full")
        full_loaded = OsmoImage.load(full_path)
        assert full_loaded.runtime.provenance_mode == "full"
        assert "source_framework" in full_loaded.runtime.provenance


def test_kms_backend_sign_and_verify(monkeypatch):
    with tempfile.TemporaryDirectory() as tmpdir:
        monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(Path(tmpdir) / "trust"))
        monkeypatch.setenv("OSMOSIS_SIGNER_BACKEND", "aws-kms")
        monkeypatch.setenv("OSMOSIS_KMS_KEY_REF", "projects/demo/keys/model-osmo")

        image = OsmoImage.build(_make_model(6), name="kms-signed")
        out = Path(tmpdir) / "kms.osmoimg"
        image.save(out)

        verify = OsmoImage.verify(out)
        assert verify["trust_status"] == "trusted"
        assert verify["signer_backend"] == "aws-kms"
        assert verify["kms_provider"] == "aws-kms"
        assert verify["kms_key_ref"] == "projects/demo/keys/model-osmo"

        inspected = OsmoImage.inspect(out)
        assert inspected["runtime"]["signer_backend"] == "aws-kms"


def test_transparency_publish_is_non_blocking(monkeypatch):
    with tempfile.TemporaryDirectory() as tmpdir:
        monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(Path(tmpdir) / "trust"))
        monkeypatch.setenv("OSMOSIS_TRANSPARENCY_URL", "http://127.0.0.1:1/transparency")
        monkeypatch.setenv("OSMOSIS_TRANSPARENCY_TIMEOUT_SEC", "0.2")

        image = OsmoImage.build(_make_model(7), name="transparency")
        out = Path(tmpdir) / "transparency.osmoimg"
        image.save(out)

        assert image.runtime.transparency_status in {"local_only", "published"}

        verify = OsmoImage.verify(out)
        assert verify["trust_status"] == "trusted"
        assert verify["transparency_status"] == "local_found"


def test_v1_v2_v3_compatibility_matrix(monkeypatch):
    with tempfile.TemporaryDirectory() as tmpdir:
        monkeypatch.setenv("OSMOSIS_TRUST_HOME", str(Path(tmpdir) / "trust"))

        model = _make_model(8)

        v1_image = OsmoImage.build(model, name="compat-v1", compression_profile="legacy", trust_sign=False)
        v1_path = Path(tmpdir) / "compat-v1.osmoimg"
        v1_image.save(v1_path, compression_profile="legacy", trust_sign=False)
        v1_verify = OsmoImage.verify(v1_path)
        assert v1_verify["header_version"] == 1
        assert v1_verify["trust_status"] == "unsigned"
        assert OsmoImage.load(v1_path).parameter_count > 0
        assert OsmoImage.inspect(v1_path)["runtime"]["compression_profile"] == "legacy"

        pytest.importorskip("zstandard")
        v2_image = OsmoImage.build(model, name="compat-v2", compression_profile="lossless-max", trust_sign=False)
        v2_path = Path(tmpdir) / "compat-v2.osmoimg"
        v2_image.save(v2_path, compression_profile="lossless-max", trust_sign=False)
        v2_verify = OsmoImage.verify(v2_path)
        assert v2_verify["header_version"] == 2
        assert v2_verify["trust_status"] == "unsigned"
        assert OsmoImage.load(v2_path).parameter_count > 0
        assert OsmoImage.inspect(v2_path)["runtime"]["compression_profile"] == "lossless-max"

        v3_image = OsmoImage.build(model, name="compat-v3", compression_profile="legacy", trust_sign=True)
        v3_path = Path(tmpdir) / "compat-v3.osmoimg"
        v3_image.save(v3_path, compression_profile="legacy", trust_sign=True)
        v3_verify = OsmoImage.verify(v3_path)
        assert v3_verify["header_version"] == 3
        assert v3_verify["trust_status"] in {"trusted", "signed_untrusted"}
        assert OsmoImage.load(v3_path).parameter_count > 0
        assert OsmoImage.inspect(v3_path)["runtime"]["artifact_id"]
