<commit_before>use chrono;
use chrono::TimeZone;
use futures_cpupool::CpuPool;
use futures::{BoxFuture, Future};
use r2d2;
use r2d2_sqlite::SqliteConnectionManager;
use rusqlite;
use rand::{Rng, thread_rng};
use serde;
use linear_map::LinearMap;

use std::sync::Arc;
use std::path::Path;


#[derive(Clone, Debug, Serialize)]
pub struct Transaction {
    pub shafter: String,
    pub shaftee: String,
    pub amount: i64,
    #[serde(serialize_with = "serialize_time")]
    pub datetime: chrono::DateTime<chrono::Utc>,
    pub reason: String,
}

#[derive(Debug, Clone, Serialize)]
pub struct User {
    pub user_id: String,
    pub display_name: String,
    pub balance: i64,
}


pub trait Database: Send + Sync {
    fn get_user_by_github_id(&self, github_user_id: String)
        -> BoxFuture<Option<String>, DatabaseError>;

    fn add_user_by_github_id(&self, github_user_id: String, display_name: String)
        -> BoxFuture<String, DatabaseError>;

    fn create_token_for_user(&self, user_id: String)
        -> BoxFuture<String, DatabaseError>;

    fn delete_token(&self, token: String)
        -> BoxFuture<(), DatabaseError>;

    fn get_user_from_token(&self, token: String)
        -> BoxFuture<Option<User>, DatabaseError>;

    fn get_balance_for_user(&self, user: String)
        -> BoxFuture<i64, DatabaseError>;

    fn get_all_users(&self)
        -> BoxFuture<LinearMap<String, User>, DatabaseError>;

    fn shaft_user(&self, transaction: Transaction)
        -> BoxFuture<(), ShaftUserError>;

    fn get_last_transactions(&self, limit: u32)
        -> BoxFuture<Vec<Transaction>, DatabaseError>;
}


#[derive(Clone)]
pub struct SqliteDatabase {
    cpu_pool: CpuPool,
    db_pool: Arc<r2d2::Pool<SqliteConnectionManager>>,
}


impl SqliteDatabase {
    pub fn with_path<P: AsRef<Path>>(path: P) -> SqliteDatabase {
        let config = r2d2::Config::default();
        let manager = SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager).unwrap();

        SqliteDatabase {
            cpu_pool: CpuPool::new_num_cpus(),
            db_pool: Arc::new(pool),
        }
    }
}


impl Database for SqliteDatabase {
    fn get_user_by_github_id(&self, github_user_id: String)
        -> BoxFuture<Option<String>, DatabaseError>
    {
        let db_pool = self.db_pool.clone();

        self.cpu_pool.spawn_fn(move || -> Result<_, DatabaseError>  {
            let conn = db_pool.get()?;

            let row = conn.query_row(
                "SELECT user_id FROM github_users WHERE github_id = $1",
                &[&github_user_id],
                |row| -> String {
                    row.get(0)
                }
            )
            .map(|row| Some(row))
            .or_else(|err| {
                if let rusqlite::Error::QueryReturnedNoRows = err {
                    Ok(None)
                } else {
                    Err(err)
                }
            })?;

            Ok(row)
        }).boxed()
    }

    fn add_user_by_github_id(&self, github_user_id: String, display_name: String)
        -> BoxFuture<String, DatabaseError>
    {
        let db_pool = self.db_pool.clone();

        self.cpu_pool.spawn_fn(move || -> Result<_, DatabaseError>  {
            let conn = db_pool.get()?;

            conn.execute(
                "INSERT INTO github_users (user_id, github_id, display_name)
                VALUES ($1, $1, $2)",
                &[&github_user_id, &display_name],
            )?;

            // TODO: Handle clashing user IDs

            Ok(github_user_id)
        }).boxed()
    }

    fn create_token_for_user(&self, user_id: String)
        -> BoxFuture<String, DatabaseError>
    {<commit_msg>Insert display name in correct place<commit_after>use chrono;
use chrono::TimeZone;
use futures_cpupool::CpuPool;
use futures::{BoxFuture, Future};
use r2d2;
use r2d2_sqlite::SqliteConnectionManager;
use rusqlite;
use rand::{Rng, thread_rng};
use serde;
use linear_map::LinearMap;

use std::sync::Arc;
use std::path::Path;


#[derive(Clone, Debug, Serialize)]
pub struct Transaction {
    pub shafter: String,
    pub shaftee: String,
    pub amount: i64,
    #[serde(serialize_with = "serialize_time")]
    pub datetime: chrono::DateTime<chrono::Utc>,
    pub reason: String,
}

#[derive(Debug, Clone, Serialize)]
pub struct User {
    pub user_id: String,
    pub display_name: String,
    pub balance: i64,
}


pub trait Database: Send + Sync {
    fn get_user_by_github_id(&self, github_user_id: String)
        -> BoxFuture<Option<String>, DatabaseError>;

    fn add_user_by_github_id(&self, github_user_id: String, display_name: String)
        -> BoxFuture<String, DatabaseError>;

    fn create_token_for_user(&self, user_id: String)
        -> BoxFuture<String, DatabaseError>;

    fn delete_token(&self, token: String)
        -> BoxFuture<(), DatabaseError>;

    fn get_user_from_token(&self, token: String)
        -> BoxFuture<Option<User>, DatabaseError>;

    fn get_balance_for_user(&self, user: String)
        -> BoxFuture<i64, DatabaseError>;

    fn get_all_users(&self)
        -> BoxFuture<LinearMap<String, User>, DatabaseError>;

    fn shaft_user(&self, transaction: Transaction)
        -> BoxFuture<(), ShaftUserError>;

    fn get_last_transactions(&self, limit: u32)
        -> BoxFuture<Vec<Transaction>, DatabaseError>;
}


#[derive(Clone)]
pub struct SqliteDatabase {
    cpu_pool: CpuPool,
    db_pool: Arc<r2d2::Pool<SqliteConnectionManager>>,
}


impl SqliteDatabase {
    pub fn with_path<P: AsRef<Path>>(path: P) -> SqliteDatabase {
        let config = r2d2::Config::default();
        let manager = SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager).unwrap();

        SqliteDatabase {
            cpu_pool: CpuPool::new_num_cpus(),
            db_pool: Arc::new(pool),
        }
    }
}


impl Database for SqliteDatabase {
    fn get_user_by_github_id(&self, github_user_id: String)
        -> BoxFuture<Option<String>, DatabaseError>
    {
        let db_pool = self.db_pool.clone();

        self.cpu_pool.spawn_fn(move || -> Result<_, DatabaseError>  {
            let conn = db_pool.get()?;

            let row = conn.query_row(
                "SELECT user_id FROM github_users WHERE github_id = $1",
                &[&github_user_id],
                |row| -> String {
                    row.get(0)
                }
            )
            .map(|row| Some(row))
            .or_else(|err| {
                if let rusqlite::Error::QueryReturnedNoRows = err {
                    Ok(None)
                } else {
                    Err(err)
                }
            })?;

            Ok(row)
        }).boxed()
    }

    fn add_user_by_github_id(&self, github_user_id: String, display_name: String)
        -> BoxFuture<String, DatabaseError>
    {
        let db_pool = self.db_pool.clone();

        self.cpu_pool.spawn_fn(move || -> Result<_, DatabaseError>  {
            let conn = db_pool.get()?;

            conn.execute(
                "INSERT INTO github_users (user_id, github_id)
                VALUES ($1, $1, $2)",
                &[&github_user_id],
            )?;

             conn.execute(
                "INSERT INTO users (user_id, display_name)
                VALUES ($1, $1, $2)",
                &[&github_user_id, &display_name],
            )?;

            Ok(github_user_id)
        }).boxed()
    }

    fn create_token_for_user(&self, user_id: String)
        -> BoxFuture<String, DatabaseError>
    {
{'fraction_non_alphanumeric': 0.12388792988978888, 'fraction_numerical': 0.006506440047802417, 'mean_word_length': 2.4220808723307585, 'pattern_counts': {'":': 0, '<': 51, '<?xml version=': 0, '>': 83, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2677958', 'n_tokens_mistral': 2444, 'n_tokens_neox': 2372, 'n_words': 496}