use std::collections::HashMap;

use crate::labels::{ForageItem, ItemLabel};
use crate::seed::SeedItem;

/// Per-user exploration instrumentation for adaptive MAB.
///
/// Tracks exploration hit rate (positive signals on explore-labeled items)
/// and per-category signal counts for UCB1 bonus computation.
#[derive(Debug, Default, Clone, serde::Serialize, serde::Deserialize)]
pub struct ExplorationStats {
    /// Number of exploration-slotted items the user signaled positively (view/save/share).
    pub exploration_hits: u32,
    /// Total number of exploration-slotted items that appeared in user feeds.
    pub exploration_total: u32,
    /// Per-category signal counts for UCB1 category bonus.
    pub category_signals: HashMap<String, u32>,
}

impl ExplorationStats {
    /// Compute the adaptive exploration ratio based on exploration hit rate.
    ///
    /// | Hit rate  | Ratio | Interpretation                                |
    /// |-----------|-------|-----------------------------------------------|
    /// | > 0.5     | 0.25  | Adventurous — user engages with novel content |
    /// | < 0.2     | 0.10  | Convergent — user prefers familiar content    |
    /// | otherwise | 0.14  | Default                                       |
    ///
    /// Returns 0.14 when no exploration data exists (cold start).
    pub fn adaptive_ratio(&self) -> f32 {
        if self.exploration_total == 0 {
            return 0.14;
        }
        let hit_rate = self.hit_rate();
        if hit_rate > 0.5 {
            0.25
        } else if hit_rate < 0.2 {
            0.10
        } else {
            0.14
        }
    }

    /// Record the outcome of an exploration-slotted item being signaled.
    ///
    /// Call this when the user signals an item that was labeled `Exploring`.
    /// `hit = true` for positive signals (view, save, share); `false` for skip.
    pub fn record_outcome(&mut self, hit: bool) {
        self.exploration_total += 1;
        if hit {
            self.exploration_hits += 1;
        }
    }

    /// Increment the signal count for a category.
    pub fn record_category_signal(&mut self, category: &str) {
        *self
            .category_signals
            .entry(category.to_owned())
            .or_insert(0) += 1;
    }

    /// Compute UCB1 exploration bonus for a given category.
    ///
    /// Categories with fewer signals receive a higher bonus, naturally surfacing
    /// underexplored content into the candidate pool:
    ///
    ///   bonus = sqrt(2 × ln(total_signals) / (category_count + 1))
    ///
    /// The `+1` denominator prevents division by zero for unseen categories.
    /// Returns 0.0 if no signals have been recorded (cold start).
    pub fn ucb1_bonus(&self, category: &str) -> f64 {
        let total: u32 = self.category_signals.values().sum();
        if total == 0 {
            return 0.0;
        }
        let category_count = self.category_signals.get(category).copied().unwrap_or(0);
        (2.0 * (total as f64).ln() / (category_count as f64 + 1.0)).sqrt()
    }

    /// Exploration hit rate (0.0 for cold start).
    pub fn hit_rate(&self) -> f32 {
        if self.exploration_total == 0 {
            0.0
        } else {
            self.exploration_hits as f32 / self.exploration_total as f32
        }
    }
}

/// Static MAB configuration (shared across all users).
pub struct MabConfig {
    /// Base exploration ratio when no per-user history exists (default 0.14 ≈ 1/7).
    ///
    /// At runtime the per-user `ExplorationStats::adaptive_ratio()` replaces this
    /// value once the user has accumulated enough exploration history.
    pub exploration_ratio: f32,
    /// Weight of semantic similarity in the warm-path ranking blend (0.0–1.0).
    ///
    /// The final score is:
    ///   `semantic_boost * semantic_similarity + (1 - semantic_boost) * rank_score`
    ///
    /// At the default 0.3: 30% of the score comes from ANN similarity to the
    /// user's preference centroid and 70% from position in the retrieval ranking.
    /// Set to 1.0 for pure-semantic ranking; 0.0 for pure-rank.
    pub semantic_boost: f32,
    /// When `true` and the user has saved items, semantically similar candidates
    /// are injected into the feed pool before MAB selection, boosting items near
    /// the user's saved content.
    pub similar_to_saved: bool,
}

impl Default for MabConfig {
    fn default() -> Self {
        Self {
            exploration_ratio: 0.14,
            semantic_boost: 0.3,
            similar_to_saved: true,
        }
    }
}

/// Epsilon-greedy item selection with per-user adaptive exploration ratio and UCB1.
///
/// Given a ranked pool (highest score first), fills `limit` slots:
/// - **Exploit slots**: top items from the pool, labeled `Match` or `Trending`.
///   These are driven purely by preference-signal scores and are never re-ordered
///   by UCB1 (preserving the user's learned preferences).
/// - **Explore slots**: items from underrepresented categories at the bottom of
///   the pool, labeled `Exploring`.  When `stats` is provided, UCB1 bonus scores
///   are applied *within* the explore candidates only, so underexplored categories
///   bubble up into the exploration slot without disrupting the exploit ranking.
///
/// `exploration_ratio` is the per-user adaptive ratio from
/// `ExplorationStats::adaptive_ratio()`.  Pass `MabConfig::exploration_ratio`
/// for cold-start users.  `stats` is `None` for cold-start users or when UCB1
/// guidance is not desired.
pub fn select(
    pool: &[tidaldb::query::retrieve::RetrieveResult],
    item_meta: &HashMap<u64, SeedItem>,
    limit: usize,
    exploration_ratio: f32,
    stats: Option<&ExplorationStats>,
) -> Vec<ForageItem> {
    if pool.is_empty() {
        return vec![];
    }

    let explore_count = ((limit as f32 * exploration_ratio).ceil() as usize).max(1);
    let exploit_count = limit.saturating_sub(explore_count);

    let pool_len = pool.len();
    let explore_start = pool_len.saturating_sub(explore_count);

    let exploit_pool = &pool[..explore_start.min(exploit_count + explore_count)];
    let explore_pool = &pool[explore_start..];

    let mut result: Vec<ForageItem> = Vec::with_capacity(limit);

    // Exploit: take top-ranked items (no UCB1 — exploit respects user preference).
    for (i, rr) in exploit_pool.iter().take(exploit_count).enumerate() {
        let id = rr.entity_id.as_u64();
        if let Some(meta) = item_meta.get(&id) {
            let label = if i < 3 {
                ItemLabel::Match
            } else if rr.score > 0.5 {
                ItemLabel::Trending
            } else {
                ItemLabel::Resurfaced
            };
            result.push(ForageItem {
                id,
                title: meta.title.clone(),
                source: meta.source.clone(),
                category: meta.category.clone(),
                reading_time_min: meta.reading_time_min,
                description: meta.description.clone(),
                label,
                score: rr.score as f32,
                url: meta.url.clone(),
                tags: vec![],
                entities: vec![],
                content_type: String::new(),
                summary: String::new(),
            });
        }
    }

    // Explore: pick items from the tail of the pool.
    // When stats are provided, UCB1 bonus scores guide selection within the
    // explore candidates — underexplored categories get priority here.
    let already_in_result: std::collections::HashSet<u64> = result.iter().map(|fi| fi.id).collect();
    let mut explore_candidates: Vec<(&tidaldb::query::retrieve::RetrieveResult, f64)> =
        explore_pool
            .iter()
            .filter(|rr| !already_in_result.contains(&rr.entity_id.as_u64()))
            .map(|rr| {
                let ucb1 = stats
                    .and_then(|s| {
                        item_meta
                            .get(&rr.entity_id.as_u64())
                            .map(|m| s.ucb1_bonus(&m.category))
                    })
                    .unwrap_or(0.0);
                (rr, rr.score + ucb1)
            })
            .collect();
    // Sort explore candidates by UCB1-adjusted score (highest first).
    explore_candidates.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    for (rr, _) in explore_candidates.iter().take(explore_count) {
        let id = rr.entity_id.as_u64();
        if let Some(meta) = item_meta.get(&id) {
            result.push(ForageItem {
                id,
                title: meta.title.clone(),
                source: meta.source.clone(),
                category: meta.category.clone(),
                reading_time_min: meta.reading_time_min,
                description: meta.description.clone(),
                label: ItemLabel::Exploring,
                score: rr.score as f32,
                url: meta.url.clone(),
                tags: vec![],
                entities: vec![],
                content_type: String::new(),
                summary: String::new(),
            });
        }
    }

    // If we didn't fill limit (pool smaller than expected), pad from exploit overflow.
    let used: std::collections::HashSet<u64> = result.iter().map(|fi| fi.id).collect();
    let remaining = limit.saturating_sub(result.len());
    if remaining > 0 {
        for rr in pool.iter().skip(exploit_count).take(remaining) {
            let id = rr.entity_id.as_u64();
            if let Some(meta) = item_meta.get(&id)
                && !used.contains(&id)
            {
                result.push(ForageItem {
                    id,
                    title: meta.title.clone(),
                    source: meta.source.clone(),
                    category: meta.category.clone(),
                    reading_time_min: meta.reading_time_min,
                    description: meta.description.clone(),
                    label: ItemLabel::Resurfaced,
                    score: rr.score as f32,
                    url: meta.url.clone(),
                    tags: vec![],
                    entities: vec![],
                    content_type: String::new(),
                    summary: String::new(),
                });
            }
        }
    }

    result.truncate(limit);
    result
}
