use pyo3::prelude::*;
use rayon::prelude::*;
use std::collections::{HashMap, HashSet};
use strsim;

// ============================================================================
// Internal compute functions (called from features.rs, no PyO3 overhead)
// ============================================================================

pub fn compute_jaro_winkler(left: &[String], right: &[String]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| strsim::jaro_winkler(l, r))
        .collect()
}

pub fn compute_levenshtein(left: &[String], right: &[String]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| strsim::normalized_levenshtein(l, r))
        .collect()
}

pub fn compute_damerau_levenshtein(left: &[String], right: &[String]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| {
            let dist = strsim::damerau_levenshtein(l, r);
            let max_len = l.chars().count().max(r.chars().count());
            if max_len == 0 {
                1.0
            } else {
                1.0 - (dist as f64 / max_len as f64)
            }
        })
        .collect()
}

pub fn compute_exact(left: &[String], right: &[String]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| if l == r { 1.0 } else { 0.0 })
        .collect()
}

pub fn compute_numeric_distance(left: &[f64], right: &[f64]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| (l - r).abs())
        .collect()
}

pub fn compute_numeric_similarity(left: &[f64], right: &[f64]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| {
            let max_abs = l.abs().max(r.abs());
            if max_abs == 0.0 {
                1.0
            } else {
                (1.0 - (l - r).abs() / max_abs).max(0.0)
            }
        })
        .collect()
}

/// Soundex code for an uppercase ASCII byte.
#[inline]
fn soundex_code(c: u8) -> u8 {
    match c {
        b'B' | b'F' | b'P' | b'V' => b'1',
        b'C' | b'G' | b'J' | b'K' | b'Q' | b'S' | b'X' | b'Z' => b'2',
        b'D' | b'T' => b'3',
        b'L' => b'4',
        b'M' | b'N' => b'5',
        b'R' => b'6',
        _ => b'0', // A, E, I, O, U, H, W, Y
    }
}

/// Zero-allocation soundex encoding into a fixed [u8; 4] buffer.
fn soundex_encode_buf(s: &str) -> [u8; 4] {
    let mut result = [b'0'; 4];
    let mut pos = 0usize;
    let mut last_code = b'0';

    for &c in s.as_bytes() {
        if !c.is_ascii_alphabetic() {
            continue;
        }
        let upper = c.to_ascii_uppercase();

        if pos == 0 {
            result[0] = upper;
            pos = 1;
            last_code = soundex_code(upper);
            continue;
        }

        let code = soundex_code(upper);
        if code != b'0' && code != last_code {
            result[pos] = code;
            pos += 1;
            if pos == 4 {
                break;
            }
        }
        last_code = code;
    }

    result
}

pub fn compute_soundex(left: &[String], right: &[String]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| {
            if soundex_encode_buf(l) == soundex_encode_buf(r) {
                1.0
            } else {
                0.0
            }
        })
        .collect()
}

/// Token-level Jaccard similarity: |intersection| / |union| of whitespace-split tokens.
pub fn compute_jaccard(left: &[String], right: &[String]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| {
            let set_l: HashSet<&str> = l.split_whitespace().collect();
            let set_r: HashSet<&str> = r.split_whitespace().collect();
            if set_l.is_empty() && set_r.is_empty() {
                return 1.0;
            }
            let intersection = set_l.intersection(&set_r).count() as f64;
            let union = set_l.union(&set_r).count() as f64;
            if union == 0.0 { 1.0 } else { intersection / union }
        })
        .collect()
}

/// Token-level cosine similarity using term frequency vectors.
pub fn compute_cosine(left: &[String], right: &[String]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| {
            let tf_l = token_freq(l);
            let tf_r = token_freq(r);
            if tf_l.is_empty() && tf_r.is_empty() {
                return 1.0;
            }
            if tf_l.is_empty() || tf_r.is_empty() {
                return 0.0;
            }

            // Accumulate norms and dot product without building a union HashSet.
            // norm_r is computed by iterating tf_r directly; norm_l and dot
            // are computed by iterating tf_l and probing tf_r.
            let mut norm_r = 0.0f64;
            for &vr in tf_r.values() {
                norm_r += vr * vr;
            }

            let mut dot = 0.0f64;
            let mut norm_l = 0.0f64;
            for (key, &vl) in &tf_l {
                norm_l += vl * vl;
                if let Some(&vr) = tf_r.get(key) {
                    dot += vl * vr;
                }
            }

            let denom = norm_l.sqrt() * norm_r.sqrt();
            if denom == 0.0 { 0.0 } else { dot / denom }
        })
        .collect()
}

fn token_freq(s: &str) -> HashMap<&str, f64> {
    let mut map = HashMap::new();
    for token in s.split_whitespace() {
        *map.entry(token).or_insert(0.0) += 1.0;
    }
    map
}

/// Sort tokens alphabetically, then compute normalized Levenshtein.
/// Handles word-order variation: "Smith John" vs "John Smith" → 1.0.
pub fn compute_token_sort_ratio(left: &[String], right: &[String]) -> Vec<f64> {
    left.par_iter()
        .zip(right.par_iter())
        .map(|(l, r)| {
            let mut toks_l: Vec<&str> = l.split_whitespace().collect();
            let mut toks_r: Vec<&str> = r.split_whitespace().collect();

            // Fast path: 0-1 tokens need no sorting or joining
            if toks_l.len() <= 1 && toks_r.len() <= 1 {
                let tl = toks_l.first().copied().unwrap_or("");
                let tr = toks_r.first().copied().unwrap_or("");
                return strsim::normalized_levenshtein(tl, tr);
            }

            toks_l.sort_unstable();
            toks_r.sort_unstable();
            let sorted_l = toks_l.join(" ");
            let sorted_r = toks_r.join(" ");
            strsim::normalized_levenshtein(&sorted_l, &sorted_r)
        })
        .collect()
}

// ============================================================================
// PyO3-exported batch functions (for direct Python use)
// ============================================================================

#[pyfunction]
pub fn jaro_winkler_batch(py: Python<'_>, left: Vec<String>, right: Vec<String>) -> PyResult<Vec<f64>> {
    validate_lengths(&left, &right)?;
    py.allow_threads(|| Ok(compute_jaro_winkler(&left, &right)))
}

#[pyfunction]
pub fn levenshtein_batch(py: Python<'_>, left: Vec<String>, right: Vec<String>) -> PyResult<Vec<f64>> {
    validate_lengths(&left, &right)?;
    py.allow_threads(|| Ok(compute_levenshtein(&left, &right)))
}

#[pyfunction]
pub fn damerau_levenshtein_batch(py: Python<'_>, left: Vec<String>, right: Vec<String>) -> PyResult<Vec<f64>> {
    validate_lengths(&left, &right)?;
    py.allow_threads(|| Ok(compute_damerau_levenshtein(&left, &right)))
}

#[pyfunction]
pub fn exact_match_batch(py: Python<'_>, left: Vec<String>, right: Vec<String>) -> PyResult<Vec<f64>> {
    validate_lengths(&left, &right)?;
    py.allow_threads(|| Ok(compute_exact(&left, &right)))
}

#[pyfunction]
pub fn numeric_distance_batch(py: Python<'_>, left: Vec<f64>, right: Vec<f64>) -> PyResult<Vec<f64>> {
    validate_lengths_f64(&left, &right)?;
    py.allow_threads(|| Ok(compute_numeric_distance(&left, &right)))
}

#[pyfunction]
pub fn numeric_similarity_batch(py: Python<'_>, left: Vec<f64>, right: Vec<f64>) -> PyResult<Vec<f64>> {
    validate_lengths_f64(&left, &right)?;
    py.allow_threads(|| Ok(compute_numeric_similarity(&left, &right)))
}

#[pyfunction]
pub fn soundex_batch(py: Python<'_>, left: Vec<String>, right: Vec<String>) -> PyResult<Vec<f64>> {
    validate_lengths(&left, &right)?;
    py.allow_threads(|| Ok(compute_soundex(&left, &right)))
}

#[pyfunction]
pub fn jaccard_batch(py: Python<'_>, left: Vec<String>, right: Vec<String>) -> PyResult<Vec<f64>> {
    validate_lengths(&left, &right)?;
    py.allow_threads(|| Ok(compute_jaccard(&left, &right)))
}

#[pyfunction]
pub fn cosine_batch(py: Python<'_>, left: Vec<String>, right: Vec<String>) -> PyResult<Vec<f64>> {
    validate_lengths(&left, &right)?;
    py.allow_threads(|| Ok(compute_cosine(&left, &right)))
}

#[pyfunction]
pub fn token_sort_ratio_batch(py: Python<'_>, left: Vec<String>, right: Vec<String>) -> PyResult<Vec<f64>> {
    validate_lengths(&left, &right)?;
    py.allow_threads(|| Ok(compute_token_sort_ratio(&left, &right)))
}

fn validate_lengths(left: &[String], right: &[String]) -> PyResult<()> {
    if left.len() != right.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "left and right must have the same length",
        ));
    }
    Ok(())
}

fn validate_lengths_f64(left: &[f64], right: &[f64]) -> PyResult<()> {
    if left.len() != right.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "left and right must have the same length",
        ));
    }
    Ok(())
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_jaro_winkler_identical() {
        assert!((strsim::jaro_winkler("hello", "hello") - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_jaro_winkler_similar() {
        assert!(strsim::jaro_winkler("john", "jon") > 0.8);
    }

    #[test]
    fn test_normalized_levenshtein_identical() {
        assert!((strsim::normalized_levenshtein("hello", "hello") - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_normalized_levenshtein_different() {
        assert!(strsim::normalized_levenshtein("abc", "xyz") < 0.5);
    }

    #[test]
    fn test_damerau_levenshtein_transposition() {
        // "ab" vs "ba" — 1 transposition, max_len=2, normalized = 1 - 1/2 = 0.5
        let scores = compute_damerau_levenshtein(
            &["ab".to_string()],
            &["ba".to_string()],
        );
        assert!((scores[0] - 0.5).abs() < 1e-10);
    }

    #[test]
    fn test_soundex_encoding() {
        assert_eq!(&soundex_encode_buf("Robert"), b"R163");
        assert_eq!(&soundex_encode_buf("Rupert"), b"R163");
        assert_eq!(&soundex_encode_buf("Smith"), b"S530");
        assert_eq!(&soundex_encode_buf("Smythe"), b"S530");
    }

    #[test]
    fn test_soundex_batch() {
        let scores = compute_soundex(
            &["Robert".to_string(), "Smith".to_string()],
            &["Rupert".to_string(), "Jones".to_string()],
        );
        assert_eq!(scores[0], 1.0); // Same soundex
        assert_eq!(scores[1], 0.0); // Different soundex
    }

    #[test]
    fn test_jaccard_identical() {
        let scores = compute_jaccard(
            &["hello world".to_string()],
            &["hello world".to_string()],
        );
        assert!((scores[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_jaccard_partial() {
        let scores = compute_jaccard(
            &["hello world foo".to_string()],
            &["hello world bar".to_string()],
        );
        // intersection = {hello, world} = 2, union = {hello, world, foo, bar} = 4
        assert!((scores[0] - 0.5).abs() < 1e-10);
    }

    #[test]
    fn test_cosine_identical() {
        let scores = compute_cosine(
            &["hello world".to_string()],
            &["hello world".to_string()],
        );
        assert!((scores[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_cosine_orthogonal() {
        let scores = compute_cosine(
            &["aaa".to_string()],
            &["bbb".to_string()],
        );
        assert!(scores[0] < 0.01);
    }

    #[test]
    fn test_token_sort_ratio_reorder() {
        let scores = compute_token_sort_ratio(
            &["John Smith".to_string()],
            &["Smith John".to_string()],
        );
        assert!((scores[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_numeric_similarity() {
        let scores = compute_numeric_similarity(&[10.0], &[10.0]);
        assert!((scores[0] - 1.0).abs() < 1e-10);

        let scores = compute_numeric_similarity(&[10.0], &[0.0]);
        assert!((scores[0] - 0.0).abs() < 1e-10);

        let scores = compute_numeric_similarity(&[10.0], &[8.0]);
        // 1 - 2/10 = 0.8
        assert!((scores[0] - 0.8).abs() < 1e-10);
    }
}
