pub mod types;
pub mod ir;
pub mod normalization;
pub mod matching;
pub mod blocking;
pub mod survivorship;
pub mod temporal;
pub mod overrides;
pub mod compiler;
pub mod spec;
pub mod observability;
#[cfg(feature = "python")]
pub mod python;

use crate::types::{NormalizedEntity, MatchDecision, Decision};
use crate::matching::MatchRule;
use crate::blocking::BlockingStrategy;
use crate::overrides::OverrideResolver;
use crate::ir::{IdentityPlan, TemporalStrategy};
use crate::temporal::{TemporalMatcher, TemporalResult};
use crate::observability::{ReconciliationTelemetry, TelemetryCollector};
use hashbrown::{HashMap, HashSet};
use uuid::Uuid;
use rayon::prelude::*;
use std::sync::Mutex;

pub struct ReconciliationEngine {
    pub rules: Vec<Box<dyn MatchRule>>,
    pub blocking_strategies: Vec<Box<dyn BlockingStrategy>>,
    pub merge_threshold: f64,
    pub review_threshold: f64,
    pub max_block_size: usize,
    pub neighborhood_window: Option<usize>,
    pub sort_fields: Vec<String>,
    pub temporal_strategy: TemporalStrategy,
    pub reference_identifiers: Vec<crate::ir::CompiledReferenceIdentifier>,
    pub relations: Vec<crate::ir::CompiledRelation>,
}

impl ReconciliationEngine {
    pub fn new(
        rules: Vec<Box<dyn MatchRule>>,
        blocking_strategies: Vec<Box<dyn BlockingStrategy>>,
        merge_threshold: f64,
        review_threshold: f64,
    ) -> Self {
        Self {
            rules,
            blocking_strategies,
            merge_threshold,
            review_threshold,
            max_block_size: 1000, 
            neighborhood_window: Some(50), 
            sort_fields: vec![],
            temporal_strategy: TemporalStrategy::default(), 
            reference_identifiers: vec![],
            relations: vec![],
        }
    }

    pub fn from_plan(plan: &IdentityPlan) -> Self {
        let rules = plan.match_graph.rules.iter()
            .map(|r| Self::compile_rule(r))
            .collect();
            
        let mut blocking_strategies: Vec<Box<dyn BlockingStrategy>> = vec![
            Box::new(crate::blocking::PlanBlocking { keys: plan.blocking.keys.clone() })
        ];

        if plan.blocking.strategy == crate::ir::BlockingStrategy::Lsh {
            let fields: Vec<String> = plan.blocking.keys.iter()
                .flat_map(|k| k.fields.clone())
                .collect();
            let bands = plan.blocking.lsh_bands.unwrap_or(20);
            let rows = plan.blocking.lsh_rows.unwrap_or(5);
            blocking_strategies.push(Box::new(crate::blocking::lsh::LshBlocking::new(fields, bands, rows)));
        }

        let mut engine = Self {
            rules,
            blocking_strategies,
            merge_threshold: plan.decision.match_threshold,
            review_threshold: plan.decision.review_threshold.unwrap_or(plan.decision.match_threshold),
            max_block_size: 1000, 
            neighborhood_window: plan.blocking.neighborhood_window,
            sort_fields: plan.blocking.sort_fields.clone(),
            temporal_strategy: TemporalStrategy::default(),
            reference_identifiers: plan.reference_identifiers.clone(),
            relations: plan.relations.clone(),
        };
        
        if let Some(first_source) = plan.sources.iter().find(|s| s.temporal.is_some()) {
            if let Some(ref t) = first_source.temporal {
                engine.temporal_strategy = t.strategy.clone();
            }
        }
        
        engine
    }

    fn compile_rule(rule: &crate::ir::CompiledRule) -> Box<dyn MatchRule> {
        match &rule.rule_type {
            crate::ir::CompiledRuleType::Exact { field, weight, case_insensitive, normalize } => {
                Box::new(crate::matching::ExactMatch {
                    field: field.clone(),
                    weight: *weight,
                    case_insensitive: *case_insensitive,
                    normalize: *normalize,
                    rule_name: rule.name.clone(),
                })
            },
            crate::ir::CompiledRuleType::Similarity { field, algorithm, threshold, weight } => {
                use crate::ir::SimilarityAlgorithm;
                match algorithm {
                     SimilarityAlgorithm::JaroWinkler => Box::new(crate::matching::FuzzyStringMatch {
                          field: field.clone(),
                          threshold: *threshold,
                          weight: *weight,
                          rule_name: rule.name.clone(),
                      }),
                      SimilarityAlgorithm::Levenshtein => Box::new(crate::matching::LevenshteinMatch {
                          field: field.clone(),
                          threshold: *threshold,
                          weight: *weight,
                          rule_name: rule.name.clone(),
                      }),
                      _ => {
                          Box::new(crate::matching::ExactMatch {
                              field: field.clone(),
                              weight: *weight,
                              case_insensitive: true,
                              normalize: true,
                              rule_name: rule.name.clone(),
                          })
                      }
                }
            },
            crate::ir::CompiledRuleType::Range { field, tolerance, weight } => {
                Box::new(crate::matching::RangeMatch {
                    field: field.clone(),
                    tolerance: *tolerance,
                    weight: *weight,
                    rule_name: rule.name.clone(),
                })
            },
            crate::ir::CompiledRuleType::Composite { operator, children } => {
                let child_rules = children.iter().map(|c| Self::compile_rule(c)).collect();
                use crate::ir::CompositeOperator;
                Box::new(crate::matching::CompositeMatch {
                    rules: child_rules,
                    weight: 1.0, 
                    require_all: match operator {
                        CompositeOperator::And => true,
                        CompositeOperator::Or => false,
                    },
                    rule_name: rule.name.clone(),
                })
            },
        }
    }

    pub fn reconcile(&self, entities: &[NormalizedEntity], override_resolver: &OverrideResolver) -> Vec<MatchDecision> {
        let entity_map: HashMap<Uuid, &NormalizedEntity> = entities.iter()
            .map(|e| (e.id, e))
            .collect();

        let mut groups: HashMap<String, Vec<Uuid>> = HashMap::new();
        for entity in entities {
            for strategy in &self.blocking_strategies {
                let keys = strategy.generate_block_key(entity);
                for key in keys {
                    groups.entry(key).or_default().push(entity.id);
                }
            }
        }

        let raw_decisions: Vec<MatchDecision> = groups.par_iter()
            .filter(|(_, ids)| ids.len() >= 2 && ids.len() <= self.max_block_size)
            .flat_map(|(_key, ids)| {
                let mut block_decisions = Vec::new();
                for i in 0..ids.len() {
                    for j in i + 1..ids.len() {
                        let id_a = ids[i];
                        let id_b = ids[j];
                        
                        let entity_a = entity_map.get(&id_a).expect("Entity not found in map");
                        let entity_b = entity_map.get(&id_b).expect("Entity not found in map");

                        if let Some(mut decision) = self.reconcile_pair(entity_a, entity_b, override_resolver) {
                            decision.blocking_key = Some(_key.clone());
                            block_decisions.push(decision);
                        }
                    }
                }
                block_decisions
            })
            .collect();

        let mut final_decisions = Vec::new();
        let mut seen_pairs = HashSet::new();
        for d in raw_decisions {
            let pair = if d.entity_a_id < d.entity_b_id { (d.entity_a_id, d.entity_b_id) } else { (d.entity_b_id, d.entity_a_id) };
            if seen_pairs.insert(pair) {
                final_decisions.push(d);
            }
        }

        if let Some(window) = self.neighborhood_window {
            if !self.sort_fields.is_empty() {
                let mut sorted_entities: Vec<_> = entities.iter().collect();
                sorted_entities.sort_by_key(|e| {
                    let mut key = String::new();
                    for field in &self.sort_fields {
                        let val = e.data.get(field).map(|v| v.to_lowercase()).unwrap_or_default();
                        key.push_str(&val);
                        key.push(':');
                    }
                    key
                });

                for i in 0..sorted_entities.len() {
                    for j in i + 1..std::cmp::min(i + window, sorted_entities.len()) {
                        let entity_a = sorted_entities[i];
                        let entity_b = sorted_entities[j];

                        let pair = if entity_a.id < entity_b.id { (entity_a.id, entity_b.id) } else { (entity_b.id, entity_a.id) };
                        if seen_pairs.insert(pair) {
                            if let Some(mut decision) = self.reconcile_pair(entity_a, entity_b, override_resolver) {
                                decision.blocking_key = Some("SORTED_NEIGHBORHOOD".to_string());
                                final_decisions.push(decision);
                            }
                        }
                    }
                }
            }
        }
        
        final_decisions
    }

    pub fn reconcile_with_telemetry(&self, entities: &[NormalizedEntity], override_resolver: &OverrideResolver) -> (Vec<MatchDecision>, ReconciliationTelemetry) {
        let entity_map: HashMap<Uuid, &NormalizedEntity> = entities.iter()
            .map(|e| (e.id, e))
            .collect();

        let mut groups: HashMap<String, Vec<Uuid>> = HashMap::new();
        for entity in entities {
            for strategy in &self.blocking_strategies {
                let keys = strategy.generate_block_key(entity);
                for key in keys {
                    groups.entry(key).or_default().push(entity.id);
                }
            }
        }

        let collector = Mutex::new(TelemetryCollector::new());
        collector.lock().unwrap().blocking_groups = groups.len();

        let raw_decisions: Vec<MatchDecision> = groups.par_iter()
            .filter(|(_, ids)| ids.len() >= 2 && ids.len() <= self.max_block_size)
            .flat_map(|(_key, ids)| {
                let mut block_decisions = Vec::new();
                for i in 0..ids.len() {
                    for j in i + 1..ids.len() {
                        let id_a = ids[i];
                        let id_b = ids[j];

                        let entity_a = entity_map.get(&id_a).expect("Entity not found in map");
                        let entity_b = entity_map.get(&id_b).expect("Entity not found in map");

                        let (maybe_decision, maybe_telemetry) = self.reconcile_pair_with_telemetry(entity_a, entity_b, override_resolver);

                        if let Some((rule_results, decision_type)) = maybe_telemetry {
                            if let Ok(mut c) = collector.lock() {
                                c.record_pair(&rule_results, &decision_type);
                            }
                        }

                        if let Some(mut decision) = maybe_decision {
                            decision.blocking_key = Some(_key.clone());
                            block_decisions.push(decision);
                        }
                    }
                }
                block_decisions
            })
            .collect();

        let mut final_decisions = Vec::new();
        let mut seen_pairs = HashSet::new();
        for d in raw_decisions {
            let pair = if d.entity_a_id < d.entity_b_id { (d.entity_a_id, d.entity_b_id) } else { (d.entity_b_id, d.entity_a_id) };
            if seen_pairs.insert(pair) {
                final_decisions.push(d);
            }
        }

        if let Some(window) = self.neighborhood_window {
            if !self.sort_fields.is_empty() {
                let mut sorted_entities: Vec<_> = entities.iter().collect();
                sorted_entities.sort_by_key(|e| {
                    let mut key = String::new();
                    for field in &self.sort_fields {
                        let val = e.data.get(field).map(|v| v.to_lowercase()).unwrap_or_default();
                        key.push_str(&val);
                        key.push(':');
                    }
                    key
                });

                for i in 0..sorted_entities.len() {
                    for j in i + 1..std::cmp::min(i + window, sorted_entities.len()) {
                        let entity_a = sorted_entities[i];
                        let entity_b = sorted_entities[j];

                        let pair = if entity_a.id < entity_b.id { (entity_a.id, entity_b.id) } else { (entity_b.id, entity_a.id) };
                        if seen_pairs.insert(pair) {
                            let (maybe_decision, maybe_telemetry) = self.reconcile_pair_with_telemetry(entity_a, entity_b, override_resolver);

                            if let Some((rule_results, decision_type)) = maybe_telemetry {
                                if let Ok(mut c) = collector.lock() {
                                    c.record_pair(&rule_results, &decision_type);
                                }
                            }

                            if let Some(mut decision) = maybe_decision {
                                decision.blocking_key = Some("SORTED_NEIGHBORHOOD".to_string());
                                final_decisions.push(decision);
                            }
                        }
                    }
                }
            }
        }

        let telemetry = collector.into_inner().unwrap().into_telemetry();
        (final_decisions, telemetry)
    }

    /// Like reconcile_pair, but returns telemetry data for all evaluated pairs.
    /// Returns (Option<MatchDecision>, telemetry_data):
    /// - MatchDecision is Some only for non-NoMerge (same filtering as reconcile_pair)
    /// - Telemetry data is Some for all pairs that pass temporal filter (including NoMerge)
    /// - Both are None for temporally-filtered pairs
    fn reconcile_pair_with_telemetry(&self, a: &NormalizedEntity, b: &NormalizedEntity, resolver: &OverrideResolver) -> (Option<MatchDecision>, Option<(Vec<crate::types::RuleResult>, String)>) {
        let (temporal_res, temporal_adj) = TemporalMatcher::evaluate(
            &self.temporal_strategy,
            a,
            b,
            None
        );

        if temporal_res == TemporalResult::Historical && temporal_adj == 0.0 {
            return (None, None);
        }

        let decision = self.evaluate_pair(a, b, resolver, temporal_adj);
        let decision_type = format!("{:?}", decision.decision).to_lowercase();
        let rule_results = decision.rule_results.clone();
        let telemetry = Some((rule_results, decision_type));

        if decision.decision != Decision::NoMerge {
            (Some(decision), telemetry)
        } else {
            (None, telemetry)
        }
    }

    fn reconcile_pair(&self, a: &NormalizedEntity, b: &NormalizedEntity, resolver: &OverrideResolver) -> Option<MatchDecision> {
        let (temporal_res, temporal_adj) = TemporalMatcher::evaluate(
            &self.temporal_strategy, 
            a, 
            b, 
            None
        );

        if temporal_res == TemporalResult::Historical && temporal_adj == 0.0 {
            return None;
        }

        let decision = self.evaluate_pair(a, b, resolver, temporal_adj);
        if decision.decision != Decision::NoMerge {
            Some(decision)
        } else {
            None
        }
    }

    pub fn cluster_decisions(&self, entities: &[NormalizedEntity], decisions: &[MatchDecision]) -> Vec<Vec<Uuid>> {
        use petgraph::unionfind::UnionFind;
        
        let id_to_idx: HashMap<Uuid, usize> = entities.iter().enumerate().map(|(i, e)| (e.id, i)).collect();
        let idx_to_id: Vec<Uuid> = entities.iter().map(|e| e.id).collect();
        
        let mut uf = UnionFind::new(entities.len());
        
        for decision in decisions {
            if decision.decision == Decision::Merge {
                if let (Some(&idx_a), Some(&idx_b)) = (id_to_idx.get(&decision.entity_a_id), id_to_idx.get(&decision.entity_b_id)) {
                    uf.union(idx_a, idx_b);
                }
            }
        }
        
        let mut groups: HashMap<usize, Vec<Uuid>> = HashMap::new();
        for i in 0..entities.len() {
            let root = uf.find(i);
            groups.entry(root).or_default().push(idx_to_id[i]);
        }
        
        groups.into_values().collect()
    }

    pub fn evaluate_pair(&self, a: &NormalizedEntity, b: &NormalizedEntity, resolver: &OverrideResolver, scale: f64) -> MatchDecision {
        if resolver.is_split(a.id, b.id) {
            return MatchDecision {
                entity_a_id: a.id,
                entity_b_id: b.id,
                decision: Decision::NoMerge,
                confidence: 0.0,
                rule_results: vec![],
                matched_on: vec!["MANUAL_SPLIT".to_string()],
                blocking_key: None,
                temporal_adjustment: Some(scale),
            };
        }

        if resolver.is_locked(a.id, b.id) {
            return MatchDecision {
                entity_a_id: a.id,
                entity_b_id: b.id,
                decision: Decision::NoMerge,
                confidence: 0.0,
                rule_results: vec![],
                matched_on: vec!["MANUAL_LOCK".to_string()],
                blocking_key: None,
                temporal_adjustment: Some(scale),
            };
        }

        if resolver.is_forced_merge(a.id, b.id) {
            return MatchDecision {
                entity_a_id: a.id,
                entity_b_id: b.id,
                decision: Decision::Merge,
                confidence: 1.0,
                rule_results: vec![],
                matched_on: vec!["MANUAL_MERGE".to_string()],
                blocking_key: None,
                temporal_adjustment: Some(scale),
            };
        }

        for rid in &self.reference_identifiers {
            let val_a = a.data.get(&rid.field);
            let val_b = b.data.get(&rid.field);
            
            if let (Some(v_a), Some(v_b)) = (val_a, val_b) {
                if v_a == v_b && !v_a.is_empty() {
                    return MatchDecision {
                        entity_a_id: a.id,
                        entity_b_id: b.id,
                        decision: Decision::Merge,
                        confidence: 1.0,
                        rule_results: vec![],
                        matched_on: vec![format!("REF_ID:{}", rid.name)],
                        blocking_key: None,
                        temporal_adjustment: Some(scale),
                    };
                }
            }
        }

        let mut rule_results = Vec::new();
        let mut total_score = 0.0;
        let mut total_weight = 0.0;
        let mut matched_on = Vec::new();

        for rule in &self.rules {
            let result = rule.evaluate(a, b);
            total_score += result.score * result.weight;
            total_weight += result.weight;
            if result.matched {
                matched_on.push(result.rule_name.clone());
            }
            rule_results.push(result);
        }

        let confidence = if total_weight > 0.0 {
            (total_score / total_weight) * scale
        } else {
            0.0
        };

        let decision = if confidence >= self.merge_threshold {
            Decision::Merge
        } else if confidence >= self.review_threshold {
            Decision::Review
        } else {
            Decision::NoMerge
        };

        MatchDecision {
            entity_a_id: a.id,
            entity_b_id: b.id,
            decision,
            confidence,
            rule_results,
            matched_on,
            blocking_key: None, 
            temporal_adjustment: Some(scale),
        }
    }

    pub fn generate_block_keys(&self, entity: &NormalizedEntity) -> Vec<String> {
        let mut keys = Vec::new();
        for strategy in &self.blocking_strategies {
            keys.extend(strategy.generate_block_key(entity));
        }
        keys
    }

    pub fn propagate_relations(
        &self,
        clusters: &[Vec<Uuid>],
        entities: &[NormalizedEntity],
        resolver: &OverrideResolver,
        initial_decisions: &[MatchDecision],
    ) -> Vec<MatchDecision> {
        let mut all_propagated_decisions = Vec::new();
        let mut current_clusters = clusters.to_vec();
        
        let mut entity_map: HashMap<Uuid, &NormalizedEntity> = HashMap::new();
        let mut entities_by_type_val: HashMap<(String, String), Vec<Uuid>> = HashMap::new();
        
        for e in entities {
            entity_map.insert(e.id, e);
            // We index everything by type and external_id (or whatever the join_key points to)
            // But wait, the join_key points to a VALUE. We should index by (type, val)
            // In the target entity, that value is usually its ID or some attribute.
            // Let's index by ALL possible attributes for the join_key lookup.
            for (_attr, val) in &e.data {
                entities_by_type_val.entry((e.entity_type.clone(), val.clone())).or_default().push(e.id);
            }
            // Also index by its own id (as string)
            entities_by_type_val.entry((e.entity_type.clone(), e.id.to_string())).or_default().push(e.id);
            // And external_id
            entities_by_type_val.entry((e.entity_type.clone(), e.external_id.clone())).or_default().push(e.id);
        }

        let max_hops = 5;
        let mut seen_pairs = HashSet::new();

        for hop in 0..max_hops {
            let mut new_decisions_this_hop = Vec::new();

            for relation in &self.relations {
                if !relation.propagate_match {
                    continue;
                }

                for cluster in &current_clusters {
                    if cluster.len() < 2 {
                        continue;
                    }

                    let mut target_ids = HashSet::new();
                    for &id in cluster {
                        if let Some(entity) = entity_map.get(&id) {
                            if entity.entity_type == relation.from_entity {
                                if let Some(target_val) = entity.data.get(&relation.join_key) {
                                    // Find all entities of to_entity type that have this target_val
                                    if let Some(found_ids) = entities_by_type_val.get(&(relation.to_entity.clone(), target_val.clone())) {
                                        for &fid in found_ids {
                                            target_ids.insert(fid);
                                        }
                                    }
                                }
                            }
                        }
                    }

                    if target_ids.len() >= 2 {
                        let mut sorted_targets: Vec<_> = target_ids.into_iter().collect();
                        sorted_targets.sort();
                        for i in 0..sorted_targets.len() {
                            for j in i + 1..sorted_targets.len() {
                                let id_a = sorted_targets[i];
                                let id_b = sorted_targets[j];

                                if seen_pairs.contains(&(id_a, id_b)) {
                                    continue;
                                }

                                if resolver.is_split(id_a, id_b) || resolver.is_locked(id_a, id_b) {
                                    continue;
                                }

                                seen_pairs.insert((id_a, id_b));
                                new_decisions_this_hop.push(MatchDecision {
                                    entity_a_id: id_a,
                                    entity_b_id: id_b,
                                    decision: Decision::Merge,
                                    confidence: 1.0,
                                    rule_results: vec![],
                                    matched_on: vec![format!("RELATION:{}:hop:{}", relation.name, hop)],
                                    blocking_key: None,
                                    temporal_adjustment: None,
                                });
                            }
                        }
                    }
                }
            }

            if new_decisions_this_hop.is_empty() {
                break;
            }

            all_propagated_decisions.extend(new_decisions_this_hop);
            
            let mut all_current_decisions = initial_decisions.to_vec();
            all_current_decisions.extend(all_propagated_decisions.clone());
            
            current_clusters = self.cluster_decisions(entities, &all_current_decisions);
        }

        all_propagated_decisions
    }
}
