integresql 0.2.0

use std::error::Error;
use std::hash::{Hash, Hasher};
use std::time::Duration;
use std::vec;

use http::Uri;
use integresql::*;
use log::debug;
use postgres::Client as SyncClient;
use serial_test::serial;
use tokio::time::sleep;
use tokio::{self, task::AbortHandle};
use tokio_postgres::{Client, NoTls};
use xxhash_rust::xxh64::Xxh64;

const TEST_BASE_URL: &str = "http://integresql:5000/api";

const SLEEP_DURATION: Duration = Duration::from_millis(500);
const DEV_DB: &str = "development";

#[tokio::test]
#[test_log::test]
#[serial]
async fn test_async_setup() {
    let admin_client = get_admin_client(DEV_DB).await;
    reset_db(&admin_client).await;
    let dbs_before = list_databases(&admin_client).await;
    assert_eq!(0, dbs_before.len(), "Expected no databases before setup");

    let manager = make_manager();
    let result = manager.get_template_db_async("KEY_1", async_setup).await;
    assert!(result.is_ok(), "Failed to set up template");
    sleep(SLEEP_DURATION).await; // Wait for the test DB setup to complete

    let expected_hash = calculate_template_hash("KEY_1");
    validate_database_setup(&admin_client, &expected_hash).await;
    validate_get_test_db(result.unwrap(), &expected_hash);
}

#[tokio::test]
#[test_log::test]
#[serial]
async fn test_async_setup_with_overrides() {
    let admin_client = get_admin_client(DEV_DB).await;
    reset_db(&admin_client).await;

    let manager = make_manager().with_overrides(Some("other_db_host".to_string()), None);
    let template = manager
        .get_template_db_async("KEY_OVERRIDES", async_setup)
        .await
        .expect("Failed to set up template");

    sleep(SLEEP_DURATION).await;

    let readonly_db = template
        .get_readonly_test_db()
        .expect("Failed to get readonly test DB");
    assert_eq!(
        readonly_db.host, "other_db_host",
        "Expected readonly DB host to use override"
    );

    let (client, handle) = connect_async(&readonly_db)
        .await
        .expect("Failed to connect using overridden host");
    client
        .query("SELECT 1", &[])
        .await
        .expect("Failed to query readonly test DB");
    drop(client);
    handle.abort();
    tokio::task::yield_now().await;
}

#[tokio::test]
#[test_log::test]
#[serial]
async fn test_async_full_workflow() {
    let admin_client = get_admin_client(DEV_DB).await;
    reset_db(&admin_client).await;
    let dbs_before = list_databases(&admin_client).await;
    assert_eq!(0, dbs_before.len(), "Expected no databases before setup");

    let manager = make_manager();
    let result = manager.get_template_db_async("KEY_2", async_setup).await;
    assert!(result.is_ok(), "Failed to set up template");
    sleep(SLEEP_DURATION).await; // Wait for the test DB setup to complete

    let expected_hash = calculate_template_hash("KEY_2");
    validate_get_test_db(result.unwrap(), &expected_hash);
}

#[tokio::test]
#[test_log::test]
#[serial]
async fn test_sync_setup() {
    let admin_client = get_admin_client(DEV_DB).await;
    reset_db(&admin_client).await;
    let dbs_before = list_databases(&admin_client).await;
    assert_eq!(0, dbs_before.len(), "Expected no databases before setup");

    let manager = make_manager();
    let result = manager.get_template_db_sync("KEY_3", sync_setup);
    assert!(result.is_ok(), "Failed to set up template");
    sleep(SLEEP_DURATION).await; // Wait for the test DB setup to complete

    let expected_hash = calculate_template_hash("KEY_3");
    validate_database_setup(&admin_client, &expected_hash).await;
}

#[tokio::test]
#[test_log::test]
#[serial]
async fn test_reset_tracking() {
    let admin_client = get_admin_client(DEV_DB).await;
    reset_db(&admin_client).await;
    let dbs_before = list_databases(&admin_client).await;
    assert_eq!(0, dbs_before.len(), "Expected no databases before setup");

    let manager = make_manager();
    let keys = vec!["KEY_1", "KEY_2", "KEY_3"];
    for key in &keys {
        let result = manager.get_template_db_sync(*key, sync_setup);
        assert!(result.is_ok(), "Failed to set up template for {}", key);
        sleep(SLEEP_DURATION).await; // Wait for the test DB setup to complete
        let template_hash = calculate_template_hash(*key);
        let expected_prefix = format!("integresql_test_{}_", template_hash);
        let dbs_after = list_databases(&admin_client).await;
        let test_db_count = dbs_after
            .iter()
            .filter(|db| db.starts_with(&expected_prefix))
            .count();
        assert!(
            test_db_count >= 4,
            "Expected at least 4 test databases for {}",
            key
        );
    }

    manager
        .clear_db_tracking()
        .expect("Failed to delete all templates");
    sleep(Duration::from_secs(4)).await; // Wait for the test DB setup to complete
    let dbs_after = list_databases(&admin_client).await;
    let test_db_count = dbs_after
        .iter()
        .filter(|db| db.starts_with("integresql_test_"))
        .count();
    assert_eq!(0, test_db_count, "Expected no test databases after reset");

    let template_db_count = dbs_after
        .iter()
        .filter(|db| db.starts_with("integresql_template_"))
        .count();
    assert_eq!(
        3, template_db_count,
        "Expected 3 template databases after reset"
    );
}

fn validate_get_test_db(template: TemplateDb, template_hash: &str) {
    let writable_db = template.get_writable_test_db().unwrap();
    let readonly_db = template.get_readonly_test_db().unwrap();
    let both = vec![writable_db, readonly_db];
    let (host, user, password) = get_env_settings();
    for db in both {
        assert_eq!(db.port, 5432);
        assert_eq!(db.template_hash(), template_hash);
        assert_eq!(db.host, host, "Test database host does not match expected");
        assert_eq!(
            db.username, user,
            "Test database user does not match expected"
        );
        assert_eq!(
            db.password, password,
            "Test database password does not match expected"
        );
        let expected_prefix = format!("integresql_test_{}_", template_hash);
        assert!(
            db.database.starts_with(&expected_prefix),
            "Test database name does not start with expected prefix"
        );
    }
}

async fn validate_database_setup(admin_client: &Client, template_hash: &str) {
    let dbs_after = list_databases(admin_client).await;
    let template_dbs = dbs_after
        .iter()
        .filter(|db| db.starts_with("integresql_template_"))
        .collect::<Vec<_>>();
    let test_dbs = dbs_after
        .iter()
        .filter(|db| db.starts_with("integresql_test_"))
        .collect::<Vec<_>>();
    assert_eq!(
        template_dbs.len(),
        1,
        "Expected exactly one template database, found: {:?}",
        template_dbs
    );
    assert_eq!(
        test_dbs.len(),
        dbs_after.len() - 1,
        "Expected {} test databases, found: {:?}",
        dbs_after.len() - 1,
        test_dbs
    );

    let template_db = format!("integresql_template_{}", template_hash);
    assert_eq!(
        template_dbs[0], &template_db,
        "Expected template database name to be {}, found: {}",
        template_db, template_dbs[0]
    );

    let test_db_prefix = format!("integresql_test_{}_", template_hash);
    let all_have_expected_prefix = test_dbs.iter().all(|db| db.starts_with(&test_db_prefix));
    assert!(
        all_have_expected_prefix,
        "Not all test databases have the expected prefix: {}",
        test_db_prefix
    );

    for db in test_dbs {
        let client = get_admin_client(db).await;
        let values = client
            .query("SELECT name FROM test_schema.test_table", &[])
            .await
            .expect("Failed to query test table")
            .into_iter()
            .map(|row| row.get::<_, String>(0))
            .collect::<Vec<_>>();
        assert_eq!(
            values.len(),
            3,
            "Expected 3 rows in test_table, found: {}",
            values.len()
        );
        let expected_values = vec!["value one", "value two", "value three"];
        assert_eq!(
            values, expected_values,
            "Test table values do not match expected: {:?}",
            values
        );
    }
}

async fn list_databases(client: &Client) -> Vec<String> {
    let rows = client
        .query(
            "SELECT datname FROM pg_database WHERE datname LIKE 'integresql_%'",
            &[],
        )
        .await
        .expect("Failed to query databases");
    rows.into_iter()
        .map(|row| row.get::<_, String>(0))
        .collect()
}

async fn reset_db(client: &Client) {
    let agent = get_http_client();
    agent
        .delete("http://integresql:5000/api/v1/admin/templates")
        .call()
        .expect("Failed to send reset request to integresql");
    let remaining_dbs = client
        .query(
            "SELECT datname FROM pg_database WHERE datname LIKE 'integresql_%'",
            &[],
        )
        .await
        .expect("Failed to query remaining databases")
        .into_iter()
        .map(|row| row.get::<_, String>(0))
        .collect::<Vec<_>>();
    for db in remaining_dbs {
        client
            .execute(&format!("DROP DATABASE IF EXISTS {}", db), &[])
            .await
            .expect("Failed to drop database");
    }
}

const TIMEOUT: Duration = Duration::from_secs(2);

fn get_http_client() -> ureq::Agent {
    ureq::Agent::config_builder()
        .timeout_global(Some(TIMEOUT))
        .build()
        .new_agent()
}

fn get_env_settings() -> (String, String, String) {
    let host = std::env::var("PGHOST").expect("PGHOST must be set");
    let user = std::env::var("PGUSER").expect("PGUSER must be set");
    let password = std::env::var("PGPASSWORD").expect("PGPASSWORD must be set");
    (host, user, password)
}

async fn get_admin_client(db_name: &str) -> Client {
    let (host, user, password) = get_env_settings();
    let conn_str = format!(
        "host={} user={} password={} dbname={}",
        host, user, password, db_name
    );
    let (client, connection) = tokio_postgres::connect(&conn_str, NoTls)
        .await
        .expect("Failed to connect to the database");
    debug!("Opening postgres connection to {}", db_name);
    tokio::spawn(async move {
        if let Err(e) = connection.await {
            eprintln!("Connection error: {}", e);
        }
        debug!("Closing postgres connection");
    });
    client
}

fn sync_setup(config: ConnectionSettings) -> InitializeResult {
    let thread = std::thread::spawn(move || {
        let mut client = connect_sync(&config);
        client
            .execute("CREATE SCHEMA IF NOT EXISTS test_schema", &[])
            .expect("Failed to create schema");
        client
            .execute(
                "CREATE TABLE IF NOT EXISTS test_schema.test_table (id SERIAL PRIMARY KEY, name TEXT)",
                &[],
            )
            .expect("Failed to create table");
        client.execute("INSERT INTO test_schema.test_table (name) VALUES ('value one'), ('value two'), ('value three');", &[])
            .expect("Failed to insert values");
    });
    match thread.join() {
        Ok(_) => Ok(()),
        Err(_) => Err("panic in sync_setup thread".into()),
    }
}

async fn async_setup(config: ConnectionSettings) -> InitializeResult {
    let (client, abort_handle) = connect_async(&config).await?;
    client
        .execute("CREATE SCHEMA IF NOT EXISTS test_schema", &[])
        .await?;
    client
        .execute(
            "CREATE TABLE IF NOT EXISTS test_schema.test_table (id SERIAL PRIMARY KEY, name TEXT)",
            &[],
        )
        .await?;
    client.execute("INSERT INTO test_schema.test_table (name) VALUES ('value one'), ('value two'), ('value three');", &[]).await?;
    abort_handle.abort();
    tokio::task::yield_now().await; // Ensure the connection is closed before returning
    Ok(())
}

fn connect_sync(config: &ConnectionSettings) -> SyncClient {
    let conn_str = config.to_libpq_url();
    SyncClient::connect(&conn_str, NoTls).expect("Failed to connect to the database")
}

async fn connect_async(
    config: &ConnectionSettings,
) -> Result<(Client, AbortHandle), tokio_postgres::Error> {
    let conn_str = config.to_libpq_url();
    let (client, connection) = tokio_postgres::connect(&conn_str, NoTls).await?;
    let handle = tokio::spawn(async move {
        debug!("Connecting to database: {}", conn_str);
        if let Err(e) = connection.await {
            eprintln!("Connection error: {}", e);
        }
        debug!("Closing postgres connection");
    });
    Ok((client, handle.abort_handle()))
}

fn make_manager() -> DbManager {
    let url = Uri::from_static(TEST_BASE_URL);
    let timeout = Duration::from_secs(10);
    DbManager::new(url, timeout)
}

fn calculate_template_hash(key: &str) -> String {
    let mut hasher = Xxh64::new(0);
    key.hash(&mut hasher);
    let hash = hasher.finish();
    format!("{:016x}", hash)
}