//! Diversity selection algorithm with greedy multi-stage relaxation.
//!
//! After the profile executor scores and sorts candidates, the `DiversitySelector`
//! applies post-hoc constraints (max-per-creator, format-mix) using a greedy
//! multi-stage relaxation strategy. This ensures the result set maintains content
//! variety without sacrificing result count.
//!
//! The selector guarantees two invariants:
//! - **INV-RANK-5:** `result.len() == min(target_count, candidates.len())` always.
//! - **INV-RANK-6:** Global score order is preserved across the entire result, not
//!   just within creator groups. The final emit pass walks `candidates` (which is
//!   already sorted descending by score) and retains only accepted IDs.
//!
//! Relaxation stages:
//! - Stage 0: full constraints
//! - Stage 1: `max_per_creator` doubled
//! - Stage 2: `format_mix` ignored
//! - Stage 3: accept anything (fill to target)

use std::collections::{HashMap, HashSet};

use super::constraints::{ConstraintViolation, DiversityConstraints, DiversityResult};
use crate::ranking::executor::ScoredCandidate;

// -- Selector ---------------------------------------------------------------

/// Stateless diversity selector. Applies constraints via greedy multi-stage
/// relaxation.
pub struct DiversitySelector;

impl DiversitySelector {
    /// Select up to `target_count` candidates satisfying diversity constraints.
    ///
    /// Uses a three-stage relaxation strategy:
    /// - Stage 0: full constraints
    /// - Stage 1: `max_per_creator` doubled (if used)
    /// - Stage 2: `format_mix` ignored
    /// - Stage 3: accept anything (fill to target)
    ///
    /// INV-RANK-5: `result.selected.len() == min(target_count, candidates.len())` always.
    /// INV-RANK-6: score order preserved within same-creator groups.
    #[must_use]
    pub fn select(
        candidates: &[ScoredCandidate],
        constraints: &DiversityConstraints,
        target_count: usize,
    ) -> DiversityResult {
        debug_assert!(
            candidates
                .windows(2)
                .all(|w| w[0].score >= w[1].score - f64::EPSILON),
            "DiversitySelector::select requires candidates in descending score order"
        );

        let max_needed = target_count.min(candidates.len());

        // Fast path 1: no active constraints -- return top-N directly.
        if constraints.max_per_creator.is_none() && constraints.format_mix_max_fraction.is_none() {
            return DiversityResult {
                selected: candidates[..max_needed].to_vec(),
                constraints_satisfied: true,
                violations: vec![],
            };
        }

        // Fast path 2: fewer candidates than target -- return all, report violations.
        if candidates.len() <= target_count {
            let violations = collect_violations(candidates, constraints);
            return DiversityResult {
                selected: candidates.to_vec(),
                constraints_satisfied: violations.is_empty(),
                violations,
            };
        }

        // Collect accepted entity IDs across all relaxation stages, then emit
        // in a single pass over `candidates` (which is score-sorted descending).
        // This preserves global score order -- INV-RANK-6 holds for all items,
        // not just within same-creator groups.
        let mut accepted: HashSet<u64> = HashSet::new();

        // Stage 0: greedy pass with full constraints.
        let stage0 = DiversityConstraints {
            max_per_creator: constraints.max_per_creator,
            format_mix_max_fraction: constraints.format_mix_max_fraction,
            ..DiversityConstraints::new()
        };
        for c in greedy_select(candidates, &stage0, max_needed) {
            accepted.insert(c.entity_id.as_u64());
        }

        // Stage 1: double max_per_creator if we did not fill target.
        if accepted.len() < max_needed {
            let stage1 = DiversityConstraints {
                max_per_creator: constraints.max_per_creator.map(|n| n.saturating_mul(2)),
                format_mix_max_fraction: constraints.format_mix_max_fraction,
                ..DiversityConstraints::new()
            };
            let remaining: Vec<_> = candidates
                .iter()
                .filter(|c| !accepted.contains(&c.entity_id.as_u64()))
                .cloned()
                .collect();
            for c in greedy_select(&remaining, &stage1, max_needed - accepted.len()) {
                accepted.insert(c.entity_id.as_u64());
            }
        }

        // Stage 2: ignore format_mix, keep doubled max_per_creator.
        if accepted.len() < max_needed {
            let stage2 = DiversityConstraints {
                max_per_creator: constraints.max_per_creator.map(|n| n.saturating_mul(2)),
                format_mix_max_fraction: None,
                ..DiversityConstraints::new()
            };
            let remaining: Vec<_> = candidates
                .iter()
                .filter(|c| !accepted.contains(&c.entity_id.as_u64()))
                .cloned()
                .collect();
            for c in greedy_select(&remaining, &stage2, max_needed - accepted.len()) {
                accepted.insert(c.entity_id.as_u64());
            }
        }

        // Stage 3: accept anything to fill to target.
        if accepted.len() < max_needed {
            for c in candidates {
                if accepted.len() >= max_needed {
                    break;
                }
                accepted.insert(c.entity_id.as_u64());
            }
        }

        // Emit in original score order by walking the sorted candidate list once.
        let selected: Vec<ScoredCandidate> = candidates
            .iter()
            .filter(|c| accepted.contains(&c.entity_id.as_u64()))
            .cloned()
            .collect();

        let violations = collect_violations(&selected, constraints);
        let constraints_satisfied = violations.is_empty();
        DiversityResult {
            selected,
            constraints_satisfied,
            violations,
        }
    }
}

// -- Internal helpers -------------------------------------------------------

/// Greedy selection: iterate candidates in score order, accepting each candidate
/// only if it does not violate the given constraints.
///
/// Candidates are already in descending score order. By iterating in this
/// order and appending to `selected`, each creator's items appear in score
/// order within the output -- satisfying INV-RANK-6.
fn greedy_select(
    candidates: &[ScoredCandidate],
    constraints: &DiversityConstraints,
    limit: usize,
) -> Vec<ScoredCandidate> {
    let mut selected = Vec::with_capacity(limit);
    let mut creator_counts: HashMap<u64, usize> = HashMap::new();
    let mut format_counts: HashMap<String, usize> = HashMap::new();

    for candidate in candidates {
        if selected.len() >= limit {
            break;
        }

        // Check max_per_creator.
        if let Some(max) = constraints.max_per_creator
            && let Some(creator_id) = candidate.creator_id
        {
            let count = creator_counts
                .get(&creator_id.as_u64())
                .copied()
                .unwrap_or(0);
            if count >= max {
                continue;
            }
        }

        // Check format_mix_max_fraction.
        // Convert fraction to an absolute count limit against the target (limit),
        // ensuring at least 1 of each format is always allowed.
        if let Some(max_frac) = constraints.format_mix_max_fraction
            && let Some(ref fmt) = candidate.format
        {
            #[allow(
                clippy::cast_precision_loss,
                clippy::cast_possible_truncation,
                clippy::cast_sign_loss
            )]
            let max_allowed = (max_frac * limit as f64).floor().max(1.0) as usize;
            let fmt_count = format_counts.get(fmt).copied().unwrap_or(0) + 1;
            if fmt_count > max_allowed {
                continue;
            }
        }

        // Accept candidate.
        if let Some(creator_id) = candidate.creator_id {
            *creator_counts.entry(creator_id.as_u64()).or_insert(0) += 1;
        }
        if let Some(ref fmt) = candidate.format {
            *format_counts.entry(fmt.clone()).or_insert(0) += 1;
        }
        selected.push(candidate.clone());
    }
    selected
}

/// Collect constraint violations for the given selected candidates against
/// the original constraints.
fn collect_violations(
    selected: &[ScoredCandidate],
    constraints: &DiversityConstraints,
) -> Vec<ConstraintViolation> {
    let mut violations = Vec::new();

    if let Some(max) = constraints.max_per_creator {
        let mut counts: HashMap<u64, usize> = HashMap::new();
        for c in selected {
            if let Some(cid) = c.creator_id {
                *counts.entry(cid.as_u64()).or_insert(0) += 1;
            }
        }
        for (cid, count) in &counts {
            if *count > max {
                violations.push(ConstraintViolation {
                    constraint: "max_per_creator".into(),
                    detail: format!("creator {cid} has {count} items (max {max})"),
                });
            }
        }
    }

    if let Some(max_frac) = constraints.format_mix_max_fraction {
        let total = selected.len();
        if total > 0 {
            let mut fmt_counts: HashMap<&str, usize> = HashMap::new();
            for c in selected {
                if let Some(ref fmt) = c.format {
                    *fmt_counts.entry(fmt.as_str()).or_insert(0) += 1;
                }
            }
            for (fmt, count) in &fmt_counts {
                #[allow(clippy::cast_precision_loss)]
                let frac = *count as f64 / total as f64;
                if frac > max_frac {
                    violations.push(ConstraintViolation {
                        constraint: "format_mix".into(),
                        detail: format!("format '{fmt}' is {frac:.2} (max {max_frac:.2})"),
                    });
                }
            }
        }
    }

    violations
}

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

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::float_cmp,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss
)]
mod tests {
    use super::*;
    use crate::schema::EntityId;

    fn make_candidate(
        id: u64,
        score: f64,
        creator: Option<u64>,
        format: Option<&str>,
    ) -> ScoredCandidate {
        ScoredCandidate {
            entity_id: EntityId::new(id),
            score,
            signal_snapshot: vec![],
            creator_id: creator.map(EntityId::new),
            format: format.map(String::from),
        }
    }

    #[test]
    fn select_no_constraints_returns_top_n() {
        let candidates: Vec<_> = (0..10)
            .map(|i| make_candidate(i + 1, (10 - i) as f64, Some(i + 1), Some("video")))
            .collect();
        let constraints = DiversityConstraints::new();
        let result = DiversitySelector::select(&candidates, &constraints, 5);

        assert_eq!(result.selected.len(), 5);
        assert!(result.constraints_satisfied);
        assert!(result.violations.is_empty());
        // Should be the top 5 by score (10, 9, 8, 7, 6).
        for (i, c) in result.selected.iter().enumerate() {
            assert_eq!(c.score, (10 - i) as f64);
        }
    }

    #[test]
    fn select_max_per_creator_respected() {
        // 10 candidates all from the same creator.
        let candidates: Vec<_> = (0..10)
            .map(|i| make_candidate(i + 1, (10 - i) as f64, Some(1), Some("video")))
            .collect();
        let constraints = DiversityConstraints::new().max_per_creator(2);
        let result = DiversitySelector::select(&candidates, &constraints, 5);

        // Stage 0 picks 2, stage 1 picks 2 more (doubled to 4), stage 2 picks 0
        // more (no format constraint), stage 3 fills the rest. Total = 5.
        assert_eq!(result.selected.len(), 5);
        // Constraints are NOT satisfied because the final set has 5 from 1 creator
        // but max_per_creator is 2.
        assert!(!result.constraints_satisfied);
    }

    #[test]
    fn select_format_mix_respected() {
        // 4 "video" + 4 "audio", all different creators, interleaved by score.
        let candidates = vec![
            make_candidate(1, 10.0, Some(1), Some("video")),
            make_candidate(2, 9.0, Some(2), Some("audio")),
            make_candidate(3, 8.0, Some(3), Some("video")),
            make_candidate(4, 7.0, Some(4), Some("audio")),
            make_candidate(5, 6.0, Some(5), Some("video")),
            make_candidate(6, 5.0, Some(6), Some("audio")),
            make_candidate(7, 4.0, Some(7), Some("video")),
            make_candidate(8, 3.0, Some(8), Some("audio")),
        ];

        let constraints = DiversityConstraints::new().format_mix(0.5);
        let result = DiversitySelector::select(&candidates, &constraints, 4);

        // With max_fraction=0.5, at most 2 of 4 can be video.
        let video_count = result
            .selected
            .iter()
            .filter(|c| c.format.as_deref() == Some("video"))
            .count();
        assert!(video_count <= 2, "video_count={video_count}, expected <= 2");
        assert_eq!(result.selected.len(), 4);
        assert!(result.constraints_satisfied);
    }

    #[test]
    fn select_invariant_count() {
        // INV-RANK-5: always min(target, candidates.len()) results.
        // 10 candidates with 10 distinct creators, max_per_creator=1 -> can fill 5.
        let candidates: Vec<_> = (0..10)
            .map(|i| make_candidate(i + 1, (10 - i) as f64, Some(i + 1), Some("video")))
            .collect();
        let constraints = DiversityConstraints::new().max_per_creator(1);
        let result = DiversitySelector::select(&candidates, &constraints, 5);

        assert_eq!(result.selected.len(), 5);
    }

    #[test]
    fn select_unsatisfiable_constraints_stage3_fill() {
        // 5 candidates all same creator, max_per_creator=1, target=5.
        // Stage 0 picks 1, stage 1 picks 1 more (doubled to 2), stage 2 picks 0,
        // stage 3 fills remaining 3. Total = 5.
        let candidates: Vec<_> = (0..5)
            .map(|i| make_candidate(i + 1, (5 - i) as f64, Some(1), Some("video")))
            .collect();
        let constraints = DiversityConstraints::new().max_per_creator(1);
        let result = DiversitySelector::select(&candidates, &constraints, 5);

        // INV-RANK-5: must have exactly 5.
        assert_eq!(result.selected.len(), 5);
        // Constraints cannot be satisfied: 5 items from 1 creator with max=1.
        assert!(!result.constraints_satisfied);
    }

    #[test]
    fn select_fewer_candidates_than_target() {
        // INV-RANK-5: min(target=10, candidates=3) = 3.
        let candidates: Vec<_> = (0..3)
            .map(|i| make_candidate(i + 1, (3 - i) as f64, Some(i + 1), Some("video")))
            .collect();
        let constraints = DiversityConstraints::new().max_per_creator(1);
        let result = DiversitySelector::select(&candidates, &constraints, 10);

        assert_eq!(result.selected.len(), 3);
    }

    #[test]
    fn select_score_order_preserved() {
        // 10 candidates with distinct scores, all different creators, max_per_creator=1.
        let candidates: Vec<_> = (0..10)
            .map(|i| make_candidate(i + 1, (100 - i * 7) as f64, Some(i + 1), Some("video")))
            .collect();
        let constraints = DiversityConstraints::new().max_per_creator(1);
        let result = DiversitySelector::select(&candidates, &constraints, 10);

        assert_eq!(result.selected.len(), 10);
        // Scores should be in non-increasing order.
        for window in result.selected.windows(2) {
            assert!(
                window[0].score >= window[1].score,
                "scores not in descending order: {} < {}",
                window[0].score,
                window[1].score
            );
        }
    }

    // -- Property test helpers ----------------------------------------------

    fn build_candidates_by_creator(
        n_creators: usize,
        items_per_creator: usize,
    ) -> Vec<ScoredCandidate> {
        let mut candidates = Vec::new();
        let mut score = 1000.0_f64;
        for creator in 0..n_creators {
            for _ in 0..items_per_creator {
                #[allow(clippy::cast_possible_truncation)]
                let entity_id = EntityId::new(score as u64);
                candidates.push(ScoredCandidate {
                    entity_id,
                    score,
                    signal_snapshot: vec![],
                    creator_id: Some(EntityId::new(creator as u64 + 1)),
                    format: Some("video".into()),
                });
                score -= 1.0;
            }
        }
        candidates
    }

    fn build_candidates_format_split(n: usize) -> Vec<ScoredCandidate> {
        (0..n)
            .map(|i| ScoredCandidate {
                entity_id: EntityId::new(i as u64 + 1),
                score: (n - i) as f64,
                signal_snapshot: vec![],
                creator_id: Some(EntityId::new(i as u64 + 1)),
                format: Some(if i % 2 == 0 {
                    "video".into()
                } else {
                    "audio".into()
                }),
            })
            .collect()
    }

    fn build_candidates_single_creator(n: usize) -> Vec<ScoredCandidate> {
        (0..n)
            .map(|i| ScoredCandidate {
                entity_id: EntityId::new(i as u64 + 1),
                score: (n - i) as f64,
                signal_snapshot: vec![],
                creator_id: Some(EntityId::new(1)),
                format: Some("video".into()),
            })
            .collect()
    }

    // -- Property tests -----------------------------------------------------

    mod proptests {
        use super::*;
        use proptest::prelude::*;

        proptest! {
            #[test]
            fn prop_max_per_creator_respected_when_satisfied(
                n_creators in 1usize..20,
                items_per_creator in 1usize..10,
                max_per_creator in 1usize..5,
                target in 1usize..50,
            ) {
                let candidates = build_candidates_by_creator(n_creators, items_per_creator);
                let constraints = DiversityConstraints::new().max_per_creator(max_per_creator);
                let result = DiversitySelector::select(&candidates, &constraints, target);

                // INV-RANK-5: result count == min(target, candidates.len()).
                let expected_count = target.min(candidates.len());
                prop_assert_eq!(result.selected.len(), expected_count);

                // When constraints_satisfied: max_per_creator not exceeded.
                if result.constraints_satisfied {
                    let mut counts: HashMap<u64, usize> = HashMap::new();
                    for c in &result.selected {
                        if let Some(cid) = c.creator_id {
                            *counts.entry(cid.as_u64()).or_insert(0) += 1;
                        }
                    }
                    for (_cid, count) in counts {
                        prop_assert!(count <= max_per_creator);
                    }
                }
            }

            #[test]
            fn prop_format_mix_not_exceeded_when_satisfied(
                n_items in 10usize..100,
                max_frac_pct in 30usize..90,
                target in 5usize..30,
            ) {
                let max_frac = max_frac_pct as f64 / 100.0;
                let candidates = build_candidates_format_split(n_items);
                let constraints = DiversityConstraints::new().format_mix(max_frac);
                let result = DiversitySelector::select(&candidates, &constraints, target);

                prop_assert_eq!(result.selected.len(), target.min(candidates.len()));

                if result.constraints_satisfied && !result.selected.is_empty() {
                    let total = result.selected.len();
                    let mut fmt_counts: HashMap<&str, usize> = HashMap::new();
                    for c in &result.selected {
                        if let Some(ref fmt) = c.format {
                            *fmt_counts.entry(fmt.as_str()).or_insert(0) += 1;
                        }
                    }
                    for (_fmt, count) in fmt_counts {
                        prop_assert!(
                            count as f64 / total as f64 <= max_frac + f64::EPSILON
                        );
                    }
                }
            }

            #[test]
            fn prop_result_count_invariant(
                n in 0usize..200,
                target in 0usize..100,
                max_per_creator in 1usize..5,
            ) {
                // INV-RANK-5: always min(target, n) results.
                let candidates = build_candidates_single_creator(n);
                let constraints = DiversityConstraints::new().max_per_creator(max_per_creator);
                let result = DiversitySelector::select(&candidates, &constraints, target);
                prop_assert_eq!(result.selected.len(), target.min(n));
            }

            #[test]
            fn prop_unsatisfiable_still_fills(
                n in 5usize..50,
                target in 3usize..20,
            ) {
                // All same creator, max_per_creator=1 -- unsatisfiable if target > 1.
                let candidates = build_candidates_single_creator(n);
                let constraints = DiversityConstraints::new().max_per_creator(1);
                let result = DiversitySelector::select(&candidates, &constraints, target);
                // INV-RANK-5: always filled.
                prop_assert_eq!(result.selected.len(), target.min(n));
            }

            #[test]
            fn prop_score_order_within_creator(
                n_creators in 2usize..10,
                items_per_creator in 2usize..8,
                max_per_creator in 1usize..4,
            ) {
                // INV-RANK-6: within each creator's selected items, score order is preserved.
                let candidates = build_candidates_by_creator(n_creators, items_per_creator);
                let constraints = DiversityConstraints::new().max_per_creator(max_per_creator);
                let target = n_creators * items_per_creator;
                let result = DiversitySelector::select(&candidates, &constraints, target);

                let mut creator_scores: HashMap<u64, Vec<f64>> = HashMap::new();
                for c in &result.selected {
                    if let Some(cid) = c.creator_id {
                        creator_scores.entry(cid.as_u64()).or_default().push(c.score);
                    }
                }
                for scores in creator_scores.values() {
                    for window in scores.windows(2) {
                        prop_assert!(
                            window[0] >= window[1] - f64::EPSILON,
                            "scores not in desc order: {} < {}",
                            window[0],
                            window[1]
                        );
                    }
                }
            }
        }
    }
}
