use std::collections::HashMap;

use anyhow::bail;
use async_trait::async_trait;
use itertools::Itertools;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use sqlx::AnyPool;
use tiberius::Config;
use tokio_util::compat::TokioAsyncWriteCompatExt;

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,
) -> 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!("COPY {} FROM $1", upload_node.table_name)),
            DBType::Mysql => Some(format!(
                "LOAD DATA INFILE ? INTO {}",
                upload_node.table_name,
            )),
            DBType::Mssql => Some(format!("BULK INSERT {} FROM ?", upload_node.table_name)),
            _ => None,
        };
        if let Some(insert_from_file_query) = insert_from_file_query {
            let result = executor
                .execute_query(
                    &insert_from_file_query,
                    &vec![upload_node.file_path.clone()],
                )
                .await?;
            rows_affected = Some(result);
        }
    }

    if rows_affected == None {
        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",
            upload_node.table_name,
            table_columns.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().map(|_| "?").join(",")).as_str();
            let mut values = result.iter().map(|value| value.to_owned()).collect_vec();
            params.append(&mut values);
            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(&query_template, &params).await?;
        }
        rows_affected = Some(running_row_total);
    }

    if let Some(post_script) = &upload_node.post_script {
        executor.execute_query(post_script, &vec![]).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,
    Sqlite,
}

#[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();
        if upload_node.db_type == DBType::Mssql {
            let config = Config::from_jdbc_string(&upload_node.connection_string)?;
            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).await?;
        } else {
            let mut pool = AnyPool::connect(&upload_node.connection_string).await?;
            upload_file_bulk(&mut pool, &upload_node, BIND_LIMIT).await?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::graph::upload_to_db::{DBType, UploadNode, UploadNodeRunner};

    #[test]
    pub fn check_basic_upload() {
        let upload_ode = UploadNodeRunner {
            upload_node: UploadNode {
                file_path: "".to_owned(),
                table_name: "".to_string(),
                column_mappings: None,
                post_script: None,
                db_type: DBType::Mysql,
                connection_string: "".to_string(),
            }
        };

    }

    #[test]
    pub fn check_batch_upload() {}
}
