tegakituesday.com/src/models/database.rs

use std::collections::HashSet;
use std::{fs::File, io::Read, collections::HashMap, path::Path};

use poise::serenity_prelude::{SerenityError, Http};
use r2d2::{Pool, PooledConnection};
use r2d2_sqlite::SqliteConnectionManager;
use r2d2_sqlite::rusqlite::{self, params};
use derive_more::From;

use crate::{utils::get_challenge_number, models::User};

use super::{LegacySubmission, Submission};

pub struct Database {
    // Must be Arc because Connection contains RefCell,
    // which cannot be shared between threads safely
    connection_pool: Pool<SqliteConnectionManager>,
}

const DATABASE_FILENAME: &str = "database.db";

#[derive(From, Debug)]
pub enum DatabaseError {
    Rusqlite(rusqlite::Error),
    Pool(r2d2::Error),
}

#[derive(From, Debug)]
pub enum LoadLegacyError {
    Database(DatabaseError),
    Serenity(SerenityError),
}

type Result<T> = std::result::Result<T, DatabaseError>;

impl Database {
    pub fn file_exists() -> bool {
        Path::new(DATABASE_FILENAME).exists()
    }

    pub fn new(
        testing: bool,
    ) -> Result<Self> {
        let connection_manager = if testing {
            SqliteConnectionManager::memory()
        } else {
            SqliteConnectionManager::file(DATABASE_FILENAME)
        };
        let connection_pool = Pool::new(connection_manager)?;
        let conn = connection_pool.get()?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS Submission (
                id        INTEGER PRIMARY KEY,
                author_id INTEGER NOT NULL,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                image     TEXT NOT NULL,
                challenge INTEGER NOT NULL,
                FOREIGN KEY (author_id) REFERENCES User(id)
            )",
            params![],
        )?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS User (
                id            INTEGER PRIMARY KEY,
                name          TEXT NOT NULL,
                discriminator INTEGER NOT NULL,
                avatar        TEXT,
                deleted       INTEGER NOT NULL
            )",
            params![],
        )?;
        Ok(Self { connection_pool })
    }

    fn conn(&self) -> std::result::Result<PooledConnection<SqliteConnectionManager>, r2d2::Error> {
        self.connection_pool.get()
    }

    pub fn has_submitted(&self, user_id: u64) -> Result<bool> {
        Ok(self.conn()?
            .prepare("SELECT 1 FROM User WHERE id = ?1 LIMIT 1")?
            .query_row(params![user_id], |_row| Ok(()))
            .is_ok())
    }

    pub async fn load_legacy(&self, http: &Http) -> std::result::Result<(), LoadLegacyError> {
        let latest_challenge = get_challenge_number();
        // HashMap of archived users that are no longer sharing a server with 字ちゃん
        // Their historical usernames and discriminators will be used
        let mut archived_users = HashMap::new();
        let conn = self.conn().map_err(DatabaseError::Pool)?;
        for n in 1..=latest_challenge {
            println!("Loading legacy challenge {n}/{latest_challenge}...");
            let mut file = File::open(format!("data/challenges/{n}.json")).unwrap();
            let mut contents = String::new();
            file.read_to_string(&mut contents).unwrap();
            for legacy in serde_json::from_str::<Vec<LegacySubmission>>(&contents).unwrap() {
                let id = legacy.id;
                if !self.has_submitted(id)? {
                    println!("Fetching user {id}...");
                    // If it already contains the archived user,
                    // overwrite write their data since they may have updated
                    // their username/discriminator since their previous submission
                    match archived_users.get(&id) {
                        Some(User { deleted, .. }) => {
                            archived_users.insert(id, User {
                                id,
                                avatar: None,
                                deleted: *deleted,
                                ..User::from_username(&legacy.username)
                            });
                        },
                        None => match User::fetch(http, id).await {
                            Ok(User { deleted: true, .. }) => {
                                archived_users.insert(id, User {
                                    id,
                                    avatar: None,
                                    deleted: true,
                                    ..User::from_username(&legacy.username)
                                });
                            },
                            Ok(user) => {
                                conn.execute(
                                    "INSERT INTO User(id, name, discriminator, avatar, deleted) VALUES (?1, ?2, ?3, ?4, ?5)",
                                    params![user.id, user.name, user.discriminator, user.avatar, user.deleted]
                                ).map_err(DatabaseError::Rusqlite)?;
                            },
                            Err(error) if error.to_string().eq("Unknown User") => {
                                // This will also be called in the case of an invalid user ID
                                println!("Failed to fetch user {id}, adding to archive");
                                archived_users.insert(id, User {
                                    id,
                                    avatar: None,
                                    deleted: false,
                                    ..User::from_username(&legacy.username)
                                });
                            },
                            Err(error) => return Err(LoadLegacyError::Serenity(error)),
                        },
                    };
                }
                for submission in legacy.parse(n as u32).unwrap() {
                    conn.execute(
                        "INSERT INTO Submission(author_id, timestamp, image, challenge) VALUES (?1, ?2, ?3, ?4)",
                        params![
                            &submission.author_id,
                            &submission.timestamp,
                            &submission.image,
                            n,
                        ]
                    ).map_err(DatabaseError::Rusqlite)?;
                }
            }
        }
        for (_id, user) in archived_users {
            conn.execute(
                "INSERT INTO USER (id, name, discriminator, avatar, deleted) VALUES (?1, ?2, ?3, ?4, ?5)",
                params![user.id, user.name, user.discriminator, user.avatar, user.deleted]
            ).map_err(DatabaseError::Rusqlite)?;
        }
        Ok(())
    }
    
    pub fn get_challenge_user_data(&self, challenge: u32) -> Result<(Vec<Submission>, HashMap<String, User>)> {
        let submissions = self.get_submissions(challenge)?;
        let users = self.get_users({
            let mut user_ids = HashSet::new();
            for submission in &submissions {
                user_ids.insert(submission.author_id);
            }
            user_ids
        })?;
        Ok((submissions, users))
    }

    pub fn get_submissions(&self, challenge: u32) -> Result<Vec<Submission>> {
        Ok(self.conn()?
            .prepare("SELECT author_id, timestamp, image FROM Submission WHERE challenge = ?1")?
            .query_map(params![challenge], |row| {
                Ok(Submission {
                    author_id: row.get(0)?,
                    timestamp: row.get(1)?,
                    image: row.get(2)?,
                    challenge,
                })
            })?
            .collect::<std::result::Result<Vec<Submission>, rusqlite::Error>>()?)
    }

    pub fn get_user_by_name(&self, name: &str) -> Result<Option<User>> {
        Ok(self.conn()?
            .prepare("SELECT id, discriminator, avatar, deleted FROM User where name = ?1")?
            .query_row(params![name], |row| {
                Ok(Some(User {
                    id: row.get(0)?,
                    name: name.to_owned(),
                    discriminator: row.get(1)?,
                    avatar: row.get(2)?,
                    deleted: row.get(3)?,
                }))
            })?)
    }

    pub fn get_submissions_by_user_name(&self, name: &str) -> Result<Vec<Submission>> {
        Ok(self.conn()?
            .prepare("
                SELECT author_id, timestamp, image, challenge
                FROM Submission s
                JOIN User u ON s.author_id = u.id
                WHERE u.name = ?1",
            )?
            .query_map(params![name], |row| {
                Ok(Submission {
                    author_id: row.get(0)?,
                    timestamp: row.get(1)?,
                    image: row.get(2)?,
                    challenge: row.get(3)?,
                })
            })?
            .collect::<std::result::Result<Vec<Submission>, rusqlite::Error>>()?)
    }

    fn get_users(&self, users: HashSet<u64>) -> Result<HashMap<String, User>> {
        let conn = self.conn()?;
        // Not sure why derive_more::From is unable to convert these errors
        users
            .iter()
            // u64 must be converted to String for templates
            .map(|id| -> Result<(String, User)> {
                match conn.prepare("SELECT name, discriminator, avatar, deleted FROM User WHERE id = ?1") {
                    Ok(mut statement) => statement.query_row(params![id], |row| {
                        Ok((id.to_string(), User {
                            id: *id,
                            name: row.get(0)?,
                            discriminator: row.get(1)?,
                            avatar: row.get(2)?,
                            deleted: row.get(3)?,
                        }))
                    }).map_err(DatabaseError::Rusqlite),
                    Err(error) => Err(DatabaseError::Rusqlite(error)),
                }
            })
            .collect()
    }

    #[allow(dead_code)]
    pub fn refresh_users(&self) -> rusqlite::Result<()> {
        // Periodically refresh all changable user data (name, discriminator, avatar)
        // Ideally this should run periodically.
        todo!()
    }

    #[allow(dead_code, unused_variables)]
    pub fn insert_submission(&self, submission: &Submission) -> rusqlite::Result<()> {
        // For new submissions only
        todo!()
    }
}