// SPDX-License-Identifier: MIT OR Apache-2.0

//! Cross-validation of `.pth` parsing against `PyTorch` reference data.
//!
//! Each test loads an `AlgZoo` `.pth` file with `parse_pth()`, then compares
//! tensor names, shapes, dtypes, and **raw bytes** against a JSON manifest
//! generated by `PyTorch` (`generate_pth_reference.py`). A mismatch at any
//! byte means the pickle VM or tensor extraction diverged from `PyTorch`.
//!
//! # Fixtures
//!
//! | File | Source | Architecture | Format |
//! |------|--------|-------------|--------|
//! | `algzoo_rnn_small.pth` | `AlgZoo` `2nd_argmax_2_2` | RNN (3 tensors, 10 params) | Newer `PyTorch` (`archive/` prefix) |
//! | `algzoo_transformer_small.pth` | `AlgZoo` `longest_cycle_2_3` | Transformer (7 tensors, 50 params) | Newer `PyTorch` (`archive/` prefix) |
//! | `algzoo_rnn_blog.pth` | `AlgZoo` blog example | RNN (3 tensors, 432 params) | Older `PyTorch` (model-name prefix) |

#![cfg(feature = "pth")]
#![allow(
    clippy::panic,
    clippy::indexing_slicing,
    clippy::unwrap_used,
    clippy::as_conversions,
    clippy::wildcard_enum_match_arm
)]

use std::path::PathBuf;

use anamnesis::parse_pth;

/// Minimal JSON reference structure (deserialized with serde).
#[derive(serde::Deserialize)]
struct Reference {
    num_tensors: usize,
    tensors: Vec<TensorRef>,
}

#[derive(serde::Deserialize)]
struct TensorRef {
    name: String,
    shape: Vec<usize>,
    dtype: String,
    data_hex: String,
}

fn fixture_dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("fixtures")
        .join("pth_reference")
}

fn load_reference(json_name: &str) -> Reference {
    let path = fixture_dir().join(json_name);
    let data = std::fs::read_to_string(&path)
        .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
    serde_json::from_str(&data)
        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", path.display()))
}

fn hex_to_bytes(hex: &str) -> Vec<u8> {
    (0..hex.len())
        .step_by(2)
        .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap())
        .collect()
}

/// Core cross-validation: parse the .pth file with our Rust parser, then
/// compare every tensor against the Python-generated reference.
fn cross_validate(pth_name: &str, json_name: &str) {
    let pth_path = fixture_dir().join(pth_name);
    let parsed =
        parse_pth(&pth_path).unwrap_or_else(|e| panic!("parse_pth({pth_name}) failed: {e}"));
    let tensors = parsed.tensors().unwrap();

    let reference = load_reference(json_name);

    assert_eq!(
        tensors.len(),
        reference.num_tensors,
        "tensor count mismatch for {pth_name}: rust={}, python={}",
        tensors.len(),
        reference.num_tensors
    );

    for (rust_tensor, py_ref) in tensors.iter().zip(reference.tensors.iter()) {
        // Name
        assert_eq!(
            rust_tensor.name, py_ref.name,
            "tensor name mismatch in {pth_name}"
        );

        // Shape
        assert_eq!(
            rust_tensor.shape, py_ref.shape,
            "shape mismatch for tensor `{}` in {pth_name}: rust={:?}, python={:?}",
            rust_tensor.name, rust_tensor.shape, py_ref.shape
        );

        // Dtype
        assert_eq!(
            rust_tensor.dtype.to_string(),
            py_ref.dtype,
            "dtype mismatch for tensor `{}` in {pth_name}: rust={}, python={}",
            rust_tensor.name,
            rust_tensor.dtype,
            py_ref.dtype
        );

        // Raw bytes — the definitive test: byte-exact match
        let expected_bytes = hex_to_bytes(&py_ref.data_hex);
        assert_eq!(
            rust_tensor.data.len(),
            expected_bytes.len(),
            "data length mismatch for tensor `{}` in {pth_name}: rust={}, python={}",
            rust_tensor.name,
            rust_tensor.data.len(),
            expected_bytes.len()
        );
        assert_eq!(
            rust_tensor.data.as_ref(),
            expected_bytes.as_slice(),
            "DATA MISMATCH for tensor `{}` in {pth_name}: \
             first differing byte at position {}",
            rust_tensor.name,
            rust_tensor
                .data
                .iter()
                .zip(expected_bytes.iter())
                .position(|(a, b)| a != b)
                .unwrap_or(0)
        );
    }
}

// -- `AlgZoo` RNN (newer PyTorch format, archive/ prefix) ----------------------

#[test]
fn cross_validate_algzoo_rnn_small() {
    cross_validate("algzoo_rnn_small.pth", "algzoo_rnn_small_reference.json");
}

// -- `AlgZoo` Transformer (newer PyTorch format, archive/ prefix) --------------

#[test]
fn cross_validate_algzoo_transformer_small() {
    cross_validate(
        "algzoo_transformer_small.pth",
        "algzoo_transformer_small_reference.json",
    );
}

// -- `AlgZoo` RNN blog post (older PyTorch format, model-name prefix) ----------

#[test]
fn cross_validate_algzoo_rnn_blog() {
    cross_validate("algzoo_rnn_blog.pth", "algzoo_rnn_blog_reference.json");
}

// -- Roundtrip: .pth → .safetensors → read back → compare ------------------

/// Parse a `.pth`, convert to `.safetensors`, read the `.safetensors` back,
/// and verify all tensor data matches the `PyTorch` reference byte-for-byte.
fn roundtrip_to_safetensors(pth_name: &str, json_name: &str) {
    let pth_path = fixture_dir().join(pth_name);
    let parsed = parse_pth(&pth_path).unwrap();
    let tensors = parsed.tensors().unwrap();

    // Write to a temporary .safetensors file.
    let tmp = tempfile::NamedTempFile::new().unwrap();
    anamnesis::pth_to_safetensors(&tensors, tmp.path()).unwrap();

    // Read the .safetensors back.
    let st_data = std::fs::read(tmp.path()).unwrap();
    let st = safetensors::SafeTensors::deserialize(&st_data).unwrap();

    // Compare against Python reference.
    let reference = load_reference(json_name);
    assert_eq!(st.len(), reference.num_tensors);

    for py_ref in &reference.tensors {
        let st_tensor = st
            .tensor(&py_ref.name)
            .unwrap_or_else(|e| panic!("tensor `{}` not in safetensors: {e}", py_ref.name));

        assert_eq!(
            st_tensor.shape().to_vec(),
            py_ref.shape,
            "shape mismatch after roundtrip for tensor `{}`",
            py_ref.name
        );

        let expected_bytes = hex_to_bytes(&py_ref.data_hex);
        assert_eq!(
            st_tensor.data(),
            expected_bytes.as_slice(),
            "DATA MISMATCH after roundtrip for tensor `{}`",
            py_ref.name
        );
    }
}

#[test]
fn roundtrip_algzoo_rnn_small() {
    roundtrip_to_safetensors("algzoo_rnn_small.pth", "algzoo_rnn_small_reference.json");
}

#[test]
fn roundtrip_algzoo_transformer_small() {
    roundtrip_to_safetensors(
        "algzoo_transformer_small.pth",
        "algzoo_transformer_small_reference.json",
    );
}

#[test]
fn roundtrip_algzoo_rnn_blog() {
    roundtrip_to_safetensors("algzoo_rnn_blog.pth", "algzoo_rnn_blog_reference.json");
}
