#[cfg(feature = "duckdb")]
use duckdb::{params, params_from_iter, Connection};
use futures::TryStreamExt;
use futures_io::{AsyncRead, AsyncWrite};
use sqlx::{Any, Column, Executor, Pool, Row};
use std::borrow::Borrow;
use std::collections::BTreeMap;
use tiberius::{Client, Query};

pub trait QueryExecutor {
    /// Retrieve data from a database
    async fn get_rows(
        &mut self,
        query: impl AsRef<str>,
        params: &[String],
        // TODO: This is looking pretty ugly, simpler way to handle it? Maybe with an iterator?
        row_consumer: &mut impl FnMut(BTreeMap<String, String>) -> anyhow::Result<()>,
    ) -> anyhow::Result<()>;

    /// Run a query that returns no results (e.g. bulk insert, insert)
    async fn execute_query(
        &mut self,
        query: impl AsRef<str>,
        params: &[String],
    ) -> anyhow::Result<u64>;

    /// Execute an unprepared query. Avoid where possible as sql injection is possible if not used carefully
    async fn execute_unchecked(&mut self, query: impl AsRef<str>) -> anyhow::Result<u64>;
}

impl<S: AsyncRead + AsyncWrite + Unpin + Send> QueryExecutor for Client<S> {
    async fn get_rows(
        &mut self,
        query: impl AsRef<str>,
        params: &[String],
        row_consumer: &mut impl FnMut(BTreeMap<String, String>) -> anyhow::Result<()>,
    ) -> anyhow::Result<()> {
        let mut query = Query::new(query.as_ref());
        for param in params {
            query.bind(param);
        }
        let query_result = query.query(self).await?;
        let mut query_stream = query_result.into_row_stream();

        while let Some(row) = query_stream.try_next().await? {
            let mut returned_row = BTreeMap::new();
            // TODO: Check how empty columns are handled by tiberius
            for column in row.columns().into_iter() {
                returned_row.insert(
                    column.name().to_owned(),
                    row.get(column.name()).unwrap_or_else(|| "").to_owned(),
                );
            }
            row_consumer(returned_row)?;
        }
        Ok(())
    }

    async fn execute_query(
        &mut self,
        query: impl AsRef<str>,
        params: &[String],
    ) -> anyhow::Result<u64> {
        let mut query = Query::new(query.as_ref());
        for param in params {
            query.bind(param);
        }
        let result = query.execute(self).await?;
        if result.rows_affected().len() == 0 {
            return Ok(0);
        }
        Ok(result.rows_affected()[0])
    }

    async fn execute_unchecked(&mut self, query: impl AsRef<str>) -> anyhow::Result<u64> {
        let result = self.execute(query.as_ref(), &[]).await?;
        Ok(result.rows_affected()[0])
    }
}

impl QueryExecutor for Pool<Any> {
    async fn get_rows(
        &mut self,
        query: impl AsRef<str>,
        params: &[String],
        row_consumer: &mut impl FnMut(BTreeMap<String, String>) -> anyhow::Result<()>,
    ) -> anyhow::Result<()> {
        let mut query = sqlx::query(query.as_ref());
        for param in params {
            query = query.bind(param);
        }
        let mut rows = query.fetch(self.borrow());

        while let Some(row) = rows.try_next().await? {
            let mut returned_row = BTreeMap::new();
            for column in row.columns().into_iter() {
                returned_row.insert(column.name().to_owned(), row.get(column.name()));
            }
            row_consumer(returned_row)?;
        }
        Ok(())
    }

    async fn execute_query(
        &mut self,
        query: impl AsRef<str>,
        params: &[String],
    ) -> anyhow::Result<u64> {
        let mut query = sqlx::query(query.as_ref());
        for param in params {
            query = query.bind(param);
        }
        let result = query.execute(self.borrow()).await?;
        Ok(result.rows_affected())
    }

    async fn execute_unchecked(&mut self, query: impl AsRef<str>) -> anyhow::Result<u64> {
        let result = self.execute(query.as_ref()).await?;
        Ok(result.rows_affected())
    }
}

#[cfg(feature = "duckdb")]
impl QueryExecutor for Connection {
    async fn get_rows(
        &mut self,
        query: impl AsRef<str>,
        params: &[String],
        row_consumer: &mut impl FnMut(BTreeMap<String, String>) -> anyhow::Result<()>,
    ) -> anyhow::Result<()> {
        let mut stmt = self.prepare(query.as_ref())?;
        let columns = stmt.column_names();
        let mut row_iter = stmt.query_map(params_from_iter(params.iter()), |row| {
            let mut row_map: BTreeMap<String, String> = BTreeMap::new();
            for i in 0..columns.len() {
                row_map.insert(columns[i].clone(), row.get(i)?);
            }
            Ok(row_map)
        })?;
        for row in row_iter {
            let row = row?;
            row_consumer(row)?;
        }
        Ok(())
    }

    async fn execute_query(
        &mut self,
        query: impl AsRef<str>,
        params: &[String],
    ) -> anyhow::Result<u64> {
        let rows_processed = self.execute(query.as_ref(), params_from_iter(params.iter()))?;
        Ok(rows_processed as u64)
    }

    async fn execute_unchecked(&mut self, query: impl AsRef<str>) -> anyhow::Result<u64> {
        let rows_processed = self.execute(query.as_ref(), params![])?;
        Ok(rows_processed as u64)
    }
}
