mod api;
mod commands;
mod listener;

use ::serenity::all::{GatewayIntents, UserId};
use dotenvy::dotenv;
use poise::{Framework, FrameworkOptions, serenity_prelude as serenity};
use std::{collections::HashSet, env, sync::Arc};
use surrealdb::Surreal;
use surrealdb::engine::remote::ws::{Client, Wss};
use surrealdb::opt::auth::Root;
use tracing::info;

type Error = anyhow::Error;
use tracing_subscriber::{EnvFilter, FmtSubscriber};
type Context<'a> = poise::Context<'a, Data, Error>;

struct Data {
    db: Surreal<Client>,
    ai_chat: Arc<commands::fun::AiChatManager>,
}

#[tokio::main()]
async fn main() -> Result<(), Error> {
    dotenv().ok();

    let subscriber = FmtSubscriber::builder()
        .with_env_filter(
            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
        )
        .finish();

    tracing::subscriber::set_global_default(subscriber)?;

    let token = env::var("DISCORD_TOKEN")?;

    let api_key = env::var("API_KEY").expect("Expected API_KEY in environment");

    let surreal_address =
        env::var("SURREAL_ADDRESS").expect("Expected SURREAL_ADDRESS in environment");
    let surreal_user = env::var("SURREAL_USER").expect("Expected SURREAL_USER in environment");
    let surreal_pass = env::var("SURREAL_PASS").expect("Expected SURREAL_PASS in environment");
    let surreal_ns = env::var("SURREAL_NS").expect("Expected SURREAL_NS in environment");
    let surreal_db = env::var("SURREAL_DB").expect("Expected SURREAL_DB in environment");

    let ollama_url = env::var("OLLAMA_SERVER_URL").expect("Expected OLLAMA_SERVER_URL in environment");
    let ollama_model = env::var("OLLAMA_MODEL").unwrap_or_else(|_| "llama3".to_string());
    let ai_chat_cooldown_ms = env::var("AI_CHAT_COOLDOWN_MS")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
        .unwrap_or(1500);

    let ignore_rude: Vec<String> = env::var("IGNORE_RUDE")
        .unwrap_or_default()
        .split(',')
        .map(|s| s.trim())
        .filter(|s| !s.is_empty())
        .map(|s| s.to_string())
        .collect();

    let ai_chat_manager = Arc::new(commands::fun::AiChatManager::new(
        ollama_url,
        ollama_model,
        std::time::Duration::from_millis(ai_chat_cooldown_ms),
        ignore_rude,
    ));

    let db = Surreal::new::<Wss>(&surreal_address).await?;

    db.signin(Root {
        username: &surreal_user,
        password: &surreal_pass,
    })
    .await?;
    info!("Successfully signed in to SurrealDB");

    db.use_ns(&surreal_ns).use_db(&surreal_db).await?;

    let db_clone = db.clone();
    let ai_chat_clone = ai_chat_manager.clone();

    let owner_id = env::var("BOT_OWNER_ID")
        .expect("Expected BOT_OWNER_ID in environment")
        .parse::<u64>()?;
    let mut owners = HashSet::new();
    owners.insert(UserId::new(owner_id));

    let framework = Framework::builder()
        .options(FrameworkOptions::<Data, Error> {
            owners,
            commands: vec![
                commands::level::leaderboard(),
                commands::level::set_levelup_message_channel(),
                commands::level::set_levelup_message(),
                commands::fun::say(),
                commands::fun::urban(),
                commands::fun::ai_chat(),
                commands::utility::auto_response(),
                commands::utility::view_auto_responses(),
                commands::utility::delete_auto_response(),
                commands::utility::edit_auto_response(),
                commands::utility::summary(),
                commands::beta::add_beta_testing(),
                commands::beta::list_beta_testing(),
                commands::beta::remove_beta_testing(),
            ],
            prefix_options: poise::PrefixFrameworkOptions {
                prefix: Some("!".into()),
                edit_tracker: Some(Arc::new(poise::EditTracker::for_timespan(
                    std::time::Duration::from_secs(600),
                ))),
                case_insensitive_commands: true,
                ..Default::default()
            },
            ..Default::default()
        })
        .setup(move |context, _ready, framework| {
            Box::pin(async move {
                poise::builtins::register_globally(context, &framework.options().commands).await?;
                Ok(Data {
                    db: db_clone,
                    ai_chat: ai_chat_clone,
                })
            })
        })
        .build();

    let mut client = serenity::Client::builder(token, GatewayIntents::all())
        .event_handler(listener::Handler)
        .framework(framework)
        .await?;

    {
        let mut data = client.data.write().await;
        data.insert::<commands::level::DbKey>(db.clone());
        data.insert::<commands::utility::DbKey>(db.clone());
        data.insert::<commands::fun::AiChatKey>(ai_chat_manager.clone());
    }

    // Get API port from environment, default to 8080
    let api_port = env::var("API_PORT")
        .ok()
        .and_then(|v| v.parse::<u16>().ok())
        .unwrap_or(8080);

    // Run bot and API server concurrently
    let cache = client.cache.clone();
    let db_for_api = db.clone();
    let bot_task = tokio::spawn(async move {
        if let Err(why) = client.start_autosharded().await {
            eprintln!("An error occurred while running the client: {why}");
        }
    });

    let api_task = api::start_api_server(cache, db_for_api, api_key, api_port);

    // Wait for API to start, then run bot
    tokio::select! {
        _ = bot_task => {
            info!("Bot task finished");
        }
        result = api_task => {
            match result {
                Ok(_) => info!("API server finished"),
                Err(e) => eprintln!("API server error: {}", e),
            }
        }
    }

    Ok(())
}
