//! Test Huffman code generation against C++ jpegli reference.
//!
//! Uses test data captured from C++ CreateHuffmanTree() function.

use jpegli::huffman::build_code_lengths;
use std::fs;

/// Parse a single test case from the C++ test data.
#[derive(Debug)]
struct HuffmanTestCase {
    input_length: usize,
    tree_limit: u8,
    input_data: Vec<u64>,
    expected_depth: Vec<u8>,
}

fn parse_test_cases(data: &str) -> Vec<HuffmanTestCase> {
    let mut cases = Vec::new();

    for line in data.lines() {
        let line = line.trim();
        if line.is_empty() || !line.starts_with('{') {
            continue;
        }

        // Remove trailing comma if present
        let line = line.trim_end_matches(',');

        // Parse JSON-like format manually (simple parser)
        if let Some(case) = parse_test_case(line) {
            cases.push(case);
        }
    }

    cases
}

fn parse_test_case(json: &str) -> Option<HuffmanTestCase> {
    // Extract input_length
    let input_length = extract_number(json, "\"input_length\":")?;

    // Extract tree_limit
    let tree_limit = extract_number(json, "\"tree_limit\":")? as u8;

    // Extract input_data array
    let input_data = extract_array(json, "\"input_data\":")?;

    // Extract output_depth array
    let expected_depth = extract_array(json, "\"output_depth\":")?
        .into_iter()
        .map(|x| x as u8)
        .collect();

    Some(HuffmanTestCase {
        input_length,
        tree_limit,
        input_data,
        expected_depth,
    })
}

fn extract_number(json: &str, key: &str) -> Option<usize> {
    let start = json.find(key)? + key.len();
    let rest = &json[start..];
    let end = rest.find(|c: char| c == ',' || c == '}')?;
    rest[..end].trim().parse().ok()
}

fn extract_array(json: &str, key: &str) -> Option<Vec<u64>> {
    let start = json.find(key)? + key.len();
    let rest = &json[start..];
    let arr_start = rest.find('[')?;
    let arr_end = rest.find(']')?;
    let arr_content = &rest[arr_start + 1..arr_end];

    Some(
        arr_content
            .split(',')
            .filter_map(|s| s.trim().parse().ok())
            .collect(),
    )
}

/// Test that Rust build_code_lengths matches C++ CreateHuffmanTree.
#[test]
fn test_huffman_cpp_reference() {
    let testdata_path =
        match jpegli::test_utils::get_cpp_testdata_path("CreateHuffmanTree.testdata") {
            Some(p) => p,
            None => {
                eprintln!("Skipping test: CreateHuffmanTree.testdata not found");
                eprintln!("Set CPP_TESTDATA_DIR env var or generate with:");
                eprintln!("  GENERATE_RUST_TEST_DATA=1 ./build/tools/cjpegli input.png output.jpg");
                return;
            }
        };

    let data = match fs::read_to_string(&testdata_path) {
        Ok(d) => d,
        Err(_) => {
            eprintln!("Skipping test: {:?} could not be read", testdata_path);
            return;
        }
    };

    let cases = parse_test_cases(&data);
    println!("Loaded {} test cases from C++ data", cases.len());

    let mut total = 0;
    let mut passed = 0;
    let mut max_diff = 0u8;

    for (i, case) in cases.iter().enumerate() {
        total += 1;

        // Run Rust implementation
        let rust_depths = build_code_lengths(&case.input_data, case.tree_limit);

        // Compare with C++ expected output
        let mut case_matches = true;
        let mut case_max_diff = 0u8;

        for (j, (&rust_d, &cpp_d)) in rust_depths
            .iter()
            .zip(case.expected_depth.iter())
            .enumerate()
        {
            let diff = (rust_d as i16 - cpp_d as i16).abs() as u8;
            if diff > 0 {
                case_matches = false;
                case_max_diff = case_max_diff.max(diff);
                if diff > max_diff {
                    max_diff = diff;
                    println!(
                        "Case {}, symbol {}: Rust depth={}, C++ depth={}, diff={}",
                        i, j, rust_d, cpp_d, diff
                    );
                }
            }
        }

        if case_matches {
            passed += 1;
        }
    }

    println!("\nResults: {}/{} cases match exactly", passed, total);
    println!("Maximum depth difference: {}", max_diff);

    // For now, just report - we'll tighten this after fixing the algorithm
    if passed < total {
        println!(
            "\nNOTE: {} cases differ - algorithm needs to be ported from C++",
            total - passed
        );
    }
}

/// Test specific edge cases.
#[test]
fn test_huffman_single_symbol() {
    // Single symbol should get depth 1
    let freqs = [100u64, 0, 0, 0];
    let depths = build_code_lengths(&freqs, 16);
    assert_eq!(depths[0], 1);
    assert_eq!(depths[1], 0);
    assert_eq!(depths[2], 0);
    assert_eq!(depths[3], 0);
}

#[test]
fn test_huffman_two_symbols() {
    // Two symbols should both get depth 1
    let freqs = [100u64, 50, 0, 0];
    let depths = build_code_lengths(&freqs, 16);
    assert_eq!(depths[0], 1);
    assert_eq!(depths[1], 1);
    assert_eq!(depths[2], 0);
    assert_eq!(depths[3], 0);
}

#[test]
fn test_huffman_tree_limit() {
    // Create frequencies that would naturally produce deep tree
    // Powers of 2: 1, 2, 4, 8, 16, 32, 64, 128, 256, 512
    let freqs: Vec<u64> = (0..20).map(|i| 1u64 << i).collect();

    let depths = build_code_lengths(&freqs, 16);

    // No depth should exceed tree_limit
    for &d in &depths {
        assert!(d <= 16, "Depth {} exceeds tree limit 16", d);
    }

    // All non-zero frequencies should have non-zero depth
    for (i, (&f, &d)) in freqs.iter().zip(depths.iter()).enumerate() {
        if f > 0 {
            assert!(d > 0, "Symbol {} has freq {} but depth 0", i, f);
        }
    }
}

/// Validate that generated codes form a valid prefix-free code.
#[test]
fn test_huffman_valid_code() {
    let freqs = [100u64, 50, 25, 10, 5, 1];
    let depths = build_code_lengths(&freqs, 16);

    // Kraft inequality: sum of 2^(-depth) should equal 1 for complete code
    let mut kraft_sum = 0.0f64;
    for &d in &depths {
        if d > 0 {
            kraft_sum += 2.0f64.powi(-(d as i32));
        }
    }

    // Should be exactly 1.0 for a complete code (allow small tolerance)
    assert!(
        (kraft_sum - 1.0).abs() < 1e-10,
        "Kraft sum {} != 1.0 - codes are invalid",
        kraft_sum
    );
}
