botcat-capoo 1.0.3

use rand::{
    distr::{Distribution, weighted::WeightedIndex},
    rng,
    seq::IndexedRandom,
};

use crate::{
    api::send_message::{send_group_message, send_private_message},
    context::AppContext,
    error::Error,
    model::{
        event::{BotEvent, PokeEvent},
        message::segment::MessageSegment,
    },
};

pub async fn handle_poke(
    ctx: &AppContext,
    root: &BotEvent,
    event: &PokeEvent,
) -> Result<(), Error> {
    if event.target_id == root.self_id {
        let segment = MessageSegment::Text {
            text: draw_reply()?.to_string(),
        };
        let message = vec![segment];
        if let Some(group_id) = event.group_id {
            // Handle group poke
            send_group_message(&ctx.config, group_id, message).await;
            Ok(())
        } else {
            // Handle private poke
            send_private_message(&ctx.config, event.user_id, message).await;
            Ok(())
        }
    } else {
        // Do nothing if the bot is not the target
        Ok(())
    }
}

#[derive(Debug, Clone)]
enum Rarity {
    Common,
    Uncommon,
    Rare,
    Epic,
}

struct ReplyPool {
    common_replies: Vec<String>,
    uncommon_replies: Vec<String>,
    rare_replies: Vec<String>,
    epic_replies: Vec<String>,
}

impl ReplyPool {
    fn new() -> Self {
        Self {
            common_replies: vec![
                "喵？".to_string(), // 55%
            ],
            uncommon_replies: vec![
                "哈！".to_string(),         // 12.5%
                "呼噜～呼噜～".to_string(), // 12.5%
            ],
            rare_replies: vec![
                "嘎！".to_string(),   // 3%
                "呱！".to_string(),   // 3%
                "嗷呜！".to_string(), // 3%
                "咕～".to_string(),   // 3%
                "mua".to_string(),    // 3%
            ],
            epic_replies: vec![
                "哈基米～哈基米～".to_string(),    // 1%
                "哈基米喔那咩路多！".to_string(),  // 1%
                "咕咕嘎嘎咕咕嘎嘎！".to_string(),  // 1%
                "Das war ein Befehl!".to_string(), // 1%
                "我到河北省来！".to_string(),      // 1%
            ],
        }
    }

    pub fn random_reply(&self) -> Option<(Rarity, String)> {
        let mut rng = rng();

        let weights = [55, 25, 15, 5];
        let dist = WeightedIndex::new(weights).ok()?;
        let rarity_index = dist.sample(&mut rng);

        let (rarity, replies) = match rarity_index {
            0 => (Rarity::Common, &self.common_replies),
            1 => (Rarity::Uncommon, &self.uncommon_replies),
            2 => (Rarity::Rare, &self.rare_replies),
            3 => (Rarity::Epic, &self.epic_replies),
            _ => return None,
        };

        let reply = replies.choose(&mut rng)?.to_string();
        Some((rarity, reply))
    }
}

fn draw_reply() -> Result<String, Error> {
    let pool = ReplyPool::new();
    if let Some((_, reply)) = pool.random_reply() {
        Ok(reply)
    } else {
        Err(Error::Other {
            desc: "Failed to draw a random reply".to_string(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_draw_reply() {
        let mut counts = std::collections::HashMap::new();
        for _ in 0..10000 {
            let reply = draw_reply().expect("Failed to draw reply");
            *counts.entry(reply).or_insert(0) += 1;
        }
        for (reply, count) in counts {
            println!("{}\t| {}", count, reply);
        }
    }
}