path: root/src/main.rs

use std::collections::{HashMap, HashSet};
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 music;
mod net;
mod openrouter;
mod recording;
mod tasks;
mod tts;
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);

/// Hard caps so autocomplete traffic cannot grow memory without bound.
const AC_CACHE_MAX_ENTRIES: usize = 2_000;
const AC_SEARCH_CACHE_MAX_ENTRIES: usize = 256;

/// 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 queue and current track per guild.
    pub music: Arc<music::MusicState>,
    /// TTS engine: piper model + audio cache.
    pub tts: Arc<tts::TtsState>,
}

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,
    ) {
        let mut cache = self.ac_cache.lock().await;
        if cache.len() >= AC_CACHE_MAX_ENTRIES {
            // Entries are only useful for the command invocation that
            // triggered the autocomplete, so dropping all of them is safe.
            cache.clear();
        }
        cache.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 mut cache = self.ac_search_cache.lock().await;
        let (hits, cached_at) = cache.get(query)?;
        if cached_at.elapsed() > AC_SEARCH_CACHE_TTL {
            cache.remove(query);
            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>) {
        let mut cache = self.ac_search_cache.lock().await;
        cache.retain(|_, (_, cached_at)| cached_at.elapsed() <= AC_SEARCH_CACHE_TTL);
        if cache.len() >= AC_SEARCH_CACHE_MAX_ENTRIES
            && let Some(oldest) = cache
                .iter()
                .min_by_key(|(_, (_, cached_at))| *cached_at)
                .map(|(query, _)| query.clone())
        {
            cache.remove(&oldest);
        }
        cache.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(&format!(
                    "{},songbird::driver::tasks::udp_rx=off",
                    // Known songbird noise: opus InvalidPacket on some
                    // incoming SSRCs; harmless for recording, huge span dumps.
                    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?;
    // Connect timeout only: body timeouts are set per-call where a finite
    // deadline makes sense (LLM/search/web), while media downloads may
    // legitimately take a long time.
    let http_client = reqwest::Client::builder()
        .connect_timeout(Duration::from_secs(10))
        .build()?;

    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(())
                })
            },
            on_error: |error| {
                Box::pin(async move {
                    match error {
                        poise::FrameworkError::Command { ctx, error, .. } => {
                            tracing::error!(command = %ctx.command().name, "Command failed: {error}");
                            if let Err(e) = ctx
                                .send(
                                    poise::CreateReply::default()
                                        .content(format!("Błąd: {error}"))
                                        .ephemeral(true),
                                )
                                .await
                            {
                                tracing::debug!(
                                    "Could not send error reply (interaction likely expired): {e}"
                                );
                            }
                        }
                        other => {
                            if let Err(e) = poise::builtins::on_error(other).await {
                                tracing::warn!(
                                    "Error while sending error message (interaction likely expired): {e}"
                                );
                            }
                        }
                    }
                })
            },
            ..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,
                    config.recording_enabled,
                    config.recording_window_seconds,
                )
                .await?;

                let music = music::MusicState::new(
                    vm.clone(),
                    ctx.http.clone(),
                    http_client.clone(),
                );

                let tts = Arc::new(tts::TtsState::new(config.piper_model_path()));

                let data = Data {
                    tts,
                    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,
                };

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

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

                // Discord keeps global commands that disappear from the
                // registration payload. Delete leftovers from renamed or
                // removed commands so the move under `/music` is visible.
                let expected: HashSet<String> = framework
                    .options()
                    .commands
                    .iter()
                    .map(|command| {
                        command
                            .context_menu_name
                            .clone()
                            .unwrap_or_else(|| command.name.clone())
                    })
                    .collect();
                match ctx.http.get_global_commands().await {
                    Ok(existing) => {
                        for command in existing {
                            if !expected.contains(&command.name) {
                                match ctx.http.delete_global_command(command.id).await {
                                    Ok(()) => tracing::info!(
                                        command = %command.name,
                                        "removed stale global command"
                                    ),
                                    Err(e) => tracing::warn!(
                                        command = %command.name,
                                        error = %e,
                                        "failed to remove stale global command"
                                    ),
                                }
                            }
                        }
                    }
                    Err(e) => {
                        tracing::warn!(error = %e, "could not list global commands for cleanup")
                    }
                }

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

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

    // Configure songbird: when recording is enabled, decode incoming voice to
    // mono 48 kHz i16 so that `CoreEvent::VoiceTick` carries `decoded_voice`
    // (without this, `decoded_voice` is `None` and there's nothing to record).
    let mut songbird_config = songbird::Config::default();
    if config_for_client.recording_enabled {
        songbird_config.decode_mode =
            songbird::driver::DecodeMode::Decode(songbird::driver::DecodeConfig::new(
                songbird::driver::Channels::Mono,
                songbird::driver::SampleRate::Hz48000,
            ));
    }

    let mut client =
        serenity::ClientBuilder::new(&config_for_client.token, serenity::GatewayIntents::all())
            .framework(framework)
            .register_songbird_from_config(songbird_config)
            .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));
        }
        _ => {}
    }
}