path: root/src/main.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use poise::serenity_prelude as serenity;
use songbird::SerenityInit;
use tracing_subscriber::filter::EnvFilter;

mod commands;
mod config;
mod copyparty;
mod database;
mod error;
mod events;
mod lang;
mod openrouter;
mod tasks;
mod voice;

use database::DbPool;

/// How long a search cache entry stays fresh (10 seconds).
const AC_SEARCH_CACHE_TTL: Duration = Duration::from_secs(10);

/// Per-user autocomplete stash: (UserId, filename) → SearchHit.
type AcCache = tokio::sync::Mutex<HashMap<(serenity::UserId, String), copyparty::SearchHit>>;

/// Shared autocomplete search cache: query → (hits, cached_at).
type AcSearchCache = Arc<
    tokio::sync::Mutex<HashMap<String, (Vec<copyparty::SearchHit>, Instant)>>,
>;

pub struct Data {
    pub db_pool: DbPool,
    pub config: config::Config,
    pub voice_manager: voice::VoiceManager,
    pub copyparty: copyparty::CopypartyClient,
    pub http_client: reqwest::Client,
    pub volume: std::sync::atomic::AtomicU64,
    pub triggers: tokio::sync::Mutex<Vec<(database::Trigger, Vec<String>, Vec<String>)>>,
    /// Autocomplete lookup, keyed per-user to avoid cross-command collisions.
    ///
    /// Mirrors Python's per-cog `_cached_urls[user_id]` map: when two users
    /// search the same filename simultaneously, each invocation must resolve its
    /// own stashed `SearchHit`, so the cache is keyed by `(UserId, filename)`.
    pub ac_cache: AcCache,
    /// Autocomplete search results cache (query → hits).
    ///
    /// Populated by background tasks when the first search times out, so the
    /// next keystroke's autocomplete can serve real results immediately.
    pub ac_search_cache: AcSearchCache,
    /// Active music track per guild (only one music source at a time).
    pub music_tracks: tokio::sync::Mutex<HashMap<u64, songbird::tracks::TrackHandle>>,
}

impl Data {
    pub async fn reload_triggers(&self) {
        let triggers = database::load_all_triggers_with_keywords_and_responses(&self.db_pool)
            .await
            .unwrap_or_default();
        *self.triggers.lock().await = triggers;
    }

    pub async fn ac_stash(
        &self,
        user_id: serenity::UserId,
        key: String,
        hit: copyparty::SearchHit,
    ) {
        self.ac_cache.lock().await.insert((user_id, key), hit);
    }

    pub async fn ac_lookup(
        &self,
        user_id: serenity::UserId,
        key: &str,
    ) -> Option<copyparty::SearchHit> {
        self.ac_cache
            .lock()
            .await
            .get(&(user_id, key.to_string()))
            .cloned()
    }

    /// Look up search cache for an autocomplete query. Returns `None` if
    /// uncached or stale.
    pub async fn ac_search_cache_lookup(
        &self,
        query: &str,
    ) -> Option<Vec<copyparty::SearchHit>> {
        let cache = self.ac_search_cache.lock().await;
        let (hits, cached_at) = cache.get(query)?;
        if cached_at.elapsed() > AC_SEARCH_CACHE_TTL {
            None
        } else {
            Some(hits.clone())
        }
    }

    /// Store search results into the shared cache.
    pub async fn ac_search_cache_store(&self, query: String, hits: Vec<copyparty::SearchHit>) {
        self.ac_search_cache
            .lock()
            .await
            .insert(query, (hits, Instant::now()));
    }
}

#[allow(dead_code)]
type Error = error::Error;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let _ = rustls::crypto::ring::default_provider().install_default();

    tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::builder()
                .parse(std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()))
                .unwrap(),
        )
        .init();

    dotenvy::dotenv().ok();
    config::check_env_vars();
    let config = config::Config::from_env();

    let db_pool = database::init_pool(&config.database_path).await?;
    let http_client = reqwest::Client::new();

    let copyparty = copyparty::CopypartyClient::new(
        config.copyparty_url.clone(),
        config.copyparty_url_upload.clone(),
        config.copyparty_password.clone(),
    );

    let triggers = database::load_all_triggers_with_keywords_and_responses(&db_pool)
        .await
        .unwrap_or_default();

    let commands_list = commands::all_commands();
    let config_for_client = config.clone();

    let framework = poise::Framework::builder()
        .options(poise::FrameworkOptions {
            commands: commands_list,
            event_handler: |ctx, event, _framework, user_data| {
                Box::pin(async move {
                    handle_event(ctx, event, user_data).await;
                    Ok(())
                })
            },
            ..Default::default()
        })
        .setup(move |ctx, _ready, framework| {
            let config = config.clone();
            let db_pool = db_pool.clone();
            let http_client = http_client.clone();
            let copyparty = copyparty.clone();
            let triggers = triggers.clone();
            Box::pin(async move {
                let vm = voice::VoiceManager::new(ctx).await?;

                let data = Data {
                    config,
                    db_pool,
                    voice_manager: vm,
                    copyparty,
                    http_client,
                    volume: std::sync::atomic::AtomicU64::new(1.0f64.to_bits()),
                    triggers: tokio::sync::Mutex::new(triggers),
                    ac_cache: tokio::sync::Mutex::new(HashMap::new()),
                    ac_search_cache: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
                    music_tracks: tokio::sync::Mutex::new(HashMap::new()),
                };

                // Spawn cron scheduler loop.
                {
                    let cron_ctx = crate::tasks::CronContext {
                        pool: data.db_pool.clone(),
                        http: ctx.http.clone(),
                        voice_manager: data.voice_manager.clone(),
                        copyparty: data.copyparty.clone(),
                    };
                    tokio::spawn(async move {
                        crate::tasks::cron_scheduler_loop(cron_ctx).await;
                    });
                }

                poise::builtins::register_globally(ctx, &framework.options().commands).await?;

                events::handle_ready(ctx).await.map_err(|e| {
                    tracing::error!("on_ready handler: {e}");
                    e
                })?;

                Ok(data)
            })
        })
        .build();

    let mut client =
        serenity::ClientBuilder::new(&config_for_client.token, serenity::GatewayIntents::all())
            .framework(framework)
            .register_songbird()
            .await?;

    client.start().await?;
    Ok(())
}

async fn handle_event(ctx: &serenity::Context, event: &serenity::FullEvent, user_data: &Data) {
    match event {
        serenity::FullEvent::Ready {
            data_about_bot: _, ..
        } => {}
        serenity::FullEvent::VoiceStateUpdate { old, new } => {
            events::handle_voice_state_update(ctx, user_data, old.clone(), new)
                .await
                .unwrap_or_else(|e| tracing::error!("{}: {e}", lang::L_EV_VOICE_STATE_ERROR));
        }
        serenity::FullEvent::Message { new_message } => {
            events::handle_message(ctx, user_data, new_message)
                .await
                .unwrap_or_else(|e| tracing::error!("{}: {e}", lang::L_EV_MESSAGE_HANDLER_ERROR));
        }
        _ => {}
    }
}