use std::{
  fmt::Write,
  io::{self, Write as _},
  str::FromStr,
  string::String,
  sync::Arc,
};

use async_trait::async_trait;
use color_eyre::eyre::{self, Result};
use futures::stream::StreamExt;
use sqlparser::ast::Statement;
use sqlx::{
  Column, Either, Row, ValueRef,
  pool::PoolConnection,
  postgres::{PgConnectOptions, PgConnection, PgPoolOptions, PgSslMode, Postgres},
  types::Uuid,
};
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tracing::Instrument;

use super::{
  Database, DbTaskResult, Driver, Header, Headers, QueryResultsWithMetadata, QueryTask, Rows,
  Value, builtin_functions, vec_to_string,
};
use crate::completion::{TableColumns, TableRef, table_columns_from_rows};

type PostgresTransaction = sqlx::Transaction<'static, Postgres>;
type ConnectionTask = JoinHandle<Result<(Arc<Mutex<PoolConnection<Postgres>>>, i32)>>;
type TransactionAcquireTask = JoinHandle<Result<(PostgresTransaction, i32)>>;
type TransactionTask = JoinHandle<(QueryResultsWithMetadata, PostgresTransaction)>;
enum PostgresTask {
  QueryConnect {
    handle: ConnectionTask,
    first_query: String,
    statement_type: Option<Statement>,
  },
  Query(QueryTask),
  TxConnect {
    handle: TransactionAcquireTask,
    first_query: String,
    statement_type: Statement,
    display_statement_type: Box<Statement>,
  },
  TxStart(TransactionTask),
  TxPending(Box<(PostgresTransaction, QueryResultsWithMetadata)>),
}

const MENU_QUERY: &str =
  "select menu_items.table_schema, menu_items.object_name, menu_items.object_kind
      from (
        select n.nspname as table_schema,
          c.relname as object_name,
          case
            when c.relkind in ('r', 'p', 'f') then 'table'
            when c.relkind = 'v' then 'view'
            when c.relkind = 'm' then 'materialized_view'
          end as object_kind
        from pg_class c
        join pg_namespace n on n.oid = c.relnamespace
        where n.nspname != 'pg_catalog'
          and n.nspname != 'information_schema'
          and c.relkind in ('r', 'p', 'f', 'v', 'm')
        union all
        select n.nspname as table_schema,
          format('%s(%s)', p.proname, pg_get_function_identity_arguments(p.oid)) as object_name,
          'function' as object_kind
        from pg_proc p
        join pg_namespace n on n.oid = p.pronamespace
        where n.nspname != 'pg_catalog'
          and n.nspname != 'information_schema'
          and p.prokind = 'f'
      ) menu_items
      order by menu_items.table_schema, menu_items.object_kind, menu_items.object_name asc";

const SIMPLE_MENU_QUERY: &str = "select table_schema, table_name as object_name,
        case
          when table_type = 'MATERIALIZED VIEW' then 'materialized_view'
          when table_type = 'VIEW' then 'view'
          else 'table'
        end as object_kind
      from information_schema.tables
      where table_schema not in ('pg_catalog', 'information_schema')
      order by table_schema, object_kind, object_name asc";

#[derive(Default)]
pub struct PostgresDriver {
  pool: Option<Arc<sqlx::Pool<Postgres>>>,
  task: Option<PostgresTask>,
  querying_conn: Option<Arc<Mutex<PoolConnection<Postgres>>>>,
  querying_pid: Option<String>,
}

#[async_trait(?Send)]
impl Database for PostgresDriver {
  fn builtin_functions(&self) -> &'static [&'static str] {
    builtin_functions::POSTGRESQL
  }

  async fn init(&mut self, args: crate::cli::Cli) -> Result<String> {
    let opts = super::postgresql::PostgresDriver::build_connection_opts(args)?;
    let pool = Arc::new(PgPoolOptions::new().max_connections(3).connect_with(opts.clone()).await?);
    self.pool = Some(pool);
    Ok(format!("{}/{}", opts.get_port(), opts.get_database().unwrap_or("postgres")))
  }

  // since it's possible for raw_sql to execute multiple queries in a single string,
  // we only execute the first one and then drop the rest.
  async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()> {
    let (first_query, statement_type) = match bypass_parser {
      true => (query, None),
      false => {
        let (first, stmt) = super::get_first_query(query, Driver::Postgres)?;
        (first, Some(stmt))
      },
    };
    let pool = self.pool.clone().unwrap();
    self.task = Some(PostgresTask::QueryConnect {
      handle: tokio::spawn(
        async move {
          let conn = Arc::new(Mutex::new(pool.acquire().await?));
          let pid = connection_pid(conn.lock().await.as_mut()).await?;
          Ok((conn, pid))
        }
        .in_current_span(),
      ),
      first_query,
      statement_type,
    });
    Ok(())
  }

  async fn abort_query(&mut self) -> Result<bool> {
    match self.task.take() {
      Some(task) => {
        match task {
          PostgresTask::QueryConnect { handle, .. } => handle.abort(),
          PostgresTask::Query(handle) => handle.abort(),
          PostgresTask::TxConnect { handle, .. } => handle.abort(),
          PostgresTask::TxStart(handle) => handle.abort(),
          _ => {},
        };
        if let Some(pid) = self.querying_pid.take() {
          let result = sqlx::raw_sql(&format!("SELECT pg_cancel_backend({pid})"))
            .fetch_one(&*self.pool.clone().unwrap())
            .await;
          let msg = match &result {
            Ok(_) => "Successfully killed".to_string(),
            Err(e) => format!("Failed to kill: {e:?}"),
          };

          let success = result.as_ref().is_ok_and(|r| r.try_get::<bool, _>(0).unwrap_or(false));
          let status_string = result.map_or("ERROR".to_string(), |r| {
            r.try_get_unchecked::<String, _>(0).unwrap_or("ERROR".to_string())
          });

          if !success {
            log::warn!("Unexpected response when cancelling backend process with PID {pid}: {msg}");
            log::warn!("Status: {status_string}");
          } else {
            log::info!("Tried to cancel backend process with PID {pid}: {msg}");
            log::info!("Status: {status_string}");
          }
        }
        self.querying_conn = None;
        Ok(true)
      },
      _ => {
        self.querying_conn = None;
        self.querying_pid = None;
        Ok(false)
      },
    }
  }

  async fn get_query_results(&mut self) -> Result<DbTaskResult> {
    let (task_result, next_task) = match self.task.take() {
      None => (DbTaskResult::NoTask, None),
      Some(PostgresTask::QueryConnect { handle, first_query, statement_type }) => {
        if !handle.is_finished() {
          (
            DbTaskResult::Pending,
            Some(PostgresTask::QueryConnect { handle, first_query, statement_type }),
          )
        } else {
          match handle.await? {
            Ok((conn, pid)) => {
              log::info!("Starting query with PID {}", pid.clone());
              self.querying_conn = Some(conn.clone());
              self.querying_pid = Some(pid.to_string());
              (
                DbTaskResult::Pending,
                Some(PostgresTask::Query(spawn_query_task(conn, first_query, statement_type))),
              )
            },
            Err(e) => {
              log::error!("Connection acquisition failed: {e:?}");
              (DbTaskResult::Finished(QueryResultsWithMetadata::new(Err(e), statement_type)), None)
            },
          }
        }
      },
      Some(PostgresTask::Query(handle)) => {
        if !handle.is_finished() {
          (DbTaskResult::Pending, Some(PostgresTask::Query(handle)))
        } else {
          let result = handle.await?;
          self.querying_conn = None;
          self.querying_pid = None;
          (DbTaskResult::Finished(result), None)
        }
      },
      Some(PostgresTask::TxConnect {
        handle,
        first_query,
        statement_type,
        display_statement_type,
      }) => {
        if !handle.is_finished() {
          (
            DbTaskResult::Pending,
            Some(PostgresTask::TxConnect {
              handle,
              first_query,
              statement_type,
              display_statement_type,
            }),
          )
        } else {
          match handle.await? {
            Ok((tx, pid)) => {
              log::info!("Starting transaction with PID {}", pid.clone());
              self.querying_pid = Some(pid.to_string());
              (
                DbTaskResult::Pending,
                Some(PostgresTask::TxStart(spawn_tx_task(
                  tx,
                  first_query,
                  statement_type,
                  *display_statement_type,
                ))),
              )
            },
            Err(e) => {
              log::error!("Transaction didn't start: {e:?}");
              (
                DbTaskResult::Finished(QueryResultsWithMetadata::with_display_statement_type(
                  Err(e),
                  Some(statement_type),
                  Some(*display_statement_type),
                )),
                None,
              )
            },
          }
        }
      },
      Some(PostgresTask::TxStart(handle)) => {
        if !handle.is_finished() {
          (DbTaskResult::Pending, Some(PostgresTask::TxStart(handle)))
        } else {
          let (result, tx) = handle.await?;
          let rows_affected = match &result.results {
            Ok(rows) => rows.rows_affected,
            _ => None,
          };
          match result {
            // if tx failed to start, return the error immediately
            QueryResultsWithMetadata {
              results: Err(e),
              statement_type,
              display_statement_type,
            } => {
              log::error!("Transaction didn't start: {e:?}");
              self.querying_conn = None;
              self.querying_pid = None;
              (
                DbTaskResult::Finished(QueryResultsWithMetadata::with_display_statement_type(
                  Err(e),
                  statement_type,
                  display_statement_type,
                )),
                None,
              )
            },
            _ => (
              DbTaskResult::ConfirmTx(
                rows_affected,
                result.statement_type.clone(),
                result.display_statement_type.clone(),
              ),
              Some(PostgresTask::TxPending(Box::new((tx, result)))),
            ),
          }
        }
      },
      Some(PostgresTask::TxPending(b)) => (DbTaskResult::Pending, Some(PostgresTask::TxPending(b))),
    };
    self.task = next_task;
    Ok(task_result)
  }

  async fn start_tx(&mut self, query: String) -> Result<()> {
    let (first_query, statement_type) = super::get_first_query(query, Driver::Postgres)?;
    let display_statement_type = super::get_display_statement_for_execution_type(&statement_type);
    let statement_type = super::get_statement_for_execution_type(&statement_type);
    let pool = self.pool.clone().unwrap();
    self.task = Some(PostgresTask::TxConnect {
      handle: tokio::spawn(
        async move {
          let mut tx = pool.begin().await?;
          let pid = connection_pid(&mut tx).await?;
          Ok((tx, pid))
        }
        .in_current_span(),
      ),
      first_query,
      statement_type,
      display_statement_type: Box::new(display_statement_type),
    });
    Ok(())
  }

  async fn commit_tx(&mut self) -> Result<Option<QueryResultsWithMetadata>> {
    if !matches!(self.task, Some(PostgresTask::TxPending(_))) {
      Ok(None)
    } else {
      match self.task.take() {
        Some(PostgresTask::TxPending(b)) => {
          b.0.commit().await?;
          self.querying_conn = None;
          self.querying_pid = None;
          Ok(Some(b.1))
        },
        _ => Ok(None),
      }
    }
  }

  async fn rollback_tx(&mut self) -> Result<()> {
    if let Some(PostgresTask::TxPending(b)) = self.task.take() {
      b.0.rollback().await?;
    }
    self.querying_conn = None;
    self.querying_pid = None;
    Ok(())
  }

  fn start_load_menu(&self) -> Result<JoinHandle<Result<Rows>>> {
    let pool = self.pool.clone().unwrap();
    Ok(tokio::spawn(async move {
      match query_with_pool(pool.clone(), MENU_QUERY.to_owned()).await {
        Ok(rows) => Ok(rows),
        Err(e) => {
          log::warn!("Falling back to simple PostgreSQL menu query: {e:?}");
          query_with_pool(pool, SIMPLE_MENU_QUERY.to_owned()).await
        },
      }
    }))
  }

  fn start_load_columns(
    &self,
    tables: Vec<TableRef>,
  ) -> Result<JoinHandle<Result<Vec<TableColumns>>>> {
    let pool = self.pool.clone().unwrap();
    Ok(tokio::spawn(async move {
      let mut output = Vec::with_capacity(tables.len());
      for table in tables {
        let schema = table.schema.replace('\'', "''");
        let name = table.table.replace('\'', "''");
        let query = format!(
          "select column_name, data_type from information_schema.columns where table_schema = '{schema}' and table_name = '{name}' order by ordinal_position"
        );
        output.push(table_columns_from_rows(table, query_with_pool(pool.clone(), query).await?));
      }
      Ok(output)
    }))
  }

  fn preview_rows_query(&self, schema: &str, table: &str) -> String {
    format!("select * from \"{schema}\".\"{table}\" limit 100")
  }

  fn preview_columns_query(&self, schema: &str, table: &str) -> String {
    format!(
      "select column_name, * from information_schema.columns where table_schema = '{schema}' and table_name = '{table}'"
    )
  }

  fn preview_constraints_query(&self, schema: &str, table: &str) -> String {
    format!(
      "select constraint_name, * from information_schema.table_constraints where table_schema = '{schema}' and table_name = '{table}'"
    )
  }

  fn preview_indexes_query(&self, schema: &str, table: &str) -> String {
    format!(
      "select indexname, indexdef, * from pg_indexes where schemaname = '{schema}' and tablename = '{table}'"
    )
  }

  fn preview_policies_query(&self, schema: &str, table: &str) -> String {
    format!("select * from pg_policies where schemaname = '{schema}' and tablename = '{table}'")
  }

  fn preview_view_definition_query(&self, schema: &str, view: &str, materialized: bool) -> String {
    let relkind = if materialized { "m" } else { "v" };
    format!(
      "select pg_get_viewdef(c.oid, true) as definition
        from pg_class c
        join pg_namespace n on n.oid = c.relnamespace
        where n.nspname = '{schema}'
          and c.relname = '{view}'
          and c.relkind = '{relkind}'"
    )
  }

  fn preview_function_definition_query(&self, schema: &str, function: &str) -> String {
    format!(
      "select pg_get_functiondef(p.oid) as definition
        from pg_proc p
        join pg_namespace n on n.oid = p.pronamespace
        where n.nspname = '{schema}'
          and format('%s(%s)', p.proname, pg_get_function_identity_arguments(p.oid)) = '{function}'
          and p.prokind = 'f'"
    )
  }
}

impl PostgresDriver {
  pub fn new() -> Self {
    Self { pool: None, task: None, querying_conn: None, querying_pid: None }
  }

  fn build_connection_opts(
    args: crate::cli::Cli,
  ) -> Result<<<sqlx::Postgres as sqlx::Database>::Connection as sqlx::Connection>::Options> {
    let ssl_required = args.ssl_required;
    let opts = match args.connection_url {
      Some(url) => PgConnectOptions::from_str(url.trim().trim_start_matches("jdbc:"))?,
      None => {
        let mut opts = PgConnectOptions::new();

        if let Some(user) = args.user {
          opts = opts.username(&user);
        } else {
          let mut user: String = String::new();
          print!("username: ");
          io::stdout().flush().unwrap();
          io::stdin().read_line(&mut user).unwrap();
          user = user.trim().to_string();
          if !user.is_empty() {
            opts = opts.username(&user);
          }
        }

        if let Some(password) = args.password {
          opts = opts.password(&password);
        } else {
          let mut password =
            rpassword::prompt_password(format!("password for user {}: ", opts.get_username()))
              .unwrap();
          password = password.trim().to_string();
          if !password.is_empty() {
            opts = opts.password(&password);
          }
        }

        if let Some(host) = args.host {
          opts = opts.host(&host);
        } else {
          let mut host: String = String::new();
          print!("host (ex. localhost): ");
          io::stdout().flush().unwrap();
          io::stdin().read_line(&mut host).unwrap();
          host = host.trim().to_string();
          if !host.is_empty() {
            opts = opts.host(&host);
          }
        }

        if let Some(port) = args.port {
          opts = opts.port(port);
        } else {
          let mut port: String = String::new();
          print!("port (ex. 5432): ");
          io::stdout().flush().unwrap();
          io::stdin().read_line(&mut port).unwrap();
          port = port.trim().to_string();
          if !port.is_empty() {
            opts = opts.port(port.parse()?);
          }
        }

        if let Some(database) = args.database {
          opts = opts.database(&database);
        } else {
          let mut database: String = String::new();
          print!("database (ex. postgres): ");
          io::stdout().flush().unwrap();
          io::stdin().read_line(&mut database).unwrap();
          database = database.trim().to_string();
          if !database.is_empty() {
            opts = opts.database(&database);
          }
        }

        opts
      },
    };

    Ok(if ssl_required {
      match opts.get_ssl_mode() {
        PgSslMode::VerifyCa | PgSslMode::VerifyFull => opts,
        _ => opts.ssl_mode(PgSslMode::Require),
      }
    } else {
      opts
    })
  }
}

async fn query_with_pool(pool: Arc<sqlx::Pool<Postgres>>, query: String) -> Result<Rows> {
  query_with_stream(&*pool.clone(), &query).await
}

async fn query_with_conn(conn: &mut PgConnection, query: String) -> Result<Rows> {
  query_with_stream(conn, &query).await
}

fn spawn_query_task(
  conn: Arc<Mutex<PoolConnection<Postgres>>>,
  first_query: String,
  statement_type: Option<Statement>,
) -> QueryTask {
  tokio::spawn(async move {
    let results = query_with_conn(conn.lock().await.as_mut(), first_query.clone()).await;
    match results {
      Ok(ref rows) => {
        log::info!("{:?} rows, {:?} affected", rows.rows.len(), rows.rows_affected);
      },
      Err(ref e) => {
        log::error!("{e:?}");
      },
    };
    QueryResultsWithMetadata::new(results, statement_type.clone())
  })
}

fn spawn_tx_task(
  tx: PostgresTransaction,
  first_query: String,
  statement_type: Statement,
  display_statement_type: Statement,
) -> TransactionTask {
  tokio::spawn(async move {
    let (results, tx) = query_with_tx(tx, &first_query).await;
    match results {
      Ok(Either::Left(rows_affected)) => {
        log::info!("{rows_affected:?} rows affected");
        (
          QueryResultsWithMetadata {
            results: Ok(Rows { headers: vec![], rows: vec![], rows_affected: Some(rows_affected) }),
            statement_type: Some(statement_type),
            display_statement_type: Some(display_statement_type),
          },
          tx,
        )
      },
      Ok(Either::Right(rows)) => {
        log::info!("{:?} rows affected", rows.rows_affected);
        (
          QueryResultsWithMetadata::with_display_statement_type(
            Ok(rows),
            Some(statement_type),
            Some(display_statement_type),
          ),
          tx,
        )
      },
      Err(e) => {
        log::error!("{e:?}");
        (
          QueryResultsWithMetadata::with_display_statement_type(
            Err(e),
            Some(statement_type),
            Some(display_statement_type),
          ),
          tx,
        )
      },
    }
  })
}

async fn query_with_stream<'a, E>(e: E, query: &'a str) -> Result<Rows>
where
  E: sqlx::Executor<'a, Database = sqlx::Postgres>,
{
  let mut stream = sqlx::raw_sql(query).fetch_many(e);
  let mut query_rows = vec![];
  let mut query_rows_affected: Option<u64> = None;
  let mut headers: Headers = vec![];
  while let Some(item) = stream.next().await {
    match item {
      Ok(Either::Left(result)) => {
        // For non-SELECT queries
        query_rows_affected = Some(result.rows_affected());
      },
      Ok(Either::Right(row)) => {
        // For SELECT queries
        query_rows.push(row_to_vec(&row));
        if headers.is_empty() {
          headers = get_headers(&row);
        }
      },
      Err(e) => return Err(eyre::Report::new(e)),
    }
  }
  Ok(Rows { rows_affected: query_rows_affected, headers, rows: query_rows })
}

async fn query_with_tx(
  mut tx: PostgresTransaction,
  query: &str,
) -> (Result<Either<u64, Rows>>, PostgresTransaction)
where
  for<'c> <sqlx::Postgres as sqlx::Database>::Arguments<'c>:
    sqlx::IntoArguments<'c, sqlx::Postgres>,
  for<'c> &'c mut <sqlx::Postgres as sqlx::Database>::Connection:
    sqlx::Executor<'c, Database = sqlx::Postgres>,
{
  let first_query = super::get_first_query(query.to_string(), Driver::Postgres);
  match first_query {
    Ok((first_query, statement_type)) => match statement_type {
      statement_type if super::should_stream_tx_results(&statement_type) => {
        let result = query_with_stream(&mut *tx, &first_query).await;
        match result {
          Ok(result) => (Ok(Either::Right(result)), tx),
          Err(e) => (Err(e), tx),
        }
      },
      _ => {
        let result = sqlx::query(&first_query).execute(&mut *tx).await;
        match result {
          Ok(result) => (Ok(Either::Left(result.rows_affected())), tx),
          Err(e) => (Err(e.into()), tx),
        }
      },
    },
    Err(e) => (Err(eyre::Report::new(e)), tx),
  }
}

async fn connection_pid(conn: &mut PgConnection) -> Result<i32>
where
  for<'c> &'c mut <sqlx::Postgres as sqlx::Database>::Connection:
    sqlx::Executor<'c, Database = sqlx::Postgres>,
{
  Ok(sqlx::query_scalar("SELECT pg_backend_pid()").fetch_one(conn).await?)
}

fn get_headers(row: &<sqlx::Postgres as sqlx::Database>::Row) -> Headers {
  row
    .columns()
    .iter()
    .map(|col| Header { name: col.name().to_string(), type_name: col.type_info().to_string() })
    .collect()
}

fn row_to_vec(row: &<sqlx::Postgres as sqlx::Database>::Row) -> Vec<String> {
  row.columns().iter().map(|col| parse_value(row, col).unwrap().string).collect()
}

// parsed based on https://docs.rs/sqlx/latest/sqlx/postgres/types/index.html
fn parse_value(
  row: &<Postgres as sqlx::Database>::Row,
  col: &<Postgres as sqlx::Database>::Column,
) -> Option<Value> {
  let col_type = col.type_info().to_string();
  if row.try_get_raw(col.ordinal()).is_ok_and(|v| v.is_null()) {
    return Some(Value { parse_error: false, string: "NULL".to_string(), is_null: true });
  }
  match col_type.to_uppercase().as_str() {
    "TIMESTAMPTZ" => {
      Some(row.try_get::<chrono::DateTime<chrono::Utc>, usize>(col.ordinal()).map_or(
        Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
        |received| Value { parse_error: false, string: received.to_string(), is_null: false },
      ))
    },
    "TIMESTAMP" => Some(row.try_get::<chrono::NaiveDateTime, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "DATE" => Some(row.try_get::<chrono::NaiveDate, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "TIME" => Some(row.try_get::<chrono::NaiveTime, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "UUID" => Some(row.try_get::<Uuid, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "INET" | "CIDR" => Some(row.try_get::<std::net::IpAddr, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "JSON" | "JSONB" => Some(row.try_get::<serde_json::Value, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "BOOL" => Some(row.try_get::<bool, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "SMALLINT" | "SMALLSERIAL" | "INT2" => Some(row.try_get::<i16, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "INT" | "SERIAL" | "INT4" => Some(row.try_get::<i32, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "BIGINT" | "BIGSERIAL" | "INT8" => Some(row.try_get::<i64, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "REAL" | "FLOAT4" => Some(row.try_get::<f32, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "DOUBLE PRECISION" | "FLOAT8" => Some(row.try_get::<f64, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value { parse_error: false, string: received.to_string(), is_null: false },
    )),
    "TEXT" | "VARCHAR" | "NAME" | "CITEXT" | "BPCHAR" | "CHAR" => {
      Some(row.try_get::<String, usize>(col.ordinal()).map_or(
        Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
        |received| Value { parse_error: false, string: received.to_string(), is_null: false },
      ))
    },
    "BYTEA" => Some(row.try_get::<Vec<u8>, usize>(col.ordinal()).map_or(
      Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
      |received| Value {
        parse_error: false,
        string: received.iter().fold(String::new(), |mut output, b| {
          let _ = write!(output, "{b:02X}");
          output
        }),
        is_null: false,
      },
    )),
    "VOID" => Some(Value { parse_error: false, string: "".to_string(), is_null: false }),
    _ if col_type.to_uppercase().ends_with("[]") => {
      let array_type = col_type.to_uppercase().replace("[]", "");
      match array_type.as_str() {
        "TIMESTAMPTZ" => {
          Some(row.try_get::<Vec<chrono::DateTime<chrono::Utc>>, usize>(col.ordinal()).map_or(
            Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
            |received| Value {
              parse_error: false,
              string: vec_to_string(received),
              is_null: false,
            },
          ))
        },
        "TIMESTAMP" => {
          Some(row.try_get::<Vec<chrono::NaiveDateTime>, usize>(col.ordinal()).map_or(
            Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
            |received| Value {
              parse_error: false,
              string: vec_to_string(received),
              is_null: false,
            },
          ))
        },
        "DATE" => Some(row.try_get::<Vec<chrono::NaiveDate>, usize>(col.ordinal()).map_or(
          Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
          |received| Value { parse_error: false, string: vec_to_string(received), is_null: false },
        )),
        "TIME" => Some(row.try_get::<Vec<chrono::NaiveTime>, usize>(col.ordinal()).map_or(
          Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
          |received| Value { parse_error: false, string: vec_to_string(received), is_null: false },
        )),
        "UUID" => Some(row.try_get::<Vec<Uuid>, usize>(col.ordinal()).map_or(
          Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
          |received| Value { parse_error: false, string: vec_to_string(received), is_null: false },
        )),
        "INET" | "CIDR" => Some(row.try_get::<Vec<std::net::IpAddr>, usize>(col.ordinal()).map_or(
          Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
          |received| Value { parse_error: false, string: vec_to_string(received), is_null: false },
        )),
        "JSON" | "JSONB" => {
          Some(row.try_get::<Vec<serde_json::Value>, usize>(col.ordinal()).map_or(
            Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
            |received| Value {
              parse_error: false,
              string: vec_to_string(received),
              is_null: false,
            },
          ))
        },
        "BOOL" => Some(row.try_get::<Vec<bool>, usize>(col.ordinal()).map_or(
          Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
          |received| Value { parse_error: false, string: vec_to_string(received), is_null: false },
        )),
        "SMALLINT" | "SMALLSERIAL" | "INT2" => {
          Some(row.try_get::<Vec<i16>, usize>(col.ordinal()).map_or(
            Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
            |received| Value {
              parse_error: false,
              string: vec_to_string(received),
              is_null: false,
            },
          ))
        },
        "INT" | "SERIAL" | "INT4" => Some(row.try_get::<Vec<i32>, usize>(col.ordinal()).map_or(
          Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
          |received| Value { parse_error: false, string: vec_to_string(received), is_null: false },
        )),
        "BIGINT" | "BIGSERIAL" | "INT8" => {
          Some(row.try_get::<Vec<i64>, usize>(col.ordinal()).map_or(
            Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
            |received| Value {
              parse_error: false,
              string: vec_to_string(received),
              is_null: false,
            },
          ))
        },
        "REAL" | "FLOAT4" => Some(row.try_get::<Vec<f32>, usize>(col.ordinal()).map_or(
          Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
          |received| Value { parse_error: false, string: vec_to_string(received), is_null: false },
        )),
        "DOUBLE PRECISION" | "FLOAT8" => {
          Some(row.try_get::<Vec<f64>, usize>(col.ordinal()).map_or(
            Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
            |received| Value {
              parse_error: false,
              string: vec_to_string(received),
              is_null: false,
            },
          ))
        },
        "TEXT" | "VARCHAR" | "NAME" | "CITEXT" | "BPCHAR" | "CHAR" => {
          Some(row.try_get::<Vec<String>, usize>(col.ordinal()).map_or(
            Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
            |received| Value {
              parse_error: false,
              string: vec_to_string(received),
              is_null: false,
            },
          ))
        },
        "BYTEA" => Some(row.try_get::<Vec<u8>, usize>(col.ordinal()).map_or(
          Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
          |received| Value {
            parse_error: false,
            string: received.iter().fold(String::new(), |mut output, b| {
              let _ = write!(output, "{b:02X}");
              output
            }),
            is_null: false,
          },
        )),
        _ => {
          // try to cast custom or other types to strings
          Some(row.try_get_unchecked::<Vec<String>, usize>(col.ordinal()).map_or(
            Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
            |received| Value {
              parse_error: false,
              string: vec_to_string(received),
              is_null: false,
            },
          ))
        },
      }
    },
    _ => {
      // try to cast custom or other types to strings
      Some(row.try_get_unchecked::<String, usize>(col.ordinal()).map_or(
        Value { parse_error: true, string: "_ERROR_".to_string(), is_null: false },
        |received| Value { parse_error: false, string: received.to_string(), is_null: false },
      ))
    },
  }
}

#[cfg(test)]
mod tests {
  use clap::Parser;
  use sqlparser::{ast::Statement, dialect::PostgreSqlDialect, parser::ParserError};
  use sqlx::postgres::PgSslMode;

  use super::*;
  use crate::database::{ExecutionType, ParseError, get_execution_type, get_first_query};

  #[test]
  fn ssl_required_overrides_url_ssl_mode() {
    let args = crate::cli::Cli::parse_from([
      "rainfrog",
      "--url",
      "postgres://localhost/postgres?sslmode=disable",
      "--ssl-required",
    ]);
    let opts = PostgresDriver::build_connection_opts(args).unwrap();
    assert!(matches!(opts.get_ssl_mode(), PgSslMode::Require));
  }

  #[test]
  fn ssl_required_preserves_stricter_url_ssl_modes() {
    for (mode, expected) in
      [("verify-ca", PgSslMode::VerifyCa), ("verify-full", PgSslMode::VerifyFull)]
    {
      let args = crate::cli::Cli::parse_from([
        "rainfrog",
        "--url",
        &format!("postgres://localhost/postgres?sslmode={mode}"),
        "--ssl-required",
      ]);
      let opts = PostgresDriver::build_connection_opts(args).unwrap();
      assert!(
        matches!(
          (opts.get_ssl_mode(), expected),
          (PgSslMode::VerifyCa, PgSslMode::VerifyCa)
            | (PgSslMode::VerifyFull, PgSslMode::VerifyFull)
        ),
        "mode {mode} was not preserved"
      );
    }
  }

  #[test]
  fn ssl_required_applies_to_structured_options() {
    let args = crate::cli::Cli::parse_from([
      "rainfrog",
      "--driver",
      "postgres",
      "--username",
      "user",
      "--password",
      "password",
      "--host",
      "localhost",
      "--port",
      "5432",
      "--database",
      "postgres",
      "--ssl-required",
    ]);
    let opts = PostgresDriver::build_connection_opts(args).unwrap();
    assert!(matches!(opts.get_ssl_mode(), PgSslMode::Require));
  }

  #[test]
  fn test_get_first_query() {
    type TestCase = (&'static str, Result<(String, Box<dyn Fn(Statement) -> bool>), ParseError>);

    let test_cases: Vec<TestCase> = vec![
      // single query
      (
        "SELECT * FROM users;",
        Ok(("SELECT * FROM users".to_string(), Box::new(|s| matches!(s, Statement::Query(_))))),
      ),
      // multiple queries
      (
        "SELECT * FROM users; DELETE FROM posts;",
        Err(ParseError::MoreThanOneStatement("Only one statement allowed per query".to_owned())),
      ),
      // empty query
      ("", Err(ParseError::EmptyQuery("Parsed query is empty".to_owned()))),
      // syntax error
      (
        "SELEC * FORM users;",
        Err(ParseError::SqlParserError(ParserError::ParserError(
          "Expected: an SQL statement, found: SELEC at Line: 1, Column: 1".to_owned(),
        ))),
      ),
      // lowercase
      (
        "select * from \"public\".\"users\"",
        Ok((
          "SELECT * FROM \"public\".\"users\"".to_owned(),
          Box::new(|s| matches!(s, Statement::Query(_))),
        )),
      ),
      // newlines
      (
        "select *\nfrom users;",
        Ok(("SELECT * FROM users".to_owned(), Box::new(|s| matches!(s, Statement::Query(_))))),
      ),
      // comment-only
      ("-- select * from users;", Err(ParseError::EmptyQuery("Parsed query is empty".to_owned()))),
      // commented line(s)
      (
        "-- select blah;\nselect * from users",
        Ok(("SELECT * FROM users".to_owned(), Box::new(|s| matches!(s, Statement::Query(_))))),
      ),
      (
        "-- select blah;\nselect * from users\n-- insert blah",
        Ok(("SELECT * FROM users".to_owned(), Box::new(|s| matches!(s, Statement::Query(_))))),
      ),
      // update
      (
        "UPDATE users SET name = 'John' WHERE id = 1",
        Ok((
          "UPDATE users SET name = 'John' WHERE id = 1".to_owned(),
          Box::new(|s| matches!(s, Statement::Update { .. })),
        )),
      ),
      // delete
      (
        "DELETE FROM users WHERE id = 1",
        Ok((
          "DELETE FROM users WHERE id = 1".to_owned(),
          Box::new(|s| matches!(s, Statement::Delete(_))),
        )),
      ),
      // drop
      (
        "DROP TABLE users",
        Ok(("DROP TABLE users".to_owned(), Box::new(|s| matches!(s, Statement::Drop { .. })))),
      ),
      // explain
      (
        "EXPLAIN SELECT * FROM users",
        Ok((
          "EXPLAIN SELECT * FROM users".to_owned(),
          Box::new(|s| matches!(s, Statement::Explain { .. })),
        )),
      ),
    ];

    let dialect = Box::new(PostgreSqlDialect {});

    for (input, expected_output) in test_cases {
      let result = get_first_query(input.to_string(), Driver::Postgres);
      match (result, expected_output) {
        (Ok((query, statement_type)), Ok((expected_query, match_statement))) => {
          assert_eq!(query, expected_query);
          assert!(match_statement(statement_type));
        },
        (Err(ParseError::EmptyQuery(msg)), Err(ParseError::EmptyQuery(expected_msg))) => {
          assert_eq!(msg, expected_msg)
        },
        (
          Err(ParseError::MoreThanOneStatement(msg)),
          Err(ParseError::MoreThanOneStatement(expected_msg)),
        ) => {
          assert_eq!(msg, expected_msg)
        },
        (Err(ParseError::SqlParserError(msg)), Err(ParseError::SqlParserError(expected_msg))) => {
          assert_eq!(msg, expected_msg)
        },
        _ => panic!("Unexpected result for input: {input}"),
      }
    }
  }

  #[test]
  fn test_execution_type_postgres() {
    let test_cases = vec![
      ("DELETE FROM users WHERE id = 1", ExecutionType::Transaction),
      ("DROP TABLE users", ExecutionType::Confirm),
      ("UPDATE users SET name = 'John' WHERE id = 1", ExecutionType::Transaction),
      ("SELECT * FROM users", ExecutionType::Normal),
      ("INSERT INTO users (name) VALUES ('John')", ExecutionType::Normal),
      ("EXPLAIN ANALYZE DELETE FROM users WHERE id = 1", ExecutionType::Transaction),
      ("EXPLAIN ANALYZE DROP TABLE users", ExecutionType::Confirm),
      ("EXPLAIN SELECT * FROM users", ExecutionType::Normal),
      ("EXPLAIN ANALYZE SELECT * FROM users WHERE id = 1", ExecutionType::Normal),
    ];

    for (query, expected) in test_cases {
      assert_eq!(
        get_execution_type(query.to_string(), false, Driver::Postgres).unwrap().0,
        expected,
        "Failed for query: {query}"
      );
    }
  }
}
