#![cfg(sqlite_test_sqlcipher)]

use std::str::FromStr;

use sqlx::sqlite::SqliteQueryResult;
use sqlx::{query, Connection, SqliteConnection};
use sqlx::{sqlite::SqliteConnectOptions, ConnectOptions};
use tempfile::TempDir;

async fn new_db_url() -> anyhow::Result<(String, TempDir)> {
    let dir = TempDir::new()?;
    let filepath = dir.path().join("database.sqlite3");

    Ok((format!("sqlite://{}", filepath.display()), dir))
}

async fn fill_db(conn: &mut SqliteConnection) -> anyhow::Result<SqliteQueryResult> {
    conn.transaction(|tx| {
        Box::pin(async move {
            query(
                "
                CREATE TABLE Company(
                    Id INT PRIMARY KEY     NOT NULL,
                    Name           TEXT    NOT NULL,
                    Salary         REAL
                 );
                 ",
            )
            .execute(&mut **tx)
            .await?;

            query(
                r#"
                INSERT INTO Company(Id, Name, Salary)
                VALUES
                    (1, "aaa", 111),
                    (2, "bbb", 222)
                "#,
            )
            .execute(&mut **tx)
            .await
        })
    })
    .await
    .map_err(|e| e.into())
}

#[sqlx_macros::test]
async fn it_encrypts() -> anyhow::Result<()> {
    let (url, _dir) = new_db_url().await?;

    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("key", "the_password")
        .create_if_missing(true)
        .connect()
        .await?;

    fill_db(&mut conn).await?;

    // Create another connection without key, query should fail
    let mut conn = SqliteConnectOptions::from_str(&url)?.connect().await?;

    assert!(conn
        .transaction(|tx| {
            Box::pin(async move { query("SELECT * FROM Company;").fetch_all(&mut **tx).await })
        })
        .await
        .is_err());

    Ok(())
}

#[sqlx_macros::test]
async fn it_can_store_and_read_encrypted_data() -> anyhow::Result<()> {
    let (url, _dir) = new_db_url().await?;

    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("key", "the_password")
        .create_if_missing(true)
        .connect()
        .await?;

    fill_db(&mut conn).await?;

    // Create another connection with valid key
    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("key", "the_password")
        .connect()
        .await?;

    let result = conn
        .transaction(|tx| {
            Box::pin(async move { query("SELECT * FROM Company;").fetch_all(&mut **tx).await })
        })
        .await?;

    assert!(result.len() > 0);

    Ok(())
}

#[sqlx_macros::test]
async fn it_fails_if_password_is_incorrect() -> anyhow::Result<()> {
    let (url, _dir) = new_db_url().await?;

    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("key", "the_password")
        .create_if_missing(true)
        .connect()
        .await?;

    fill_db(&mut conn).await?;

    // Connection with invalid key should not allow to execute queries
    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("key", "BADBADBAD")
        .connect()
        .await?;

    assert!(conn
        .transaction(|tx| {
            Box::pin(async move { query("SELECT * FROM Company;").fetch_all(&mut **tx).await })
        })
        .await
        .is_err());

    Ok(())
}

#[sqlx_macros::test]
async fn it_honors_order_of_encryption_pragmas() -> anyhow::Result<()> {
    let (url, _dir) = new_db_url().await?;

    // Make call of cipher configuration mixed with other pragmas,
    // it should have no effect, encryption related pragmas should be
    // executed first and allow to establish valid connection
    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("cipher_kdf_algorithm", "PBKDF2_HMAC_SHA1")
        .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
        .pragma("cipher_page_size", "1024")
        .pragma("key", "the_password")
        .foreign_keys(true)
        .pragma("kdf_iter", "64000")
        .auto_vacuum(sqlx::sqlite::SqliteAutoVacuum::Incremental)
        .pragma("cipher_hmac_algorithm", "HMAC_SHA1")
        .create_if_missing(true)
        .connect()
        .await?;

    fill_db(&mut conn).await?;

    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("dummy", "pragma")
        // The cipher configuration set on first connection is
        // version 3 of SQLCipher, so for second it's enough to set
        // the compatibility mode.
        .pragma("cipher_compatibility", "3")
        .pragma("key", "the_password")
        .connect()
        .await?;

    let result = conn
        .transaction(|tx| {
            Box::pin(async move { query("SELECT * FROM COMPANY;").fetch_all(&mut **tx).await })
        })
        .await?;

    assert!(result.len() > 0);

    Ok(())
}

#[sqlx_macros::test]
async fn it_allows_to_rekey_the_db() -> anyhow::Result<()> {
    let (url, _dir) = new_db_url().await?;

    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("key", "the_password")
        .create_if_missing(true)
        .connect()
        .await?;

    fill_db(&mut conn).await?;

    // The 'pragma rekey' can be called at any time
    query("PRAGMA rekey = new_password;")
        .execute(&mut conn)
        .await?;

    let mut conn = SqliteConnectOptions::from_str(&url)?
        .pragma("dummy", "pragma")
        .pragma("key", "new_password")
        .connect()
        .await?;

    let result = conn
        .transaction(|tx| {
            Box::pin(async move { query("SELECT * FROM COMPANY;").fetch_all(&mut **tx).await })
        })
        .await?;

    assert!(result.len() > 0);

    Ok(())
}