use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use crate::types::NormalizedEntity;
use crate::ir::{SurvivorshipPolicy, SurvivorshipStrategy};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum AggregateFunction {
    Max,
    Min,
    Sum,
    Avg,
    Union,
    Intersection,
}

pub struct SurvivorshipEngine;

impl SurvivorshipEngine {
    pub fn construct_golden_record(
        cluster: &[NormalizedEntity],
        policy: &SurvivorshipPolicy,
    ) -> HashMap<String, String> {
        if cluster.is_empty() {
            return HashMap::new();
        }

        let mut golden = HashMap::new();
        let mut fields = HashSet::new();
        for entity in cluster {
            for key in entity.data.keys() {
                fields.insert(key.clone());
            }
        }

        for field in fields {
            let field_policy = policy.field_policies.get(&field);
            let strategy = field_policy
                .map(|p| p.strategy.clone())
                .unwrap_or_else(|| policy.default_strategy.clone());
            
            if let Some(value) = match strategy {
                SurvivorshipStrategy::MostRecent => Self::pick_by_most_recent(cluster, &field),
                SurvivorshipStrategy::MostComplete => {
                    cluster.iter()
                        .filter_map(|e| e.data.get(&field))
                        .max_by_key(|v| v.len())
                        .cloned()
                },
                SurvivorshipStrategy::SourcePriority => {
                     let priority = field_policy
                         .and_then(|p| p.source_priority.as_ref())
                         .map(|p| p.as_slice())
                         .unwrap_or(&[]);
                     Self::pick_by_source_priority(cluster, &field, priority)
                },
                SurvivorshipStrategy::Aggregate => {
                    let func_name = field_policy.and_then(|p| p.aggregate_fn.as_ref());
                    Self::aggregate_values(cluster, &field, func_name)
                },
                _ => Self::pick_by_most_recent(cluster, &field),
            } {
                golden.insert(field, value);
            }
        }
        golden
    }

    fn pick_by_most_recent(cluster: &[NormalizedEntity], field: &str) -> Option<String> {
        cluster.iter()
            .filter(|e| e.data.contains_key(field))
            .max_by(|a, b| {
                let a_time = a.valid_from.unwrap_or(a.last_updated);
                let b_time = b.valid_from.unwrap_or(b.last_updated);
                a_time.cmp(&b_time).then_with(|| a.id.cmp(&b.id))
            })
            .and_then(|e| e.data.get(field).cloned())
    }

    fn pick_by_source_priority(cluster: &[NormalizedEntity], field: &str, priority: &[String]) -> Option<String> {
        let candidates: Vec<&NormalizedEntity> = cluster.iter().filter(|e| e.data.contains_key(field)).collect();
        
        if candidates.is_empty() {
            return None;
        }

        let mut best_candidate = candidates[0];
        let mut best_rank = usize::MAX;

        for candidate in candidates {
             if let Some(rank) = priority.iter().position(|s| s == &candidate.source_name) {
                 if rank < best_rank {
                     best_rank = rank;
                     best_candidate = candidate;
                 }
             }
        }
        
        if best_rank == usize::MAX {
             return Self::pick_by_most_recent(cluster, field);
        }

        best_candidate.data.get(field).cloned()
    }
    
    fn aggregate_values(cluster: &[NormalizedEntity], field: &str, func_name: Option<&String>) -> Option<String> {
        let values: Vec<&String> = cluster.iter().filter_map(|e| e.data.get(field)).collect();
        if values.is_empty() { return None; }
        
        let func = match func_name.map(|s| s.as_str()) {
            Some("Max") => AggregateFunction::Max,
            Some("Min") => AggregateFunction::Min,
            Some("Sum") => AggregateFunction::Sum,
            Some("Avg") => AggregateFunction::Avg,
            Some("Union") => AggregateFunction::Union,
            Some("Intersection") => AggregateFunction::Intersection,
            _ => AggregateFunction::Union,
        };

        match func {
             AggregateFunction::Max => values.iter().max().cloned().cloned(),
             AggregateFunction::Min => values.iter().min().cloned().cloned(),
             AggregateFunction::Sum => {
                 let sum: f64 = values.iter().filter_map(|v| v.parse::<f64>().ok()).sum();
                 Some(sum.to_string())
             },
             AggregateFunction::Avg => {
                 let valid: Vec<f64> = values.iter().filter_map(|v| v.parse::<f64>().ok()).collect();
                 if valid.is_empty() { return None; }
                 let sum: f64 = valid.iter().sum();
                 Some((sum / valid.len() as f64).to_string())
             },
             AggregateFunction::Union => {
                 let mut unique: Vec<&String> = values.iter().copied().collect::<HashSet<_>>().into_iter().collect();
                 unique.sort();
                 Some(unique.into_iter().cloned().collect::<Vec<_>>().join(","))
             },
             AggregateFunction::Intersection => {
                  let first = values[0];
                  if values.iter().all(|&v| v == first) {
                      Some(first.clone())
                  } else {
                      None
                  }
             }
        }
    }
}
