use anyhow::bail;
use async_trait::async_trait;
use itertools::Itertools;
use log::{log, Level};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use sqlx::any::install_default_drivers;
use sqlx::{AnyConnection, AnyPool};
use std::{collections::HashMap, fmt::format};
use tiberius::{Config, EncryptionLevel};
use tokio_util::compat::TokioAsyncWriteCompatExt;

#[cfg(feature = "duckdb")]
use duckdb::Connection;

use super::{node::RunnableNode, sql::QueryExecutor};

const BIND_LIMIT: usize = 65535;

pub async fn upload_file_bulk(
    executor: &mut impl QueryExecutor,
    upload_node: &UploadNode,
    bind_limit: usize,
    db_type: &DBType,
) -> anyhow::Result<u64> {
    let mut rows_affected = None;
    if upload_node.column_mappings.is_none() {
        let insert_from_file_query = match upload_node.db_type {
            DBType::Postgres => Some(format!(
                r#"COPY "{}" FROM '{}' WITH DELIMITER ',' CSV HEADER QUOTE '"';"#,
                upload_node.table_name.replace("\"", ""),
                upload_node.file_path.replace("'", "''")
            )),
            // This can't be parameterised as mysql doesn't allow LOAD DATA
            // to be executed within a prepare. Therefore even dynamic sql
            // can't work (like it can with mssql) since mysql won't allow it
            DBType::Mysql => Some(format!(
                "LOAD DATA INFILE '{}' INTO TABLE `{}`
            FIELDS TERMINATED BY ',' ENCLOSED BY '\"'
            LINES TERMINATED BY '\n'
            IGNORE 1 LINES;",
                upload_node.file_path.replace("'", "''"),
                upload_node.table_name.replace("'", "''"),
            )),
            DBType::Mssql => Some(format!(
                "BULK INSERT [{}] FROM '{}' WITH ( FORMAT = 'CSV', FIRSTROW = 2 );",
                upload_node.table_name.replace("]", ""),
                upload_node.file_path.replace("'", "''")
            )),
            #[cfg(feature = "duckdb")]
            DBType::Duckdb => Some(format!(
                r#"COPY "{}" FROM '{}' WITH DELIMITER ',' CSV HEADER QUOTE '"';"#,
                upload_node.table_name.replace("\"", ""),
                upload_node.file_path.replace("'", "''")
            )),
        };
        if let Some(insert_from_file_query) = insert_from_file_query {
            let result = executor.execute_unchecked(&insert_from_file_query).await;
            if let Ok(result) = result {
                rows_affected = Some(result);
            } else {
                log!(
                    Level::Debug,
                    "Failed to bulk insert, trying standard sql insert instead"
                );
            }
        }
    }

    if rows_affected == None {
        log!(
            Level::Warn,
            "Failed to copy csv to database, falling back to manual insert, which will be slow"
        );
        let mut file_reader = csv::Reader::from_path(upload_node.file_path.clone())?;
        let csv_columns = file_reader
            .headers()?
            .iter()
            .map(|header| header.to_owned())
            .collect_vec();
        let table_columns = if let Some(column_mappings) = &upload_node.column_mappings {
            csv_columns
                .iter()
                .map(|column| column_mappings.get(column).unwrap_or(column).clone())
                .collect_vec()
        } else {
            csv_columns.clone()
        };
        let query_template = format!(
            "INSERT INTO {} ({}) \n",
            db_type.quote_name(&upload_node.table_name),
            table_columns
                .iter()
                .map(|column| db_type.quote_name(column))
                .join(",")
        );
        let mut params = vec![];
        let mut insert_query = "".to_owned();
        let mut num_params = 0;
        let mut running_row_total = 0;

        for result in file_reader.records() {
            let result = result?;
            insert_query = insert_query
                + format!(
                    "VALUES ({})",
                    result
                        .iter()
                        // TODO: This is done because postgres won't do an implicit cast on text in the insert,
                        //         if this is resolved we can go back to parameterised queries instead
                        //         as mysql/sql server work fine with implicit cast on parameters
                        //         Alternatively, we could inspect the table metadata, though a
                        //         good DBA should prevent the db user from accessing this information
                        //         or performing bad queries like deletes if a user somehow breaks
                        //         quote escapes - we'll need to doc this limitation
                        .map(|value| format!("'{}'", value.replace("'", "''")))
                        .join(",")
                )
                .as_str();
            num_params += csv_columns.len();
            if num_params == bind_limit {
                running_row_total += executor.execute_query(&query_template, &params).await?;
                insert_query = "".to_owned();
                params = vec![];
                num_params = 0;
            }
        }
        if !insert_query.is_empty() {
            running_row_total += executor
                .execute_query(format!("{}{}", query_template, insert_query), &params)
                .await?;
        }
        rows_affected = Some(running_row_total);
    }

    if rows_affected.is_some() {
        if let Some(post_script) = &upload_node.post_script {
            executor.execute_query(post_script, &[]).await?;
        }
    }
    match rows_affected {
        Some(rows_affected) => Ok(rows_affected),
        None => bail!(
            "Invalid state, rows_affected must be populated, or an error should have occurred"
        ),
    }
}

#[derive(Serialize, Deserialize, Clone, JsonSchema, PartialEq)]
pub enum DBType {
    Mysql,
    Postgres,
    Mssql,
    #[cfg(feature = "duckdb")]
    Duckdb,
}

impl DBType {
    pub fn quote_name(&self, quoted: impl AsRef<str>) -> String {
        let quoted = quoted.as_ref();
        match self {
            DBType::Postgres => format!(r#""{}""#, quoted.replace("\"", "")),
            DBType::Mysql => format!("`{}`", quoted.replace("`", "")),
            DBType::Mssql => format!("[{}]", quoted.replace("]", "")),
            #[cfg(feature = "duckdb")]
            DBType::Duckdb => format!(r#""{}""#, quoted.replace("\"", "")),
        }
    }
}

#[derive(Serialize, Deserialize, Clone, JsonSchema)]
pub struct UploadNode {
    file_path: String,
    table_name: String,
    // Mappings from column in file -> column in db
    column_mappings: Option<HashMap<String, String>>,
    post_script: Option<String>,
    db_type: DBType,
    connection_string: String,
}

pub struct UploadNodeRunner {
    pub upload_node: UploadNode,
}

#[async_trait]
impl RunnableNode for UploadNodeRunner {
    async fn run(&self) -> anyhow::Result<()> {
        let upload_node = self.upload_node.clone();

        #[cfg(feature = "duckdb")]
        {
            if upload_node.db_type == DBType::Duckdb {
                let mut conn = Connection::open(&upload_node.connection_string)?;
                upload_file_bulk(&mut conn, &upload_node, BIND_LIMIT, &upload_node.db_type).await?;
                Ok(())
            }
        }
        if upload_node.db_type == DBType::Mssql {
            let mut config = Config::from_jdbc_string(&upload_node.connection_string)?;
            // TODO: Restore encryption for remote hosts, doesn't work on localhost without encryption.
            config.encryption(EncryptionLevel::NotSupported);
            let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
            tcp.set_nodelay(true)?;
            let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
            upload_file_bulk(&mut client, &upload_node, BIND_LIMIT, &upload_node.db_type).await?;
        } else {
            install_default_drivers();
            let mut pool = AnyPool::connect(&upload_node.connection_string).await?;
            upload_file_bulk(&mut pool, &upload_node, BIND_LIMIT, &upload_node.db_type).await?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::graph::node::RunnableNode;
    use crate::graph::upload_to_db::{DBType, UploadNode, UploadNodeRunner};
    #[cfg(feature = "duckdb")]
    use duckdb::{params, Connection};
    use sqlx::{AnyPool, Row};
    use std::path::PathBuf;
    use std::{fs, path};
    use testcontainers::core::{IntoContainerPort, WaitFor};
    use testcontainers::runners::AsyncRunner;
    use testcontainers::{GenericImage, ImageExt};
    use tiberius::{Config, EncryptionLevel};
    use tokio_util::compat::TokioAsyncWriteCompatExt;

    #[tokio::test]
    #[ignore]
    pub async fn check_bulk_upload_mssql() -> anyhow::Result<()> {
        let container = GenericImage::new(
            "gitea.michaelpivato.dev/vato007/ingey-test-db-mssql",
            "latest",
        )
        .with_exposed_port(1433.tcp())
        .with_wait_for(WaitFor::message_on_stdout(
            "Recovery is complete.".to_owned(),
        ))
        .with_env_var("ACCEPT_EULA", "Y")
        .with_env_var("MSSQL_SA_PASSWORD", "TestOnlyContainer123")
        .with_copy_to(
            "/upload_to_db/test.csv",
            PathBuf::from("./testing/input/upload_to_db/test.csv"),
        )
        .start()
        .await?;
        let host = container.get_host().await?;
        let port = container.get_host_port_ipv4(1433).await?;
        let connection_string = format!("jdbc:sqlserver://{}:{};username=sa;password=TestOnlyContainer123;database=TestIngeyDatabase", host, port).to_owned();
        let file = "/upload_to_db/test.csv".to_owned();
        let table_name = "My Test Table".to_string();
        let upload_node = UploadNodeRunner {
            upload_node: UploadNode {
                file_path: file.to_owned(),
                table_name: table_name.clone(),
                column_mappings: None,
                post_script: None,
                db_type: DBType::Mssql,
                connection_string,
            },
        };
        upload_node.run().await?;
        let mut config = Config::from_jdbc_string(&upload_node.upload_node.connection_string)?;
        config.encryption(EncryptionLevel::NotSupported);
        let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
        tcp.set_nodelay(true)?;
        let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
        let result = client.query("SELECT * FROM [My Test Table]", &[]).await?;
        let result = result.into_first_result().await?;
        assert_eq!(1, result.len());

        let first_row = result.get(0).unwrap();
        assert_eq!(Some(1), first_row.get("column1"));
        assert_eq!(Some("Hello"), first_row.get("column2"));

        Ok(())
    }

    #[tokio::test]
    pub async fn check_bulk_upload_postgres() -> anyhow::Result<()> {
        let container = GenericImage::new(
            "gitea.michaelpivato.dev/vato007/ingey-test-db-postgres",
            "latest",
        )
        .with_exposed_port(5432.tcp())
        .with_wait_for(WaitFor::message_on_stderr(
            "database system is ready to accept connections",
        ))
        .with_env_var("POSTGRES_PASSWORD", "TestOnlyContainer123")
        .with_copy_to(
            "/upload_to_db/test.csv",
            PathBuf::from("./testing/input/upload_to_db/test.csv"),
        )
        .start()
        .await?;
        let host = container.get_host().await?;
        let port = container.get_host_port_ipv4(5432).await?;
        let connection_string = format!(
            "postgres://postgres:TestOnlyContainer123@{}:{}/testingeydatabase",
            host, port
        )
        .to_owned();
        let file = "/upload_to_db/test.csv".to_owned();
        let table_name = "My Test Table".to_string();
        let upload_node = UploadNodeRunner {
            upload_node: UploadNode {
                file_path: file.to_owned(),
                table_name: table_name.clone(),
                column_mappings: None,
                post_script: None,
                db_type: DBType::Postgres,
                connection_string,
            },
        };
        upload_node.run().await?;
        let pool = AnyPool::connect(&upload_node.upload_node.connection_string).await?;
        let result = sqlx::query(r#"SELECT * FROM "My Test Table""#)
            .fetch_one(&pool)
            .await?;

        let column1: i32 = result.try_get("column1")?;
        let column2: &str = result.try_get("column2")?;
        assert_eq!(1, column1);
        assert_eq!("Hello", column2);
        Ok(())
    }

    #[tokio::test]
    #[cfg(feature = "duckdb")]
    pub async fn check_bulk_upload_duckdb() -> anyhow::Result<()> {
        let connection_string = "./testing/output/test.duckdb".to_owned();
        fs::remove_file(&connection_string);
        let file = path::absolute("./testing/input/upload_to_db/test.csv")?
            .to_string_lossy()
            .into_owned();
        let table_name = "My Test Table".to_string();
        let upload_node = UploadNodeRunner {
            upload_node: UploadNode {
                file_path: file.to_owned(),
                table_name: table_name.clone(),
                column_mappings: None,
                post_script: None,
                db_type: DBType::Duckdb,
                connection_string,
            },
        };
        let mut conn = Connection::open(&upload_node.upload_node.connection_string)?;
        conn.execute(
            r#"CREATE TABLE "My Test Table" (
        column1 INT,
        column2 VARCHAR(200),
        );"#,
            params![],
        )?;
        upload_node.run().await?;
        // Duckdb quirk
        let mut conn = Connection::open(&upload_node.upload_node.connection_string)?;
        let mut stmt = conn.prepare(r#"SELECT column1, column2 FROM "My Test Table";"#)?;
        let mut query = stmt.query(params![])?;
        let row = query.next()?;
        assert_eq!(true, row.is_some());
        let row = row.unwrap();
        let column1: i32 = row.get(0)?;
        let column2: String = row.get(1)?;
        assert_eq!(1, column1);
        assert_eq!("Hello", column2);
        Ok(())
    }

    #[tokio::test]
    pub async fn check_bulk_upload_mysql() -> anyhow::Result<()> {
        let container = GenericImage::new(
            "gitea.michaelpivato.dev/vato007/ingey-test-db-mysql",
            "latest",
        )
        .with_exposed_port(3306.tcp())
        .with_wait_for(WaitFor::message_on_stderr("ready for connections."))
        .with_env_var("MYSQL_ROOT_PASSWORD", "TestOnlyContainer123")
        .with_copy_to(
            "/upload_to_db/test.csv",
            PathBuf::from("./testing/input/upload_to_db/test.csv"),
        )
        // https://dev.mysql.com/doc/refman/8.4/en/server-system-variables.html#sysvar_secure_file_priv
        .with_cmd(&["--secure-file-priv=/upload_to_db".to_owned()])
        .start()
        .await?;
        let host = container.get_host().await?;
        let port = container.get_host_port_ipv4(3306).await?;
        let connection_string = format!(
            "mysql://root:TestOnlyContainer123@{}:{}/TestIngeyDatabase",
            host, port
        )
        .to_owned();
        let file = "/upload_to_db/test.csv".to_owned();
        let table_name = "My Test Table".to_string();
        let upload_node = UploadNodeRunner {
            upload_node: UploadNode {
                file_path: file.to_owned(),
                table_name: table_name.clone(),
                column_mappings: None,
                post_script: None,
                db_type: DBType::Mysql,
                connection_string,
            },
        };
        upload_node.run().await?;
        let pool = AnyPool::connect(&upload_node.upload_node.connection_string).await?;
        let result = sqlx::query(r#"SELECT * FROM `My Test Table`"#)
            .fetch_one(&pool)
            .await?;

        let column1: i32 = result.try_get("column1")?;
        let column2: &str = result.try_get("column2")?;
        assert_eq!(1, column1);
        assert_eq!("Hello", column2);
        Ok(())
    }

    #[tokio::test]
    #[ignore]
    pub async fn check_insert_mssql() -> anyhow::Result<()> {
        let container = GenericImage::new(
            "gitea.michaelpivato.dev/vato007/ingey-test-db-mssql",
            "latest",
        )
        .with_exposed_port(1433.tcp())
        .with_wait_for(WaitFor::message_on_stdout(
            "Recovery is complete.".to_owned(),
        ))
        .with_env_var("ACCEPT_EULA", "Y")
        .with_env_var("MSSQL_SA_PASSWORD", "TestOnlyContainer123")
        .start()
        .await?;
        let host = container.get_host().await?;
        let port = container.get_host_port_ipv4(1433).await?;
        let connection_string = format!("jdbc:sqlserver://{}:{};username=sa;password=TestOnlyContainer123;database=TestIngeyDatabase", host, port).to_owned();
        let file = "./testing/input/upload_to_db/test.csv".to_owned();
        let table_name = "My Test Table".to_string();
        let upload_node = UploadNodeRunner {
            upload_node: UploadNode {
                file_path: file.to_owned(),
                table_name: table_name.clone(),
                column_mappings: None,
                post_script: None,
                db_type: DBType::Mssql,
                connection_string,
            },
        };
        upload_node.run().await?;
        let mut config = Config::from_jdbc_string(&upload_node.upload_node.connection_string)?;
        config.encryption(EncryptionLevel::NotSupported);
        let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
        tcp.set_nodelay(true)?;
        let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
        let result = client.query("SELECT * FROM [My Test Table]", &[]).await?;
        let result = result.into_first_result().await?;
        assert_eq!(1, result.len());

        let first_row = result.get(0).unwrap();
        assert_eq!(Some(1), first_row.get("column1"));
        assert_eq!(Some("Hello"), first_row.get("column2"));

        Ok(())
    }

    #[tokio::test]
    pub async fn check_insert_postgres() -> anyhow::Result<()> {
        let container = GenericImage::new(
            "gitea.michaelpivato.dev/vato007/ingey-test-db-postgres",
            "latest",
        )
        .with_exposed_port(5432.tcp())
        .with_wait_for(WaitFor::message_on_stderr(
            "database system is ready to accept connections",
        ))
        .with_env_var("POSTGRES_PASSWORD", "TestOnlyContainer123")
        .start()
        .await?;
        let host = container.get_host().await?;
        let port = container.get_host_port_ipv4(5432).await?;
        let connection_string = format!(
            "postgres://postgres:TestOnlyContainer123@{}:{}/testingeydatabase",
            host, port
        )
        .to_owned();
        let file = "./testing/input/upload_to_db/test.csv".to_owned();
        let table_name = "My Test Table".to_string();
        let upload_node = UploadNodeRunner {
            upload_node: UploadNode {
                file_path: file.to_owned(),
                table_name: table_name.clone(),
                column_mappings: None,
                post_script: None,
                db_type: DBType::Postgres,
                connection_string,
            },
        };
        upload_node.run().await?;
        let pool = AnyPool::connect(&upload_node.upload_node.connection_string).await?;
        let result = sqlx::query(r#"SELECT * FROM "My Test Table""#)
            .fetch_one(&pool)
            .await?;

        let column1: i32 = result.try_get("column1")?;
        let column2: &str = result.try_get("column2")?;
        assert_eq!(1, column1);
        assert_eq!("Hello", column2);
        Ok(())
    }

    #[tokio::test]
    pub async fn check_insert_mysql() -> anyhow::Result<()> {
        let container = GenericImage::new(
            "gitea.michaelpivato.dev/vato007/ingey-test-db-mysql",
            "latest",
        )
        .with_exposed_port(3306.tcp())
        .with_wait_for(WaitFor::message_on_stderr("ready for connections."))
        .with_env_var("MYSQL_ROOT_PASSWORD", "TestOnlyContainer123")
        .start()
        .await?;
        let host = container.get_host().await?;
        let port = container.get_host_port_ipv4(3306).await?;
        let connection_string = format!(
            "mysql://root:TestOnlyContainer123@{}:{}/TestIngeyDatabase",
            host, port
        )
        .to_owned();
        let file = "./testing/input/upload_to_db/test.csv".to_owned();
        let table_name = "My Test Table".to_string();
        let upload_node = UploadNodeRunner {
            upload_node: UploadNode {
                file_path: file.to_owned(),
                table_name: table_name.clone(),
                column_mappings: None,
                post_script: None,
                db_type: DBType::Mysql,
                connection_string,
            },
        };
        upload_node.run().await?;
        let pool = AnyPool::connect(&upload_node.upload_node.connection_string).await?;
        let result = sqlx::query(r#"SELECT * FROM `My Test Table`"#)
            .fetch_one(&pool)
            .await?;

        let column1: i32 = result.try_get("column1")?;
        let column2: &str = result.try_get("column2")?;
        assert_eq!(1, column1);
        assert_eq!("Hello", column2);
        Ok(())
    }
}
