use pyo3::prelude::*;
use aho_corasick::AhoCorasick;
use std::collections::HashMap;

/// kNN search using brute-force L2 distance (O(n*d) — FAISS Rust bindings deferred).
/// Returns top-k (label, distance) pairs sorted by ascending distance.
#[pyfunction]
fn knn_search(
    query: Vec<f32>,
    index_vectors: Vec<Vec<f32>>,
    labels: Vec<String>,
    k: usize,
) -> PyResult<Vec<(String, f32)>> {
    if labels.len() < index_vectors.len() {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
            "labels length ({}) must be >= index_vectors length ({})",
            labels.len(),
            index_vectors.len()
        )));
    }
    if index_vectors.is_empty() || query.is_empty() {
        return Ok(vec![]);
    }
    let k = k.min(index_vectors.len());

    let mut scored: Vec<(usize, f32)> = index_vectors
        .iter()
        .enumerate()
        .map(|(i, v)| {
            let dist: f32 = query.iter().zip(v.iter()).map(|(a, b)| (a - b).powi(2)).sum();
            (i, dist)
        })
        .collect();

    scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));

    Ok(scored[..k]
        .iter()
        .map(|(i, d)| (labels[*i].clone(), *d))
        .collect())
}

/// Extract routing features from text.
/// Returns a dict of feature_name -> count.
#[pyfunction]
fn extract_features(text: &str) -> PyResult<HashMap<String, i64>> {
    let lower = text.to_lowercase();
    let mut features = HashMap::new();

    // Token count estimate (whitespace split)
    features.insert("estimated_tokens".to_string(), text.split_whitespace().count() as i64);

    // Code markers
    let code_markers = AhoCorasick::new(["```", "def ", "fn ", "class ", "import ", "function", "const ", "let ", "var "]).unwrap();
    features.insert("code_marker_count".to_string(), code_markers.find_iter(text).count() as i64);

    // Complexity indicators
    let complexity = AhoCorasick::new(["compare", "analyze", "multiple", "synthesis", "comprehensive", "trade-off", "pros and cons", "step-by-step"]).unwrap();
    features.insert("complexity_indicators".to_string(), complexity.find_iter(&lower).count() as i64);

    // Question detection
    features.insert("is_question".to_string(), if lower.contains('?') { 1 } else { 0 });

    // System prompt length indicator
    features.insert("has_system_context".to_string(), if text.len() > 200 { 1 } else { 0 });

    Ok(features)
}

/// Score confidence from extracted features.
/// Returns a float in [0.0, 1.0].
#[pyfunction]
fn score_confidence(features: HashMap<String, i64>) -> f32 {
    let complexity = features.get("complexity_indicators").copied().unwrap_or(0);
    let code_markers = features.get("code_marker_count").copied().unwrap_or(0);
    let tokens = features.get("estimated_tokens").copied().unwrap_or(0);

    // Base confidence scaled by signal strength
    let signal_count = (complexity > 0) as i64 + (code_markers > 0) as i64 + (tokens > 50) as i64;

    match signal_count {
        0 => 0.6,
        1 => 0.72,
        2 => 0.85,
        _ => 0.93,
    }
}

/// Classify text into privacy tier using Aho-Corasick pattern matching.
/// Returns "personal", "private", or "external".
#[pyfunction]
fn classify_privacy(text: &str) -> PyResult<String> {
    let lower = text.to_lowercase();

    // Personal PII patterns (match against lowercased text)
    let personal_patterns = AhoCorasick::new([
        "my name is", "my address", "my ssn", "my social security", "my date of birth",
        "my medical", "my health", "my diagnosis", "my credit card", "my bank account",
        "my password", "my private key", "my phone", "my email",
    ]).unwrap();

    if personal_patterns.find(&lower).is_some() {
        return Ok("personal".to_string());
    }

    // Private (internal/business) patterns
    let private_patterns = AhoCorasick::new([
        "internal", "confidential", "proprietary", "secret", "our company",
        "our revenue", "our customers", "our api key", "our database",
    ]).unwrap();

    if private_patterns.find(&lower).is_some() {
        return Ok("private".to_string());
    }

    Ok("external".to_string())
}

pub const RUST_VERSION: &str = env!("CARGO_PKG_VERSION");

/// Register router_core functions into a parent module.
pub fn register_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(knn_search, m)?)?;
    m.add_function(wrap_pyfunction!(extract_features, m)?)?;
    m.add_function(wrap_pyfunction!(score_confidence, m)?)?;
    m.add_function(wrap_pyfunction!(classify_privacy, m)?)?;
    m.add("RUST_VERSION", RUST_VERSION)?;
    Ok(())
}

/// Module name must match the lib.name in Cargo.toml
#[pymodule]
fn _router_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
    register_module(m)
}
