use numpy::{PyArray1, PyReadonlyArray2};
use pyo3::prelude::*;

use crate::similarity;

const VALID_STRING_METRICS: &[&str] = &[
    "jaro_winkler",
    "levenshtein",
    "damerau_levenshtein",
    "exact",
    "soundex",
    "jaccard",
    "cosine",
    "token_sort_ratio",
];

const VALID_NUMERIC_METRICS: &[&str] = &["numeric", "numeric_similarity", "date_distance"];

/// Compute a feature matrix for multiple comparisons in a single Rust call.
///
/// Accepts string columns as list-of-lists and numeric columns as a 2D numpy
/// array (shape: n_numeric_cols × n_pairs) to avoid Python list conversion overhead.
///
/// Parameters:
/// - string_columns_left/right: list of lists for string-based metrics
/// - string_metric_types: metric name per string column
/// - numeric_left/right: 2D numpy array (n_numeric_cols, n_pairs) for numeric metrics
/// - numeric_metric_types: metric name per numeric column
/// - n_pairs: total number of pairs
///
/// Returns: flat f64 array of shape (n_pairs * n_features,) — reshape in Python.
#[pyfunction]
pub fn compute_feature_matrix<'py>(
    py: Python<'py>,
    string_columns_left: Vec<Vec<String>>,
    string_columns_right: Vec<Vec<String>>,
    string_metric_types: Vec<String>,
    numeric_left: PyReadonlyArray2<f64>,
    numeric_right: PyReadonlyArray2<f64>,
    numeric_metric_types: Vec<String>,
    n_pairs: usize,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let n_string = string_columns_left.len();
    let nl = numeric_left.as_array();
    let nr = numeric_right.as_array();
    let n_numeric = nl.nrows();
    let n_features = n_string + n_numeric;

    if n_string != string_columns_right.len() || n_string != string_metric_types.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "string columns left/right/metrics must have the same length",
        ));
    }
    if n_numeric != nr.nrows() || n_numeric != numeric_metric_types.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "numeric columns left/right/metrics must have the same length",
        ));
    }
    if n_features == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "Must provide at least one comparison",
        ));
    }

    // Validate metric names before entering allow_threads
    for (i, metric) in string_metric_types.iter().enumerate() {
        if !VALID_STRING_METRICS.contains(&metric.as_str()) {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "Unknown string metric '{}' at index {}. Valid: {:?}",
                metric, i, VALID_STRING_METRICS
            )));
        }
    }
    for (i, metric) in numeric_metric_types.iter().enumerate() {
        if !VALID_NUMERIC_METRICS.contains(&metric.as_str()) {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "Unknown numeric metric '{}' at index {}. Valid: {:?}",
                metric, i, VALID_NUMERIC_METRICS
            )));
        }
    }

    // Validate string column lengths
    for i in 0..n_string {
        if string_columns_left[i].len() != n_pairs || string_columns_right[i].len() != n_pairs {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "String column {} has wrong length (expected {})",
                i, n_pairs
            )));
        }
    }
    // Validate numeric array widths
    if n_numeric > 0 && (nl.ncols() != n_pairs || nr.ncols() != n_pairs) {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "Numeric array has wrong number of columns (expected {})",
            n_pairs
        )));
    }

    // Ensure contiguous layout for numeric arrays so row slices work
    let nl_owned;
    let nl_ref = if nl.as_slice().is_some() {
        nl
    } else {
        nl_owned = nl.as_standard_layout().into_owned();
        nl_owned.view()
    };
    let nr_owned;
    let nr_ref = if nr.as_slice().is_some() {
        nr
    } else {
        nr_owned = nr.as_standard_layout().into_owned();
        nr_owned.view()
    };

    // Extract row slices (each row = one numeric column's data) before releasing GIL
    let nl_rows: Vec<&[f64]> = (0..n_numeric)
        .map(|i| nl_ref.row(i).to_slice().unwrap())
        .collect();
    let nr_rows: Vec<&[f64]> = (0..n_numeric)
        .map(|i| nr_ref.row(i).to_slice().unwrap())
        .collect();

    // Compute all features and build row-major output in one GIL-free block
    let data: Vec<f64> = py.allow_threads(|| {
        // Compute all feature columns
        let mut cols: Vec<Vec<f64>> = Vec::with_capacity(n_features);

        for i in 0..n_string {
            let left = &string_columns_left[i];
            let right = &string_columns_right[i];
            let col = match string_metric_types[i].as_str() {
                "jaro_winkler" => similarity::compute_jaro_winkler(left, right),
                "levenshtein" => similarity::compute_levenshtein(left, right),
                "damerau_levenshtein" => similarity::compute_damerau_levenshtein(left, right),
                "exact" => similarity::compute_exact(left, right),
                "soundex" => similarity::compute_soundex(left, right),
                "jaccard" => similarity::compute_jaccard(left, right),
                "cosine" => similarity::compute_cosine(left, right),
                "token_sort_ratio" => similarity::compute_token_sort_ratio(left, right),
                _ => unreachable!(), // validated above
            };
            cols.push(col);
        }

        for i in 0..n_numeric {
            let left = nl_rows[i];
            let right = nr_rows[i];
            let col = match numeric_metric_types[i].as_str() {
                "numeric" => similarity::compute_numeric_distance(left, right),
                "numeric_similarity" => similarity::compute_numeric_similarity(left, right),
                "date_distance" => similarity::compute_numeric_distance(left, right),
                _ => unreachable!(), // validated above
            };
            cols.push(col);
        }

        // Build flat row-major output buffer
        let mut data = vec![0.0f64; n_pairs * n_features];
        for (col_idx, col) in cols.iter().enumerate() {
            for (row_idx, &val) in col.iter().enumerate() {
                data[row_idx * n_features + col_idx] = val;
            }
        }
        data
    });

    Ok(PyArray1::from_vec(py, data))
}
