//! Transitive Intelligence: compute soft relationship edges between entities
//! in different merge clusters based on inverse-frequency scoring of shared
//! attribute values.

use cannon_common::relationship::{RelationshipConfig, RelationshipEdge, SharedAttribute};
use std::collections::HashMap;
use uuid::Uuid;

/// Compute relationship edges between entities in different merge clusters.
///
/// Algorithm:
/// 1. Build entity_id -> cluster_id map from clusters
/// 2. Build global frequency map scanning all entity field values
/// 3. Build inverted index mapping (field, value) to entity IDs
/// 4. For each (field, value) with 2+ entities in DIFFERENT clusters:
///    - Skip if frequency > config.max_value_frequency (edge explosion guard)
///    - contribution = 1.0 / (1.0 + log2(frequency))
///    - Emit candidate edges between cross-cluster entity pairs
/// 5. Aggregate contributions per entity pair, sum weights, collect attributes
/// 6. Filter weight >= config.min_edge_weight
/// 7. Sort by (entity_a_id, entity_b_id) for determinism, enforce ordered pair (a < b)
/// 8. Truncate to config.max_edges (keep highest weight)
pub fn compute_relationship_edges(
    entities: &[(Uuid, HashMap<String, String>)],
    clusters: &[Vec<Uuid>],
    bridge_fields: &[String],
    config: &RelationshipConfig,
    entity_to_canonical: &HashMap<Uuid, Uuid>,
) -> Vec<RelationshipEdge> {
    if entities.is_empty() || clusters.is_empty() || bridge_fields.is_empty() {
        return Vec::new();
    }

    // 1. Build entity_id -> cluster_id map
    let mut entity_cluster: HashMap<Uuid, usize> = HashMap::new();
    for (cluster_idx, cluster) in clusters.iter().enumerate() {
        for &entity_id in cluster {
            entity_cluster.insert(entity_id, cluster_idx);
        }
    }

    // 2. Build global frequency map: (field, value) -> count
    let mut frequency_map: HashMap<(String, String), usize> = HashMap::new();

    // 3. Build inverted index: (field, value) -> Vec<entity_id>
    let mut inverted_index: HashMap<(String, String), Vec<Uuid>> = HashMap::new();

    let bridge_set: std::collections::HashSet<&str> =
        bridge_fields.iter().map(|s| s.as_str()).collect();

    for (entity_id, fields) in entities {
        for (field_name, field_value) in fields {
            if !bridge_set.contains(field_name.as_str()) {
                continue;
            }
            if field_value.is_empty() {
                continue;
            }
            let key = (field_name.clone(), field_value.clone());
            *frequency_map.entry(key.clone()).or_insert(0) += 1;
            inverted_index.entry(key).or_default().push(*entity_id);
        }
    }

    // 4 + 5. For each (field, value) with 2+ entities in different clusters,
    //         compute candidate edges and aggregate contributions.
    let mut edge_map: HashMap<(Uuid, Uuid), (f64, Vec<SharedAttribute>)> = HashMap::new();

    for ((field, value), entity_ids) in &inverted_index {
        if entity_ids.len() < 2 {
            continue;
        }

        let frequency = *frequency_map.get(&(field.clone(), value.clone())).unwrap_or(&0);
        if frequency > config.max_value_frequency {
            continue;
        }

        let contribution = 1.0 / (1.0 + (frequency as f64).log2());

        // Check for cross-cluster pairs
        for (i, &id_a) in entity_ids.iter().enumerate() {
            let cluster_a = entity_cluster.get(&id_a);

            for &id_b in &entity_ids[i + 1..] {
                let cluster_b = entity_cluster.get(&id_b);

                // Only emit edges between entities in DIFFERENT clusters
                match (cluster_a, cluster_b) {
                    (Some(ca), Some(cb)) if ca != cb => {}
                    _ => continue,
                }

                // Enforce ordered pair (a < b)
                let (lo, hi) = if id_a < id_b {
                    (id_a, id_b)
                } else {
                    (id_b, id_a)
                };

                let entry = edge_map.entry((lo, hi)).or_insert_with(|| (0.0, Vec::new()));
                entry.0 += contribution;
                entry.1.push(SharedAttribute {
                    field: field.clone(),
                    value: value.clone(),
                    frequency,
                    contribution,
                });
            }
        }
    }

    // 6. Filter by min_edge_weight
    let mut edges: Vec<RelationshipEdge> = edge_map
        .into_iter()
        .filter(|(_, (weight, _))| *weight >= config.min_edge_weight)
        .map(|((a, b), (weight, attributes))| {
            let canonical_a = entity_to_canonical.get(&a).copied().unwrap_or(a);
            let canonical_b = entity_to_canonical.get(&b).copied().unwrap_or(b);
            RelationshipEdge {
                entity_a_id: a,
                entity_b_id: b,
                canonical_a_id: canonical_a,
                canonical_b_id: canonical_b,
                weight,
                attributes,
            }
        })
        .collect();

    // 7. Sort by weight descending for truncation, then by pair for determinism
    edges.sort_by(|a, b| {
        b.weight
            .partial_cmp(&a.weight)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| a.entity_a_id.cmp(&b.entity_a_id))
            .then_with(|| a.entity_b_id.cmp(&b.entity_b_id))
    });

    // 8. Truncate to max_edges
    edges.truncate(config.max_edges);

    // Final sort by (entity_a_id, entity_b_id) for deterministic output
    edges.sort_by(|a, b| {
        a.entity_a_id
            .cmp(&b.entity_a_id)
            .then_with(|| a.entity_b_id.cmp(&b.entity_b_id))
    });

    edges
}

/// Extract the global frequency map from entity data for a set of bridge fields.
/// Returns (field_name, field_value) -> frequency.
pub fn build_frequency_map(
    entities: &[(Uuid, HashMap<String, String>)],
    bridge_fields: &[String],
) -> HashMap<(String, String), usize> {
    let mut freq: HashMap<(String, String), usize> = HashMap::new();
    let bridge_set: std::collections::HashSet<&str> =
        bridge_fields.iter().map(|s| s.as_str()).collect();

    for (_entity_id, fields) in entities {
        for (field_name, field_value) in fields {
            if !bridge_set.contains(field_name.as_str()) || field_value.is_empty() {
                continue;
            }
            *freq
                .entry((field_name.clone(), field_value.clone()))
                .or_insert(0) += 1;
        }
    }
    freq
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn make_entity(id: Uuid, fields: Vec<(&str, &str)>) -> (Uuid, HashMap<String, String>) {
        let map: HashMap<String, String> = fields
            .into_iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect();
        (id, map)
    }

    fn uuid_from_u128(val: u128) -> Uuid {
        Uuid::from_u128(val)
    }

    #[test]
    fn test_empty_input_returns_empty() {
        let config = RelationshipConfig::default();
        let canonical_map = HashMap::new();
        let result = compute_relationship_edges(&[], &[], &[], &config, &canonical_map);
        assert!(result.is_empty());
    }

    #[test]
    fn test_same_cluster_produces_no_edges() {
        let id_a = uuid_from_u128(1);
        let id_b = uuid_from_u128(2);

        let entities = vec![
            make_entity(id_a, vec![("company", "Acme")]),
            make_entity(id_b, vec![("company", "Acme")]),
        ];

        // Both entities in the same cluster
        let clusters = vec![vec![id_a, id_b]];
        let bridge_fields = vec!["company".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            ..Default::default()
        };

        let mut canonical_map = HashMap::new();
        canonical_map.insert(id_a, uuid_from_u128(100));
        canonical_map.insert(id_b, uuid_from_u128(100));

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert!(edges.is_empty(), "Same-cluster entities should produce no edges");
    }

    #[test]
    fn test_rare_shared_value_produces_high_contribution() {
        let id_a = uuid_from_u128(1);
        let id_b = uuid_from_u128(2);

        let entities = vec![
            make_entity(id_a, vec![("company", "RareCorp")]),
            make_entity(id_b, vec![("company", "RareCorp")]),
        ];

        // Different clusters
        let clusters = vec![vec![id_a], vec![id_b]];
        let bridge_fields = vec!["company".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            ..Default::default()
        };

        let mut canonical_map = HashMap::new();
        canonical_map.insert(id_a, uuid_from_u128(100));
        canonical_map.insert(id_b, uuid_from_u128(200));

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert_eq!(edges.len(), 1);

        // frequency = 2, contribution = 1.0 / (1.0 + log2(2)) = 1.0 / 2.0 = 0.5
        let expected_contribution = 1.0 / (1.0 + 2.0_f64.log2());
        assert!(
            (edges[0].weight - expected_contribution).abs() < 0.001,
            "Rare value should produce contribution of ~0.5, got {}",
            edges[0].weight
        );
    }

    #[test]
    fn test_common_shared_value_produces_low_contribution() {
        // 10 entities all sharing the same company value
        let ids: Vec<Uuid> = (1..=10).map(uuid_from_u128).collect();

        let entities: Vec<(Uuid, HashMap<String, String>)> = ids
            .iter()
            .map(|id| make_entity(*id, vec![("company", "CommonCorp")]))
            .collect();

        // Each entity in its own cluster
        let clusters: Vec<Vec<Uuid>> = ids.iter().map(|id| vec![*id]).collect();
        let bridge_fields = vec!["company".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            ..Default::default()
        };

        let canonical_map: HashMap<Uuid, Uuid> = ids
            .iter()
            .enumerate()
            .map(|(i, id)| (*id, uuid_from_u128(100 + i as u128)))
            .collect();

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);

        // frequency = 10, contribution = 1.0 / (1.0 + log2(10)) ~ 0.232
        let expected_contribution = 1.0 / (1.0 + 10.0_f64.log2());
        for edge in &edges {
            assert!(
                (edge.weight - expected_contribution).abs() < 0.001,
                "Common value should produce low contribution, got {}",
                edge.weight
            );
        }
    }

    #[test]
    fn test_frequency_over_max_is_skipped() {
        let id_a = uuid_from_u128(1);
        let id_b = uuid_from_u128(2);

        // Two entities sharing a value, but we set max_value_frequency = 1
        // so frequency 2 will be skipped
        let entities = vec![
            make_entity(id_a, vec![("company", "SkipCorp")]),
            make_entity(id_b, vec![("company", "SkipCorp")]),
        ];

        let clusters = vec![vec![id_a], vec![id_b]];
        let bridge_fields = vec!["company".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            max_value_frequency: 1,
            ..Default::default()
        };

        let mut canonical_map = HashMap::new();
        canonical_map.insert(id_a, uuid_from_u128(100));
        canonical_map.insert(id_b, uuid_from_u128(200));

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert!(
            edges.is_empty(),
            "Values exceeding max_value_frequency should be skipped"
        );
    }

    #[test]
    fn test_ordered_pair_constraint() {
        // Ensure entity_a_id < entity_b_id regardless of input order
        let id_a = uuid_from_u128(100); // higher
        let id_b = uuid_from_u128(1); // lower

        let entities = vec![
            make_entity(id_a, vec![("company", "TestCo")]),
            make_entity(id_b, vec![("company", "TestCo")]),
        ];

        let clusters = vec![vec![id_a], vec![id_b]];
        let bridge_fields = vec!["company".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            ..Default::default()
        };

        let mut canonical_map = HashMap::new();
        canonical_map.insert(id_a, uuid_from_u128(300));
        canonical_map.insert(id_b, uuid_from_u128(400));

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert_eq!(edges.len(), 1);
        assert!(
            edges[0].entity_a_id < edges[0].entity_b_id,
            "Entity A must be less than Entity B in the pair"
        );
    }

    #[test]
    fn test_max_edges_cap() {
        // Create enough entities to produce many edges, then cap at 2
        let ids: Vec<Uuid> = (1..=5).map(uuid_from_u128).collect();

        let entities: Vec<(Uuid, HashMap<String, String>)> = ids
            .iter()
            .map(|id| make_entity(*id, vec![("domain", "example.com")]))
            .collect();

        let clusters: Vec<Vec<Uuid>> = ids.iter().map(|id| vec![*id]).collect();
        let bridge_fields = vec!["domain".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            max_edges: 2,
            ..Default::default()
        };

        let canonical_map: HashMap<Uuid, Uuid> = ids
            .iter()
            .enumerate()
            .map(|(i, id)| (*id, uuid_from_u128(100 + i as u128)))
            .collect();

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert!(
            edges.len() <= 2,
            "Should not exceed max_edges cap of 2, got {}",
            edges.len()
        );
    }

    #[test]
    fn test_multiple_shared_attributes_sum_weights() {
        let id_a = uuid_from_u128(1);
        let id_b = uuid_from_u128(2);

        // Two entities sharing both company and domain
        let entities = vec![
            make_entity(
                id_a,
                vec![("company", "SharedCo"), ("domain", "shared.com")],
            ),
            make_entity(
                id_b,
                vec![("company", "SharedCo"), ("domain", "shared.com")],
            ),
        ];

        let clusters = vec![vec![id_a], vec![id_b]];
        let bridge_fields = vec!["company".to_string(), "domain".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            ..Default::default()
        };

        let mut canonical_map = HashMap::new();
        canonical_map.insert(id_a, uuid_from_u128(100));
        canonical_map.insert(id_b, uuid_from_u128(200));

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert_eq!(edges.len(), 1);

        // Each field has frequency 2, contribution = 1.0 / (1.0 + log2(2)) = 0.5
        // Total weight should be ~1.0 (sum of two contributions)
        let single_contribution = 1.0 / (1.0 + 2.0_f64.log2());
        let expected_total = single_contribution * 2.0;
        assert!(
            (edges[0].weight - expected_total).abs() < 0.001,
            "Multiple shared attributes should sum their contributions, expected {}, got {}",
            expected_total,
            edges[0].weight
        );
        assert_eq!(edges[0].attributes.len(), 2);
    }

    #[test]
    fn test_min_edge_weight_filter() {
        let id_a = uuid_from_u128(1);
        let id_b = uuid_from_u128(2);

        let entities = vec![
            make_entity(id_a, vec![("company", "WeakCo")]),
            make_entity(id_b, vec![("company", "WeakCo")]),
        ];

        let clusters = vec![vec![id_a], vec![id_b]];
        let bridge_fields = vec!["company".to_string()];
        // Set high min_edge_weight so the edge is filtered out
        let config = RelationshipConfig {
            min_edge_weight: 5.0,
            ..Default::default()
        };

        let mut canonical_map = HashMap::new();
        canonical_map.insert(id_a, uuid_from_u128(100));
        canonical_map.insert(id_b, uuid_from_u128(200));

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert!(
            edges.is_empty(),
            "Edges below min_edge_weight should be filtered out"
        );
    }

    #[test]
    fn test_empty_field_values_ignored() {
        let id_a = uuid_from_u128(1);
        let id_b = uuid_from_u128(2);

        let entities = vec![
            make_entity(id_a, vec![("company", "")]),
            make_entity(id_b, vec![("company", "")]),
        ];

        let clusters = vec![vec![id_a], vec![id_b]];
        let bridge_fields = vec!["company".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            ..Default::default()
        };

        let canonical_map = HashMap::new();

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert!(edges.is_empty(), "Empty field values should be ignored");
    }

    #[test]
    fn test_non_bridge_fields_ignored() {
        let id_a = uuid_from_u128(1);
        let id_b = uuid_from_u128(2);

        let entities = vec![
            make_entity(id_a, vec![("email", "test@example.com"), ("status", "active")]),
            make_entity(id_b, vec![("email", "test@example.com"), ("status", "active")]),
        ];

        let clusters = vec![vec![id_a], vec![id_b]];
        // Only consider "email", not "status"
        let bridge_fields = vec!["email".to_string()];
        let config = RelationshipConfig {
            min_edge_weight: 0.0,
            ..Default::default()
        };

        let mut canonical_map = HashMap::new();
        canonical_map.insert(id_a, uuid_from_u128(100));
        canonical_map.insert(id_b, uuid_from_u128(200));

        let edges =
            compute_relationship_edges(&entities, &clusters, &bridge_fields, &config, &canonical_map);
        assert_eq!(edges.len(), 1);
        // Only one attribute (email), not two
        assert_eq!(edges[0].attributes.len(), 1);
        assert_eq!(edges[0].attributes[0].field, "email");
    }

    #[test]
    fn test_build_frequency_map() {
        let entities = vec![
            make_entity(uuid_from_u128(1), vec![("company", "Acme"), ("domain", "acme.com")]),
            make_entity(uuid_from_u128(2), vec![("company", "Acme"), ("domain", "beta.com")]),
            make_entity(uuid_from_u128(3), vec![("company", "Beta")]),
        ];

        let bridge_fields = vec!["company".to_string(), "domain".to_string()];
        let freq = build_frequency_map(&entities, &bridge_fields);

        assert_eq!(
            freq.get(&("company".to_string(), "Acme".to_string())),
            Some(&2)
        );
        assert_eq!(
            freq.get(&("company".to_string(), "Beta".to_string())),
            Some(&1)
        );
        assert_eq!(
            freq.get(&("domain".to_string(), "acme.com".to_string())),
            Some(&1)
        );
        assert_eq!(
            freq.get(&("domain".to_string(), "beta.com".to_string())),
            Some(&1)
        );
    }
}
