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)
}

// ── Active learning ────────────────────────────────────────────────

/// Extract uncertain pairs (Review decisions) sorted by uncertainty.
/// Returns top N pairs with both records' data and confidence scores.
#[pyfunction]
fn get_uncertain_pairs(
    py: Python<'_>,
    decisions_json: &str,
    entities_json: &str,
    n: usize,
    scoring_method: &str,
    merge_threshold: f64,
) -> PyResult<Py<PyAny>> {
    use cannon_core::types::{MatchDecision, NormalizedEntity};

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

    // Use a string-keyed map to avoid uuid dependency
    let entity_map: std::collections::HashMap<String, &NormalizedEntity> =
        entities.iter().map(|e| (e.id.to_string(), e)).collect();

    // Filter to Review decisions and decisions near the threshold
    let mut uncertain: Vec<&MatchDecision> = decisions.iter()
        .filter(|d| {
            use cannon_core::types::Decision;
            d.decision == Decision::Review
                || (d.decision == Decision::Merge
                    && d.confidence < merge_threshold * 1.1)
                || (d.decision == Decision::NoMerge
                    && d.confidence > merge_threshold * 0.9)
        })
        .collect();

    // Sort by uncertainty: closest to the decision boundary
    if scoring_method == "fellegi_sunter" {
        // For FS: use posterior_probability, most uncertain = closest to 0.5
        uncertain.sort_by(|a, b| {
            let a_unc = (a.posterior_probability.unwrap_or(0.5) - 0.5).abs();
            let b_unc = (b.posterior_probability.unwrap_or(0.5) - 0.5).abs();
            a_unc.partial_cmp(&b_unc).unwrap_or(std::cmp::Ordering::Equal)
        });
    } else {
        // For other methods: closest to merge_threshold
        uncertain.sort_by(|a, b| {
            let a_unc = (a.confidence - merge_threshold).abs();
            let b_unc = (b.confidence - merge_threshold).abs();
            a_unc.partial_cmp(&b_unc).unwrap_or(std::cmp::Ordering::Equal)
        });
    }

    uncertain.truncate(n);

    // Build output with both records' data
    let pairs: Vec<serde_json::Value> = uncertain.iter().map(|d| {
        let a_key = d.entity_a_id.to_string();
        let b_key = d.entity_b_id.to_string();
        let a_data = entity_map.get(&a_key)
            .map(|e| serde_json::to_value(&e.data).unwrap_or_default())
            .unwrap_or_default();
        let b_data = entity_map.get(&b_key)
            .map(|e| serde_json::to_value(&e.data).unwrap_or_default())
            .unwrap_or_default();
        let a_ext_id = entity_map.get(&a_key).map(|e| e.external_id.as_str()).unwrap_or("");
        let b_ext_id = entity_map.get(&b_key).map(|e| e.external_id.as_str()).unwrap_or("");
        let a_source = entity_map.get(&a_key).map(|e| e.source_name.as_str()).unwrap_or("");
        let b_source = entity_map.get(&b_key).map(|e| e.source_name.as_str()).unwrap_or("");

        serde_json::json!({
            "entity_a_id": d.entity_a_id.to_string(),
            "entity_b_id": d.entity_b_id.to_string(),
            "a_external_id": a_ext_id,
            "b_external_id": b_ext_id,
            "a_source": a_source,
            "b_source": b_source,
            "a": a_data,
            "b": b_data,
            "confidence": d.confidence,
            "posterior_probability": d.posterior_probability,
            "decision": serde_json::to_value(&d.decision).unwrap_or_default(),
            "rule_results": d.rule_results.iter().map(|r| serde_json::json!({
                "rule_name": r.rule_name,
                "score": r.score,
                "weight": r.weight,
                "matched": r.matched,
            })).collect::<Vec<_>>(),
        })
    }).collect();

    let result = serde_json::Value::Array(pairs);
    json_value_to_py(py, &result)
}

/// Reconcile with user feedback (active learning).
/// Restores trained FS params, applies supervised EM with feedback labels, then reconciles.
#[pyfunction]
fn reconcile_with_feedback(
    py: Python<'_>,
    yaml_str: &str,
    entities_json: &str,
    feedback_json: &str,
    trained_fs_params_json: Option<&str>,
    learning_rate: f64,
) -> PyResult<Py<PyAny>> {
    use cannon_common::spec::IdentitySpec;
    use cannon_common::ir::{CompiledFellegiSunterPlan, FeedbackLabel, FeedbackDecision};
    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 entities
    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);

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

    // 2. Restore trained FS params if provided
    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);
    } else {
        // Train from scratch if no prior params
        engine.train_fellegi_sunter(&entities);
    }

    // 3. Parse feedback labels
    let feedback: Vec<FeedbackLabel> = serde_json::from_str(feedback_json)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid feedback JSON: {e}")))?;

    // 4. If FS scoring and we have feedback, run supervised EM
    if engine.fellegi_sunter_config.is_some() && !feedback.is_empty() {
        // Build entity lookup by (source_name, external_id)
        let entity_lookup: std::collections::HashMap<(String, String), &NormalizedEntity> =
            entities.iter().map(|e| ((e.source_name.clone(), e.external_id.clone()), e)).collect();

        let mut labeled_pairs: Vec<(NormalizedEntity, NormalizedEntity, bool)> = Vec::new();
        for label in &feedback {
            let a = entity_lookup.get(&(label.source_a.clone(), label.entity_a_id.clone()));
            let b = entity_lookup.get(&(label.source_b.clone(), label.entity_b_id.clone()));
            if let (Some(a), Some(b)) = (a, b) {
                let is_match = label.label == FeedbackDecision::Match;
                labeled_pairs.push(((*a).clone(), (*b).clone(), is_match));
            }
        }

        if !labeled_pairs.is_empty() {
            engine.train_fellegi_sunter_supervised(&labeled_pairs, learning_rate);
        }
    }

    // 5. For non-FS scoring with feedback, apply as overrides
    let mut overrides = Vec::new();
    if engine.fellegi_sunter_config.is_none() {
        let entity_lookup: std::collections::HashMap<(String, String), &NormalizedEntity> =
            entities.iter().map(|e| ((e.source_name.clone(), e.external_id.clone()), e)).collect();

        for label in &feedback {
            let a = entity_lookup.get(&(label.source_a.clone(), label.entity_a_id.clone()));
            let b = entity_lookup.get(&(label.source_b.clone(), label.entity_b_id.clone()));
            if let (Some(a), Some(b)) = (a, b) {
                use cannon_core::overrides::{ManualOverride, OverrideType};
                let override_type = match label.label {
                    FeedbackDecision::Match => OverrideType::Merge,
                    FeedbackDecision::NoMatch => OverrideType::Split,
                };
                overrides.push(ManualOverride {
                    tenant_id: uuid::Uuid::nil(),
                    entity_a_id: Some(a.id),
                    entity_b_id: Some(b.id),
                    override_type,
                    canonical_id: None,
                    override_data: serde_json::Value::Null,
                });
            }
        }
    }

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

    // 7. Run reconciliation
    let resolver = OverrideResolver::new(overrides);
    let (decisions, telemetry) = py.detach(|| {
        engine.reconcile_with_telemetry(&entities, &resolver)
    });

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

    // 9. Golden records
    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);

        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. 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)?)?;
    m.add_function(wrap_pyfunction!(get_uncertain_pairs, m)?)?;
    m.add_function(wrap_pyfunction!(reconcile_with_feedback, m)?)?;
    Ok(())
}
