//! Candidate generation strategies for the RETRIEVE executor.
//!
//! Contains the scan, signal-ranked, and exploration injection strategies
//! used in Stage 1 of the RETRIEVE pipeline.

use std::collections::HashSet;
use std::sync::RwLock;

use roaring::RoaringBitmap;

use crate::ranking::executor::ScoredCandidate;
use crate::schema::{EntityId, Timestamp};
use crate::signals::SignalLedger;

/// Scan the universe bitmap for all entity IDs.
///
/// When `has_user_context` is true, uses a larger multiplier (10x) to
/// ensure enough candidates survive user-context filtering (seen, blocked,
/// hard negatives).
pub(crate) fn scan_candidates(
    universe: Option<&RwLock<RoaringBitmap>>,
    limit: usize,
    has_user_context: bool,
) -> Vec<EntityId> {
    let multiplier = if has_user_context { 10 } else { 4 };
    let max_candidates = (limit * multiplier).max(200);
    let mut candidates = Vec::with_capacity(max_candidates);

    // Read-lock the universe bitmap.
    if let Some(universe) = universe
        && let Ok(bm) = universe.read()
    {
        for id_u32 in bm.iter() {
            candidates.push(EntityId::new(u64::from(id_u32)));
            if candidates.len() >= max_candidates {
                break;
            }
        }
    }

    candidates
}

/// Rank candidates by a specific signal's decay score.
pub(crate) fn signal_ranked_candidates(
    ledger: &SignalLedger,
    signal_name: &str,
    limit: usize,
) -> Vec<EntityId> {
    let max_candidates = (limit * 4).max(200);
    let Ok(type_id) = ledger.resolve_signal_type(signal_name) else {
        return Vec::new();
    };

    let now_ns = Timestamp::now().as_nanos();
    let mut scored: Vec<(EntityId, f64)> = Vec::new();

    for entry in ledger.entries() {
        let (entity_id, signal_type_id) = entry.key();
        if *signal_type_id == type_id {
            // Use decay score at index 0 with lambda=0 (no additional decay
            // beyond what was already applied at write time). This is a
            // simplified ranking for candidate generation -- the full scoring
            // happens in Stage 3 via ProfileExecutor.
            let score = entry.value().hot.current_score(0, now_ns, 0.0);
            scored.push((*entity_id, score));
        }
    }

    // O(N) partition to top-K via `select_nth_unstable_by`, then O(K log K)
    // sort of just the top K. Better than O(N log N) full sort when N >> limit.
    // NaN scores (from degenerate decay math) fall back to entity-ID order for
    // deterministic output.
    let cmp = |a: &(EntityId, f64), b: &(EntityId, f64)| {
        b.1.partial_cmp(&a.1)
            .unwrap_or_else(|| b.0.as_u64().cmp(&a.0.as_u64()))
    };
    if scored.len() > max_candidates {
        scored.select_nth_unstable_by(max_candidates - 1, cmp);
        scored.truncate(max_candidates);
    }
    scored.sort_unstable_by(cmp);

    scored.into_iter().map(|(id, _)| id).collect()
}

/// Inject exploration candidates into the scored list.
///
/// Reserves `exploration_fraction` of the result set for candidates
/// not in the top-scored results. Used by `for_you` to prevent filter
/// bubbles and expose cold-start content.
///
/// Exploration candidates are appended at the end of the scored list
/// with score 0.0 (lowest). The selection is deterministic: a hash of
/// the candidate set size + first candidate ID seeds the order.
pub(crate) fn inject_exploration(
    scored: &mut Vec<ScoredCandidate>,
    all_candidates: &[EntityId],
    exploration_fraction: f64,
) {
    if scored.is_empty() || all_candidates.is_empty() || exploration_fraction <= 0.0 {
        return;
    }

    #[allow(
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss,
        clippy::cast_precision_loss
    )]
    let exploration_slots = (exploration_fraction * scored.len() as f64).ceil() as usize;
    if exploration_slots == 0 {
        return;
    }

    // Collect entity IDs already in the scored set.
    let scored_ids: HashSet<u64> = scored.iter().map(|c| c.entity_id.as_u64()).collect();

    // Find candidates not in the scored set.
    let mut exploration_pool: Vec<EntityId> = all_candidates
        .iter()
        .filter(|id| !scored_ids.contains(&id.as_u64()))
        .copied()
        .collect();

    if exploration_pool.is_empty() {
        return;
    }

    // Deterministic shuffle using BLAKE3 hash of candidate IDs.
    // This is reproducible for the same candidate set.
    // `sort_unstable_by` is safe: hash-based ordering is not stability-sensitive.
    exploration_pool.sort_unstable_by(|a, b| {
        let hash_a = blake3::hash(&a.as_u64().to_le_bytes());
        let hash_b = blake3::hash(&b.as_u64().to_le_bytes());
        hash_a.as_bytes().cmp(hash_b.as_bytes())
    });

    let actual_slots = exploration_slots.min(exploration_pool.len());

    // Trim scored to make room for exploration candidates.
    let keep = scored.len().saturating_sub(actual_slots);
    scored.truncate(keep);

    // Append exploration candidates with score 0.0.
    for &entity_id in exploration_pool.iter().take(actual_slots) {
        scored.push(ScoredCandidate {
            entity_id,
            score: 0.0,
            signal_snapshot: vec![],
            creator_id: None,
            format: None,
        });
    }
}

// ── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::ranking::executor::ScoredCandidate;
    use crate::schema::EntityId;

    #[test]
    fn exploration_injects_random_candidates() {
        // Scored list: entities 1-5 (the "top ranked" items).
        let mut scored: Vec<ScoredCandidate> = (1..=5)
            .map(|i| ScoredCandidate {
                entity_id: EntityId::new(i),
                score: 1.0 - (i as f64 * 0.1),
                signal_snapshot: vec![],
                creator_id: None,
                format: None,
            })
            .collect();

        // Full candidate list: entities 1-10 (so 6-10 are available for exploration).
        let all_candidates: Vec<EntityId> = (1..=10).map(EntityId::new).collect();

        // 40% exploration -> ceil(0.4 * 5) = 2 exploration slots.
        inject_exploration(&mut scored, &all_candidates, 0.4);

        assert_eq!(scored.len(), 5, "total output should remain 5");

        // At least one of the exploration candidates (6-10) should be present.
        let exploration_ids: Vec<u64> = scored
            .iter()
            .filter(|c| c.entity_id.as_u64() > 5)
            .map(|c| c.entity_id.as_u64())
            .collect();
        assert_eq!(
            exploration_ids.len(),
            2,
            "expected 2 exploration candidates, got {}",
            exploration_ids.len()
        );

        // Exploration candidates should have score 0.0.
        for c in &scored {
            if c.entity_id.as_u64() > 5 {
                assert_eq!(c.score, 0.0, "exploration candidates should have score 0.0");
            }
        }
    }

    #[test]
    fn exploration_zero_fraction_is_noop() {
        let mut scored: Vec<ScoredCandidate> = (1..=5)
            .map(|i| ScoredCandidate {
                entity_id: EntityId::new(i),
                score: 1.0,
                signal_snapshot: vec![],
                creator_id: None,
                format: None,
            })
            .collect();

        let all: Vec<EntityId> = (1..=10).map(EntityId::new).collect();
        inject_exploration(&mut scored, &all, 0.0);

        // No change when exploration fraction is 0.
        assert_eq!(scored.len(), 5);
        assert!(scored.iter().all(|c| c.entity_id.as_u64() <= 5));
    }

    #[test]
    fn exploration_no_extra_candidates_is_noop() {
        // All candidates are already scored -- no exploration pool.
        let mut scored: Vec<ScoredCandidate> = (1..=5)
            .map(|i| ScoredCandidate {
                entity_id: EntityId::new(i),
                score: 1.0,
                signal_snapshot: vec![],
                creator_id: None,
                format: None,
            })
            .collect();

        let all: Vec<EntityId> = (1..=5).map(EntityId::new).collect();
        inject_exploration(&mut scored, &all, 0.5);

        // No change when there are no extra candidates for exploration.
        assert_eq!(scored.len(), 5);
    }
}
