use super::sql::QueryExecutor;
use crate::graph::node::RunnableNode;
use crate::graph::upload_to_db::DBType;
use crate::io::{DataSource, RecordSerializer};
use async_trait::async_trait;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use sqlx::any::install_default_drivers;
use sqlx::AnyPool;
use tiberius::{Config, EncryptionLevel};
use tokio_util::compat::TokioAsyncWriteCompatExt;

/**
 * Pull data from a db using a db query into a csv file that can be used by another node
 */
async fn pull_from_db(
    executor: &mut impl QueryExecutor,
    node: &PullFromDBNode,
) -> anyhow::Result<()> {
    let mut output_file = csv::Writer::from_path(node.output_data_source.path.clone())?;
    let mut first_row = true;
    executor
        .get_rows(&node.query, &node.parameters, &mut move |row| {
            if first_row {
                output_file.write_header(&row)?;
                first_row = false;
            }
            output_file.write_record(row.values())?;
            Ok(())
        })
        .await?;
    Ok(())
}

#[derive(Serialize, Deserialize, Clone, JsonSchema)]
pub struct PullFromDBNode {
    output_data_source: DataSource,
    query: String,
    parameters: Vec<String>,
    db_type: DBType,
    connection_string: String,
}

pub struct PullFromDBNodeRunner {
    pub pull_from_db_node: PullFromDBNode,
}

#[async_trait]
impl RunnableNode for PullFromDBNodeRunner {
    async fn run(&self) -> anyhow::Result<()> {
        let node = self.pull_from_db_node.clone();
        // TODO: Clean up grabbing of connection/executor so don't need to repeat this between upload/download
        match node.db_type {
            DBType::Mssql => {
                let mut config = Config::from_jdbc_string(&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?;
                pull_from_db(&mut client, &node).await?;
            }
            _ => {
                install_default_drivers();
                let mut pool = AnyPool::connect(&node.connection_string).await?;
                pull_from_db(&mut pool, &node).await?;
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::graph::node::RunnableNode;
    use crate::graph::pull_from_db::{PullFromDBNode, PullFromDBNodeRunner};
    use crate::graph::upload_to_db::DBType::Mssql;
    use crate::io::DataSource;
    use crate::io::SourceType::CSV;
    use std::fs::File;
    use std::io::Read;
    use std::path::PathBuf;
    use testcontainers::core::{IntoContainerPort, WaitFor};
    use testcontainers::runners::AsyncRunner;
    use testcontainers::{GenericImage, ImageExt};

    #[tokio::test]
    async fn test_sql_server() -> anyhow::Result<()> {
        let container = GenericImage::new("mcr.microsoft.com/mssql/server", "2022-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",
            host, port
        )
        .to_owned();

        let runner = PullFromDBNodeRunner {
            pull_from_db_node: PullFromDBNode {
                db_type: Mssql,
                query: "SELECT '1' Test".to_owned(),
                parameters: vec![],
                connection_string,
                output_data_source: DataSource {
                    path: PathBuf::from("test_pull.csv"),
                    source_type: CSV,
                },
            },
        };
        runner.run().await?;
        let mut result_contents = String::new();
        let _ = File::open("test_pull.csv")?.read_to_string(&mut result_contents)?;
        assert_eq!(
            "Test
1
",
            result_contents,
            "Should pull the correct data from sql"
        );
        Ok(())
    }
}
