mod connect_options;

use std::sync::Arc;

use async_trait::async_trait;
use color_eyre::eyre::Result;
use connect_options::OracleConnectOptions;
use oracle::{Connection, pool::Pool};
use sqlparser::ast::Statement;
use tokio::task::{self, JoinHandle};

use crate::cli::Driver;

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

type ConnectionTask = JoinHandle<Result<Arc<Connection>>>;
type TransactionTask = JoinHandle<Result<QueryResultsWithMetadata>>;
enum OracleTask {
  QueryConnect {
    handle: ConnectionTask,
    first_query: String,
    statement_type: Option<Statement>,
    display_statement_type: Option<Statement>,
    use_query_api: bool,
  },
  Query(QueryTask),
  TxConnect {
    handle: ConnectionTask,
    first_query: String,
    statement_type: Option<Statement>,
    display_statement_type: Option<Statement>,
    stream_results: bool,
  },
  TxStart(TransactionTask),
  TxPending(Box<QueryResultsWithMetadata>),
}

#[derive(Default)]
pub struct OracleDriver {
  pool: Option<Arc<oracle::pool::Pool>>,
  task: Option<OracleTask>,
  querying_conn: Option<Arc<Connection>>,
}

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

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

  async fn init(&mut self, args: crate::cli::Cli) -> Result<String> {
    let connection_opts = OracleConnectOptions::build_connection_opts(args)?;

    let (user, password, connection_string, port, db_name) =
      connection_opts.get_connection_options().map_err(|e| color_eyre::eyre::eyre!(e))?;
    let pool = Arc::new(
      oracle::pool::PoolBuilder::new(user, password, connection_string)
        .max_connections(3)
        .build()?,
    );
    self.pool = Some(pool);

    Ok(format!("{}/{}", port, db_name))
  }

  async fn start_query(&mut self, query: String, bypass_parser: bool) -> Result<()> {
    let (first_query, statement_type, display_statement_type, use_query_api) = if bypass_parser {
      (query, None, None, false)
    } else {
      let (first, stmt) = super::get_first_query(query, Driver::Oracle)?;
      let use_query_api = matches!(stmt, Statement::Query(_));
      (
        first,
        Some(super::get_statement_for_execution_type(&stmt)),
        Some(super::get_display_statement_for_execution_type(&stmt)),
        use_query_api,
      )
    };
    let pool = self.pool.clone().unwrap();

    let span = tracing::Span::current();
    self.task = Some(OracleTask::QueryConnect {
      handle: task::spawn_blocking(move || {
        let _enter = span.enter();
        Ok(Arc::new(pool.get()?))
      }),
      first_query,
      statement_type,
      display_statement_type,
      use_query_api,
    });

    Ok(())
  }

  async fn abort_query(&mut self) -> Result<bool> {
    if let Some(task) = self.task.take() {
      match task {
        OracleTask::QueryConnect { handle, .. } => handle.abort(),
        OracleTask::Query(handle) => handle.abort(),
        OracleTask::TxConnect { handle, .. } => handle.abort(),
        OracleTask::TxStart(handle) => handle.abort(),
        _ => {},
      };
      if let Some(conn) = &self.querying_conn {
        conn.break_execution()?;
      }
      self.querying_conn = None;
      Ok(true)
    } else {
      self.querying_conn = 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(OracleTask::QueryConnect {
        handle,
        first_query,
        statement_type,
        display_statement_type,
        use_query_api,
      }) => {
        if !handle.is_finished() {
          (
            DbTaskResult::Pending,
            Some(OracleTask::QueryConnect {
              handle,
              first_query,
              statement_type,
              display_statement_type,
              use_query_api,
            }),
          )
        } else {
          match handle.await? {
            Ok(conn) => {
              self.querying_conn = Some(conn.clone());
              (
                DbTaskResult::Pending,
                Some(OracleTask::Query(spawn_query_task(
                  conn,
                  first_query,
                  statement_type,
                  display_statement_type,
                  use_query_api,
                ))),
              )
            },
            Err(e) => {
              log::error!("Connection acquisition failed: {e:?}");
              (
                DbTaskResult::Finished(QueryResultsWithMetadata::with_display_statement_type(
                  Err(e),
                  statement_type,
                  display_statement_type,
                )),
                None,
              )
            },
          }
        }
      },
      Some(OracleTask::Query(handle)) => {
        if !handle.is_finished() {
          (DbTaskResult::Pending, Some(OracleTask::Query(handle)))
        } else {
          self.querying_conn = None;
          (DbTaskResult::Finished(handle.await?), None)
        }
      },
      Some(OracleTask::TxConnect {
        handle,
        first_query,
        statement_type,
        display_statement_type,
        stream_results,
      }) => {
        if !handle.is_finished() {
          (
            DbTaskResult::Pending,
            Some(OracleTask::TxConnect {
              handle,
              first_query,
              statement_type,
              display_statement_type,
              stream_results,
            }),
          )
        } else {
          match handle.await? {
            Ok(conn) => {
              self.querying_conn = Some(conn.clone());
              (
                DbTaskResult::Pending,
                Some(OracleTask::TxStart(spawn_tx_task(
                  conn,
                  first_query,
                  statement_type,
                  display_statement_type,
                  stream_results,
                ))),
              )
            },
            Err(e) => {
              log::error!("Transaction didn't start: {e:?}");
              (
                DbTaskResult::Finished(QueryResultsWithMetadata::with_display_statement_type(
                  Err(e),
                  statement_type,
                  display_statement_type,
                )),
                None,
              )
            },
          }
        }
      },
      Some(OracleTask::TxStart(handle)) => {
        if !handle.is_finished() {
          (DbTaskResult::Pending, Some(OracleTask::TxStart(handle)))
        } else {
          let result = 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;
              (
                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(OracleTask::TxPending(Box::new(result))),
            ),
          }
        }
      },
      Some(OracleTask::TxPending(handle)) => {
        (DbTaskResult::Pending, Some(OracleTask::TxPending(handle)))
      },
    };
    self.task = next_task;
    Ok(task_result)
  }

  async fn start_tx(&mut self, query: String) -> Result<()> {
    let (first_query, stmt) = super::get_first_query(query, Driver::Oracle)?;
    let statement_type = Some(super::get_statement_for_execution_type(&stmt));
    let display_statement_type = Some(super::get_display_statement_for_execution_type(&stmt));
    let pool = self.pool.clone().unwrap();

    let span = tracing::Span::current();
    self.task = Some(OracleTask::TxConnect {
      handle: task::spawn_blocking(move || {
        let _enter = span.enter();
        Ok(Arc::new(pool.get()?))
      }),
      first_query,
      statement_type,
      display_statement_type,
      stream_results: super::should_stream_tx_results(&stmt),
    });

    Ok(())
  }

  async fn commit_tx(&mut self) -> Result<Option<QueryResultsWithMetadata>> {
    if let Some(OracleTask::TxPending(b)) = self.task.take()
      && let Some(self_conn) = self.querying_conn.clone()
    {
      self_conn.commit()?;
      self.querying_conn = None;
      Ok(Some(*b))
    } else {
      Ok(None)
    }
  }

  async fn rollback_tx(&mut self) -> Result<()> {
    if let Some(OracleTask::TxPending(b)) = self.task.take()
      && let Some(self_conn) = self.querying_conn.clone()
    {
      self_conn.rollback()?;
      self.querying_conn = None;
      Ok(())
    } else {
      Ok(())
    }
  }

  fn start_load_menu(&self) -> Result<JoinHandle<Result<Rows>>> {
    let pool = self.pool.clone().unwrap();
    Ok(task::spawn_blocking(move || {
      query_with_pool(
        &pool,
        "select user, table_name, 'table' as object_kind
        from user_tables
        where tablespace_name is not null
      union all
      select user, view_name, 'view' as object_kind
        from user_views
      union all
      select user, mview_name, 'materialized_view' as object_kind
        from user_mviews
      union all
      select user, object_name, 'function' as object_kind
        from user_objects
        where object_type = 'FUNCTION'
        order by 1, 3, 2",
      )
    }))
  }

  fn start_load_columns(
    &self,
    tables: Vec<TableRef>,
  ) -> Result<JoinHandle<Result<Vec<TableColumns>>>> {
    let pool = self.pool.clone().unwrap();
    Ok(task::spawn_blocking(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 all_tab_columns where owner = upper('{schema}') and table_name = upper('{name}') order by column_id"
        );
        output.push(table_columns_from_rows(table, query_with_pool(&pool, &query)?));
      }
      Ok(output)
    }))
  }

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

  fn preview_columns_query(&self, schema: &str, table: &str) -> String {
    format!("select * from user_tab_columns where table_name = '{}' and user = '{}'", table, schema)
  }

  fn preview_constraints_query(&self, schema: &str, table: &str) -> String {
    format!("select * from user_constraints where table_name = '{}' and user = '{}'", table, schema)
  }

  fn preview_indexes_query(&self, schema: &str, table: &str) -> String {
    format!("select * from user_ind_columns where table_name = '{}' and user = '{}'", table, schema)
  }

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

  fn preview_view_definition_query(&self, schema: &str, view: &str, materialized: bool) -> String {
    if materialized {
      return format!(
        "select query as definition from user_mviews where mview_name = '{}' and user = '{}'",
        view, schema
      );
    }
    format!(
      "select text as definition from user_views where view_name = '{}' and user = '{}'",
      view, schema
    )
  }

  fn preview_function_definition_query(&self, schema: &str, function: &str) -> String {
    format!(
      "select listagg(text) within group (order by line) as definition
        from user_source
        where type = 'FUNCTION'
          and name = '{}'
          and user = '{}'",
      function, schema
    )
  }
}

fn spawn_query_task(
  query_conn: Arc<Connection>,
  first_query: String,
  statement_type: Option<Statement>,
  display_statement_type: Option<Statement>,
  use_query_api: bool,
) -> QueryTask {
  tokio::spawn(async move {
    let results = if use_query_api {
      query_with_conn(query_conn.as_ref(), &first_query)
    } else {
      execute_with_conn(query_conn.as_ref(), &first_query).and_then(|rows| {
        query_conn
          .commit()
          .map_err(|e| color_eyre::eyre::eyre!("Error committing query: {}", e))?;
        Ok(rows)
      })
    };
    QueryResultsWithMetadata::with_display_statement_type(
      results,
      statement_type,
      display_statement_type,
    )
  })
}

fn spawn_tx_task(
  query_conn: Arc<Connection>,
  first_query: String,
  statement_type: Option<Statement>,
  display_statement_type: Option<Statement>,
  stream_results: bool,
) -> TransactionTask {
  tokio::spawn(async move {
    let results = if stream_results {
      query_with_conn(query_conn.as_ref(), &first_query)
    } else {
      execute_with_conn(query_conn.as_ref(), &first_query)
    };
    match results {
      Ok(ref rows) => {
        log::info!("{:?} rows, {:?} affected", rows.rows.len(), rows.rows_affected);
      },
      Err(ref e) => {
        log::error!("{e:?}");
      },
    };
    Ok(QueryResultsWithMetadata::with_display_statement_type(
      results,
      statement_type,
      display_statement_type,
    ))
  })
}

fn query_with_pool(pool: &Pool, query: &str) -> Result<Rows> {
  let conn = pool.get()?;
  query_with_conn(&conn, query)
}

fn query_with_conn(conn: &Connection, query: &str) -> Result<Rows> {
  let mut headers = Vec::new();
  let rows = conn
    .query(query, &[])
    .map_err(|e| color_eyre::eyre::eyre!("Error executing query: {}", e))?
    .filter_map(|row| row.ok())
    .map(|row| {
      if headers.is_empty() {
        headers = get_headers(&row);
      }

      row_to_vec(&row)
    })
    .collect::<Vec<_>>();

  Ok(Rows { headers, rows, rows_affected: None })
}

fn execute_with_conn(conn: &Connection, statement: &str) -> Result<Rows> {
  let result = conn
    .execute(statement, &[])
    .map_err(|e| color_eyre::eyre::eyre!("Error executing statement: {}", e))?;
  Ok(Rows { headers: Vec::new(), rows: Vec::new(), rows_affected: result.row_count().ok() })
}

fn get_headers(row: &oracle::Row) -> Vec<Header> {
  row
    .column_info()
    .iter()
    .map(|col| Header { name: col.name().to_string(), type_name: col.oracle_type().to_string() })
    .collect()
}

fn row_to_vec(row: &oracle::Row) -> Vec<String> {
  row.sql_values().iter().map(|v| v.to_string()).collect()
}

#[cfg(test)]
mod tests {
  use sqlparser::{ast::Statement, parser::ParserError};

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

  #[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 `users`",
        Ok(("SELECT * FROM `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(_))))),
      ),
      // 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 { .. })),
        )),
      ),
    ];

    for (input, expected_output) in test_cases {
      let result = get_first_query(input.to_string(), Driver::Oracle);
      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_mysql() {
    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::Oracle).unwrap().0,
        expected,
        "Failed for query: {}",
        query
      );
    }
  }
}
