import hashlib
import json

import numpy as np
import pytest

from solomon_mlx.artifacts import BASE_REVISION, verify_release
from solomon_mlx.evaluation import compare_rows, unit
from solomon_mlx.prepare import verify_base


def test_corrupt_source_and_base_fail_closed(tmp_path):
    path = tmp_path / "input"
    path.write_bytes(b"original")
    sha = hashlib.sha256(b"original").hexdigest()
    (tmp_path / "MANIFEST.json").write_text(
        json.dumps({"files": [{"destination": "input", "staged_bytes": 8, "staged_sha256": sha}]})
    )
    assert verify_release(tmp_path)["input"] == sha
    manifest = {"revision": BASE_REVISION, "files": [{"name": "input", "size": 8, "sha256": sha}]}
    assert verify_base(tmp_path, manifest)["input"] == sha
    path.write_bytes(b"corrupt!")
    with pytest.raises(ValueError):
        verify_release(tmp_path)
    with pytest.raises(ValueError):
        verify_base(tmp_path, manifest)


def test_reserved_choice_exclusion_and_whole_question_accuracy():
    rows = [
        {
            "id": "a",
            "task": "entity",
            "n": 4,
            "head_key": "entity/state4",
            "question_id": "q",
            "gold": 0,
            "letter_logits": [5, 0, 0, 0],
        },
        {
            "id": "b",
            "task": "entity",
            "n": 4,
            "head_key": "entity/state4",
            "question_id": "q",
            "gold": 3,
            "letter_logits": [0, 0, 0, 5],
        },
    ]
    other = [{**rows[0]}, {**rows[1], "letter_logits": [5, 0, 0, 0]}]
    report = compare_rows(rows, other)
    assert report["unit_decision_agreement"] == 0.5
    assert report["mlx_whole_question_accuracy"] == 1
    assert report["cuda_whole_question_accuracy"] == 0
    assert not report["quality_gate_passed"]
    assert unit(
        {"task": "single", "n": 4, "head_key": "single/choiceR", "gold": 3, "letter_logits": [1, 2, 3, 4]}
    ) == (None, None)


def test_norm_offset_preserves_bf16_source_values_in_fp32():
    from types import SimpleNamespace

    import mlx.core as mx
    from mlx_vlm.models.qwen3_5.qwen3_5 import Model

    # Sanitization requires config only for tied embeddings and image channels.
    config = SimpleNamespace(
        text_config=SimpleNamespace(tie_word_embeddings=False), vision_config=SimpleNamespace(in_channels=3)
    )
    original = mx.array([0.02, -0.04], dtype=mx.bfloat16)
    converted = Model.sanitize(
        SimpleNamespace(config=config), {"model.language_model.norm.weight": original.astype(mx.float32)}
    )
    out = converted["language_model.model.norm.weight"]
    assert out.dtype == mx.float32
    np.testing.assert_array_equal(np.asarray(out), np.asarray(original.astype(mx.float32)) + 1)