from __future__ import annotations

import json

import numpy as np
import pytest

from native_token_compatibility import (
    compare_native_tokens,
    main,
    validate_native_tokens,
)


def _valid_tokens(frames: int = 5) -> np.ndarray:
    tokens = np.zeros((frames, 8), dtype=np.int64)
    tokens[:, 0] = np.arange(frames) % 16_384
    for index in range(1, 8):
        tokens[:, index] = (np.arange(frames) + index) % 1_024
    return tokens


def test_exact_cross_layout_comparison():
    tokens = _valid_tokens()
    result = compare_native_tokens(
        tokens,
        tokens.T,
        reference_layout="frames_first",
        recovered_layout="codebooks_first",
    )

    assert result["exact_match"] is True
    assert result["reference_shape_frames_first"] == [5, 8]
    assert result["duration_seconds_at_25hz"] == pytest.approx(0.2)
    assert all(item["agreement"] == 1.0 for item in result["per_codebook"])


def test_mismatch_reports_exact_per_codebook_agreement():
    reference = _valid_tokens(4)
    recovered = reference.copy()
    recovered[0, 3] += 1

    result = compare_native_tokens(
        reference,
        recovered,
        reference_layout="frames_first",
        recovered_layout="frames_first",
    )

    assert result["exact_match"] is False
    assert result["per_codebook"][2]["agreement"] == 1.0
    assert result["per_codebook"][3]["matching_frames"] == 3
    assert result["per_codebook"][3]["agreement"] == 0.75


def test_shape_mismatch_is_not_truncated_or_aligned():
    result = compare_native_tokens(
        _valid_tokens(4),
        _valid_tokens(3),
        reference_layout="frames_first",
        recovered_layout="frames_first",
    )
    assert result["shape_match"] is False
    assert result["frames"] is None
    assert all(item["agreement"] is None for item in result["per_codebook"])


@pytest.mark.parametrize(
    ("codebook", "value", "message"),
    [(0, 16_384, "c0"), (1, 1_024, "c1"), (7, -1, "c7")],
)
def test_codebook_ranges_are_enforced(codebook, value, message):
    tokens = _valid_tokens()
    tokens[0, codebook] = value
    with pytest.raises(ValueError, match=message):
        validate_native_tokens(tokens, layout="frames_first")


def test_layout_must_match_declared_axis():
    with pytest.raises(ValueError, match="frames_first"):
        validate_native_tokens(_valid_tokens().T, layout="frames_first")


def test_floating_point_tokens_are_rejected():
    with pytest.raises(TypeError, match="integer dtype"):
        validate_native_tokens(_valid_tokens().astype(np.float32), layout="frames_first")


def test_malformed_torch_token_file_json_cli_returns_structured_error(tmp_path, capsys):
    malformed = tmp_path / "malformed.pt"
    malformed.write_bytes(b"not a torch token file")
    valid = tmp_path / "valid.npy"
    np.save(valid, _valid_tokens())

    result = main(
        [
            str(malformed),
            str(valid),
            "--reference-layout",
            "frames_first",
            "--recovered-layout",
            "frames_first",
            "--json",
        ]
    )
    assert result == 2
    payload = json.loads(capsys.readouterr().out)
    assert payload["status"] == "ERROR"
    assert "unable to safely load token file" in payload["error"]