use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::IntoPyObjectExt;

// ── Helper: serde_json::Value → Python object ──────────────────────

fn json_value_to_py(py: Python<'_>, value: &serde_json::Value) -> PyResult<Py<PyAny>> {
    match value {
        serde_json::Value::Null => None::<bool>.into_py_any(py),
        serde_json::Value::Bool(b) => b.into_py_any(py),
        serde_json::Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                i.into_py_any(py)
            } else if let Some(f) = n.as_f64() {
                f.into_py_any(py)
            } else {
                None::<bool>.into_py_any(py)
            }
        }
        serde_json::Value::String(s) => s.into_py_any(py),
        serde_json::Value::Array(arr) => {
            let list: Vec<Py<PyAny>> = arr
                .iter()
                .map(|v| json_value_to_py(py, v))
                .collect::<PyResult<_>>()?;
            list.into_py_any(py)
        }
        serde_json::Value::Object(map) => {
            let dict = PyDict::new(py);
            for (k, v) in map {
                dict.set_item(k, json_value_to_py(py, v)?)?;
            }
            Ok(dict.unbind().into_any())
        }
    }
}

// ── Serde error translation ─────────────────────────────────────────

fn translate_serde_error(e: &serde_yaml::Error) -> Vec<String> {
    let msg = e.to_string();
    let location = e.location().map(|l| format!(" (line {})", l.line()));
    let loc = location.as_deref().unwrap_or("");

    // unknown variant 'X', expected one of ...
    if let Some(rest) = msg.strip_prefix("unknown variant `") {
        if let Some(tick_end) = rest.find('`') {
            let variant = &rest[..tick_end];
            let after = &rest[tick_end..];

            // Check for survivorship strategy context
            let survivorship_strategies = ["most_recent", "most_complete", "source_priority", "aggregate", "custom"];
            if survivorship_strategies.iter().any(|s| after.contains(s)) {
                return vec![format!(
                    "KNV-E201: Invalid survivorship strategy '{variant}'{loc}. Valid: {}",
                    survivorship_strategies.join(", ")
                )];
            }

            // Extract the expected list
            if let Some(exp_start) = after.find("expected one of ") {
                let expected = &after[exp_start + "expected one of ".len()..];
                let expected = expected.trim_end_matches(|c: char| c == '.' || c.is_whitespace());
                return vec![format!("KNV-E101: Invalid value '{variant}'{loc}. Valid options: {expected}")];
            }
        }
    }

    // missing field 'X'
    if msg.contains("missing field") {
        if let Some(start) = msg.find("missing field `") {
            let rest = &msg[start + "missing field `".len()..];
            if let Some(end) = rest.find('`') {
                let field = &rest[..end];
                // Check for survivorship override context
                if field == "field" && msg.contains("SurvivorshipOverride") {
                    return vec![format!(
                        "KNV-E203: Each survivorship override needs a 'field' name{loc}."
                    )];
                }
                return vec![format!(
                    "KNV-E102: Missing required field '{field}'{loc}. Add '{field}: <value>' to your spec."
                )];
            }
        }
    }

    // unknown field 'X', expected one of ...
    if let Some(rest) = msg.strip_prefix("unknown field `") {
        if let Some(tick_end) = rest.find('`') {
            let field = &rest[..tick_end];
            let after = &rest[tick_end..];

            // Try to find closest match for "did you mean" suggestion
            if let Some(exp_start) = after.find("expected one of ") {
                let expected_str = &after[exp_start + "expected one of ".len()..];
                let expected_str = expected_str.trim_end_matches(|c: char| c == '.' || c.is_whitespace());
                let candidates: Vec<&str> = expected_str.split(", ")
                    .map(|s| s.trim_matches('`').trim())
                    .collect();
                let suggestion = find_closest_match(field, &candidates);
                return match suggestion {
                    Some(s) => vec![format!("KNV-E103: Unknown field '{field}'{loc}. Did you mean '{s}'?")],
                    None => vec![format!("KNV-E103: Unknown field '{field}'{loc}. Expected one of: {expected_str}")],
                };
            }
        }
    }

    // invalid type: X, expected Y
    if msg.contains("invalid type:") {
        return vec![format!("KNV-E104: Type error{loc} — {msg}")];
    }

    // Fallback: pass through with a code prefix
    vec![format!("KNV-E100: {msg}")]
}

fn find_closest_match<'a>(target: &str, candidates: &[&'a str]) -> Option<&'a str> {
    let target_lower = target.to_lowercase();
    candidates.iter()
        .filter_map(|&c| {
            let c_lower = c.to_lowercase();
            let dist = edit_distance(&target_lower, &c_lower);
            if dist <= 2 { Some((c, dist)) } else { None }
        })
        .min_by_key(|(_, d)| *d)
        .map(|(c, _)| c)
}

fn edit_distance(a: &str, b: &str) -> usize {
    let a_chars: Vec<char> = a.chars().collect();
    let b_chars: Vec<char> = b.chars().collect();
    let m = a_chars.len();
    let n = b_chars.len();
    let mut dp = vec![vec![0usize; n + 1]; m + 1];
    for i in 0..=m { dp[i][0] = i; }
    for j in 0..=n { dp[0][j] = j; }
    for i in 1..=m {
        for j in 1..=n {
            let cost = if a_chars[i-1] == b_chars[j-1] { 0 } else { 1 };
            dp[i][j] = (dp[i-1][j] + 1).min(dp[i][j-1] + 1).min(dp[i-1][j-1] + cost);
        }
    }
    dp[m][n]
}

// ── PyO3 functions ─────────────────────────────────────────────────

#[pyfunction]
fn validate(yaml_str: &str) -> PyResult<Vec<String>> {
    kanoniv_core::validate_yaml(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[pyfunction]
fn validate_strict(yaml_str: &str) -> PyResult<Vec<String>> {
    // 1. Run existing permissive validator
    let mut errors = kanoniv_core::validate_yaml(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

    // 2. Attempt strict serde deserialization — catches what the permissive validator misses
    match serde_yaml::from_str::<cannon_common::spec::IdentitySpec>(yaml_str) {
        Ok(_) => {},
        Err(e) => errors.extend(translate_serde_error(&e)),
    }

    Ok(errors)
}

#[pyfunction]
fn validate_schema(yaml_str: &str) -> PyResult<Vec<String>> {
    let spec = kanoniv_core::parse_yaml(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    kanoniv_core::validate_schema(&spec)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[pyfunction]
fn validate_semantics(yaml_str: &str) -> PyResult<Vec<String>> {
    let spec = kanoniv_core::parse_yaml(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    kanoniv_core::validate_semantics(&spec)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[pyfunction]
fn parse(py: Python<'_>, yaml_str: &str) -> PyResult<Py<PyAny>> {
    let value = kanoniv_core::parse_yaml(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    json_value_to_py(py, &value)
}

#[pyfunction]
fn compile_ir(py: Python<'_>, yaml_str: &str) -> PyResult<Py<PyAny>> {
    let spec = kanoniv_core::parse_yaml(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    let ir = kanoniv_core::compile_to_ir(&spec)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    json_value_to_py(py, &ir)
}

#[pyfunction]
fn diff(py: Python<'_>, yaml_a: &str, yaml_b: &str) -> PyResult<Py<PyAny>> {
    let result = kanoniv_core::compute_diff(yaml_a, yaml_b)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    let value = serde_json::to_value(&result)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    json_value_to_py(py, &value)
}

#[pyfunction]
fn hash(yaml_str: &str) -> PyResult<String> {
    use sha2::Digest;
    let spec: serde_json::Value = serde_yaml::from_str(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    let canonical = serde_json::to_string(&spec)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    let mut hasher = sha2::Sha256::new();
    hasher.update(canonical.as_bytes());
    Ok(format!("sha256:{:x}", hasher.finalize()))
}

#[pyfunction]
fn plan(py: Python<'_>, yaml_str: &str) -> PyResult<Py<PyAny>> {
    let result = kanoniv_core::generate_plan(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    let value = serde_json::to_value(&result)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    json_value_to_py(py, &value)
}

// ── Local reconciliation ──────────────────────────────────────────

#[pyfunction]
fn reconcile_local(py: Python<'_>, yaml_str: &str, entities_json: &str) -> PyResult<Py<PyAny>> {
    use cannon_common::spec::IdentitySpec;
    use cannon_core::compiler::SpecCompiler;
    use cannon_core::engine::ReconciliationEngine;
    use cannon_core::types::NormalizedEntity;
    use cannon_core::overrides::OverrideResolver;
    use cannon_core::survivorship::SurvivorshipEngine;

    // 1. Parse spec YAML -> IdentitySpec
    let spec: IdentitySpec = serde_yaml::from_str(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid spec YAML: {e}")))?;

    // 2. Compile -> IdentityPlan
    let plan = SpecCompiler::compile(&spec)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Compilation error: {e}")))?;

    // 3. Build engine and run EM training if configured
    let mut engine = ReconciliationEngine::from_plan(&plan);

    // 4. Deserialize entities
    let entities: Vec<NormalizedEntity> = serde_json::from_str(entities_json)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid entities JSON: {e}")))?;

    // 5. Train FS parameters if EM training is configured
    engine.train_fellegi_sunter(&entities);

    // 6. Serialize trained FS parameters for persistence
    let trained_fs_params = engine.fellegi_sunter_config.as_ref()
        .and_then(|fs| serde_json::to_string(fs).ok());

    // 7. Release GIL for rayon parallelism, run reconciliation with telemetry
    let resolver = OverrideResolver::new(vec![]);
    let (decisions, telemetry) = py.detach(|| {
        engine.reconcile_with_telemetry(&entities, &resolver)
    });

    // 8. Cluster decisions
    let clusters = engine.cluster_decisions(&entities, &decisions);

    // 9. Golden records - with kanoniv_id and member_count
    let golden_records: Vec<std::collections::HashMap<String, String>> = clusters.iter().map(|cluster_ids| {
        let cluster_entities: Vec<NormalizedEntity> = cluster_ids.iter()
            .filter_map(|id| entities.iter().find(|e| &e.id == id).cloned())
            .collect();
        let mut golden = SurvivorshipEngine::construct_golden_record(&cluster_entities, &plan.survivorship);

        // Generate deterministic kanoniv_id from sorted external_ids (not UUIDs)
        use sha2::Digest;
        let mut sorted_ext_ids: Vec<String> = cluster_entities.iter()
            .map(|e| format!("{}:{}", e.source_name, e.external_id))
            .collect();
        sorted_ext_ids.sort();
        let mut hasher = sha2::Sha256::new();
        for ext_id in &sorted_ext_ids {
            hasher.update(ext_id.as_bytes());
            hasher.update(b"\0");
        }
        let hash_bytes = hasher.finalize();
        let kanoniv_id = format!("knv_{:x}", &hash_bytes).chars().take(16).collect::<String>();

        golden.insert("kanoniv_id".to_string(), kanoniv_id);
        golden.insert("member_count".to_string(), cluster_ids.len().to_string());
        golden
    }).collect();

    // 10. Serialize result to JSON -> Python dict
    let result = serde_json::json!({
        "clusters": clusters,
        "golden_records": golden_records,
        "decisions": decisions,
        "telemetry": telemetry,
        "trained_fs_params": trained_fs_params,
    });

    json_value_to_py(py, &result)
}

#[pyfunction]
fn reconcile_incremental(
    py: Python<'_>,
    yaml_str: &str,
    new_entities_json: &str,
    existing_entities_json: &str,
    existing_clusters_json: &str,
    trained_fs_params_json: Option<&str>,
) -> PyResult<Py<PyAny>> {
    use cannon_common::spec::IdentitySpec;
    use cannon_common::ir::CompiledFellegiSunterPlan;
    use cannon_core::compiler::SpecCompiler;
    use cannon_core::engine::ReconciliationEngine;
    use cannon_core::types::NormalizedEntity;
    use cannon_core::overrides::OverrideResolver;
    use cannon_core::survivorship::SurvivorshipEngine;

    // 1. Parse spec and build engine
    let spec: IdentitySpec = serde_yaml::from_str(yaml_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid spec YAML: {e}")))?;
    let plan = SpecCompiler::compile(&spec)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Compilation error: {e}")))?;
    let mut engine = ReconciliationEngine::from_plan(&plan);

    // 2. If trained FS params provided, restore them (skip EM training)
    if let Some(fs_json) = trained_fs_params_json {
        let fs: CompiledFellegiSunterPlan = serde_json::from_str(fs_json)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid trained_fs_params: {e}")))?;
        engine.fellegi_sunter_config = Some(fs);
    }

    // 3. Parse new and existing entities
    let new_entities: Vec<NormalizedEntity> = serde_json::from_str(new_entities_json)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid new entities JSON: {e}")))?;
    let existing_entities: Vec<NormalizedEntity> = serde_json::from_str(existing_entities_json)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid existing entities JSON: {e}")))?;

    // 4. Parse existing clusters (Vec<Vec<String>> of UUID strings -> Vec<Vec<Uuid>>)
    let clusters_raw: Vec<Vec<String>> = serde_json::from_str(existing_clusters_json)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid existing clusters JSON: {e}")))?;
    let existing_clusters: Vec<Vec<_>> = clusters_raw.iter()
        .map(|cluster| cluster.iter()
            .filter_map(|s| s.parse().ok())
            .collect::<Vec<_>>())
        .collect();

    // 5. Build set of existing entity IDs (using hashbrown::HashSet to match engine API)
    let existing_ids: hashbrown::HashSet<_> = existing_entities.iter().map(|e| e.id).collect();

    // 6. Combine all entities
    let mut all_entities = existing_entities;
    all_entities.extend(new_entities);

    // 7. If no trained FS params were provided, train on all entities
    if trained_fs_params_json.is_none() {
        engine.train_fellegi_sunter(&all_entities);
    }

    // 8. Serialize trained FS parameters for persistence
    let trained_fs_params = engine.fellegi_sunter_config.as_ref()
        .and_then(|fs| serde_json::to_string(fs).ok());

    // 9. Run incremental reconciliation (skip existing-vs-existing pairs)
    let resolver = OverrideResolver::new(vec![]);
    let (decisions, telemetry) = py.detach(|| {
        engine.reconcile_incremental_with_telemetry(&all_entities, &existing_ids, &resolver)
    });

    // 10. Cluster incrementally: pre-populate with existing clusters, apply new decisions
    let clusters = engine.cluster_decisions_incremental(&all_entities, &decisions, &existing_clusters);

    // 11. Golden records for all clusters
    let golden_records: Vec<std::collections::HashMap<String, String>> = clusters.iter().map(|cluster_ids| {
        let cluster_entities: Vec<NormalizedEntity> = cluster_ids.iter()
            .filter_map(|id| all_entities.iter().find(|e| &e.id == id).cloned())
            .collect();
        let mut golden = SurvivorshipEngine::construct_golden_record(&cluster_entities, &plan.survivorship);

        use sha2::Digest;
        let mut sorted_ext_ids: Vec<String> = cluster_entities.iter()
            .map(|e| format!("{}:{}", e.source_name, e.external_id))
            .collect();
        sorted_ext_ids.sort();
        let mut hasher = sha2::Sha256::new();
        for ext_id in &sorted_ext_ids {
            hasher.update(ext_id.as_bytes());
            hasher.update(b"\0");
        }
        let hash_bytes = hasher.finalize();
        let kanoniv_id = format!("knv_{:x}", &hash_bytes).chars().take(16).collect::<String>();

        golden.insert("kanoniv_id".to_string(), kanoniv_id);
        golden.insert("member_count".to_string(), cluster_ids.len().to_string());
        golden
    }).collect();

    // 12. Return result
    let result = serde_json::json!({
        "clusters": clusters,
        "golden_records": golden_records,
        "decisions": decisions,
        "telemetry": telemetry,
        "trained_fs_params": trained_fs_params,
    });

    json_value_to_py(py, &result)
}

// ── Module definition ──────────────────────────────────────────────

#[pymodule]
fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(validate, m)?)?;
    m.add_function(wrap_pyfunction!(validate_strict, m)?)?;
    m.add_function(wrap_pyfunction!(validate_schema, m)?)?;
    m.add_function(wrap_pyfunction!(validate_semantics, m)?)?;
    m.add_function(wrap_pyfunction!(parse, m)?)?;
    m.add_function(wrap_pyfunction!(compile_ir, m)?)?;
    m.add_function(wrap_pyfunction!(diff, m)?)?;
    m.add_function(wrap_pyfunction!(hash, m)?)?;
    m.add_function(wrap_pyfunction!(plan, m)?)?;
    m.add_function(wrap_pyfunction!(reconcile_local, m)?)?;
    m.add_function(wrap_pyfunction!(reconcile_incremental, m)?)?;
    Ok(())
}
