// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use auth::tests::{DatabaseAuthInfo, MockUserProvider};
use chrono::{Datelike, NaiveDate};
use common_catalog::consts::DEFAULT_SCHEMA_NAME;
use common_query::Output;
use common_recordbatch::RecordBatch;
use common_runtime::Builder as RuntimeBuilder;
use common_runtime::runtime::BuilderBuild;
use common_time::{Timestamp, Timezone};
use datafusion_expr::LogicalPlan;
use datatypes::prelude::VectorRef;
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::value::Value;
use datatypes::vectors::{Int32Vector, TimestampMicrosecondVector, TimestampSecondVector};
use mysql_async::prelude::*;
use mysql_async::{Conn, Row, SslOpts};
use query::parser::PromQuery;
use query::query_engine::DescribeResult;
use servers::error::Result;
use servers::install_default_crypto_provider;
use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
use servers::query_handler::sql::{ServerSqlQueryHandlerRef, SqlQueryHandler};
use servers::server::Server;
use servers::tls::{ReloadableTlsServerConfig, TlsOption};
use session::context::QueryContextRef;
use sql::statements::statement::Statement;
use table::TableRef;
use table::test_util::MemTable;

use crate::create_testing_sql_query_handler;
use crate::mysql::{MysqlTextRow, TestingData, all_datatype_testing_data};

#[derive(Default)]
struct MysqlOpts<'a> {
    tls: TlsOption,
    auth_info: Option<DatabaseAuthInfo<'a>>,
    reject_no_database: bool,
}

fn create_mysql_server(table: TableRef, opts: MysqlOpts<'_>) -> Result<Box<dyn Server>> {
    let query_handler = create_testing_sql_query_handler(table);
    create_mysql_server_with_query_handler(query_handler, opts)
}

fn create_mysql_server_with_query_handler(
    query_handler: ServerSqlQueryHandlerRef,
    opts: MysqlOpts<'_>,
) -> Result<Box<dyn Server>> {
    let _ = install_default_crypto_provider();
    let io_runtime = RuntimeBuilder::default()
        .worker_threads(4)
        .thread_name("mysql-io-handlers")
        .build()
        .unwrap();

    let mut provider = MockUserProvider::default();
    if let Some(auth_info) = opts.auth_info {
        provider.set_authorization_info(auth_info);
    }

    let tls_server_config = Arc::new(
        ReloadableTlsServerConfig::try_new(opts.tls.clone())
            .expect("Failed to load certificates and keys"),
    );

    Ok(MysqlServer::create_server(
        io_runtime,
        Arc::new(MysqlSpawnRef::new(query_handler, Some(Arc::new(provider)))),
        Arc::new(MysqlSpawnConfig::new(
            opts.tls.should_force_tls(),
            tls_server_config,
            0,
            opts.reject_no_database,
            10000,
        )),
        None,
    ))
}

struct SessionTimezoneQueryHandler {
    inner: ServerSqlQueryHandlerRef,
}

#[async_trait]
impl SqlQueryHandler for SessionTimezoneQueryHandler {
    async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>> {
        if query == "SET time_zone = '+08:00'" {
            query_ctx.set_timezone(Timezone::hours_mins_opt(8, 0).unwrap());
            return vec![Ok(Output::new_with_affected_rows(0))];
        }

        self.inner.do_query(query, query_ctx).await
    }

    async fn do_analyze_stream_query(
        &self,
        query: &str,
        query_ctx: QueryContextRef,
    ) -> Result<Output> {
        self.inner.do_analyze_stream_query(query, query_ctx).await
    }

    async fn do_exec_plan(
        &self,
        plan: LogicalPlan,
        stmt: Option<Statement>,
        query_ctx: QueryContextRef,
    ) -> Result<Output> {
        self.inner.do_exec_plan(plan, stmt, query_ctx).await
    }

    async fn do_promql_query(
        &self,
        query: &PromQuery,
        query_ctx: QueryContextRef,
    ) -> Vec<Result<Output>> {
        self.inner.do_promql_query(query, query_ctx).await
    }

    async fn do_describe(
        &self,
        stmt: Statement,
        query_ctx: QueryContextRef,
    ) -> Result<Option<DescribeResult>> {
        self.inner.do_describe(stmt, query_ctx).await
    }

    async fn is_valid_schema(&self, catalog: &str, schema: &str) -> Result<bool> {
        self.inner.is_valid_schema(catalog, schema).await
    }
}

#[tokio::test]
async fn test_start_mysql_server() -> Result<()> {
    let table = MemTable::default_numbers_table();

    let mut mysql_server = create_mysql_server(table, Default::default())?;
    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();

    let result = mysql_server.start(listening).await;
    assert!(
        result
            .unwrap_err()
            .to_string()
            .contains("MySQL server has been started.")
    );
    Ok(())
}

#[tokio::test]
async fn test_reject_no_database() -> Result<()> {
    common_telemetry::init_default_ut_logging();
    let table = MemTable::default_numbers_table();
    let mut mysql_server = create_mysql_server(
        table,
        MysqlOpts {
            reject_no_database: true,
            ..Default::default()
        },
    )?;
    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();
    let server_addr = mysql_server.bind_addr().unwrap();
    let server_port = server_addr.port();

    let fail = create_connection(server_port, None, false).await;
    assert!(fail.is_err());
    let _ = create_connection(server_port, Some("public"), false)
        .await
        .unwrap();
    mysql_server.shutdown().await.unwrap();

    Ok(())
}

#[tokio::test]
async fn test_schema_validation() -> Result<()> {
    async fn generate_server(auth_info: DatabaseAuthInfo<'_>) -> Result<(Box<dyn Server>, u16)> {
        let table = MemTable::default_numbers_table();
        let mut mysql_server = create_mysql_server(
            table,
            MysqlOpts {
                auth_info: Some(auth_info),
                ..Default::default()
            },
        )?;
        let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
        mysql_server.start(listening).await.unwrap();
        let server_addr = mysql_server.bind_addr().unwrap();
        Ok((mysql_server, server_addr.port()))
    }

    common_telemetry::init_default_ut_logging();
    let (mysql_server, server_port) = generate_server(DatabaseAuthInfo {
        catalog: "greptime",
        schema: "public",
        username: "greptime",
    })
    .await?;

    let _ = create_connection_default_db_name(server_port, false)
        .await
        .unwrap();
    mysql_server.shutdown().await.unwrap();

    // change to another username
    let (mysql_server, server_port) = generate_server(DatabaseAuthInfo {
        catalog: "greptime",
        schema: "public",
        username: "no_access_user",
    })
    .await?;

    let fail = create_connection_default_db_name(server_port, false).await;
    assert!(fail.is_err());
    mysql_server.shutdown().await.unwrap();

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_shutdown_mysql_server() -> Result<()> {
    common_telemetry::init_default_ut_logging();

    let table = MemTable::default_numbers_table();

    let mut mysql_server = create_mysql_server(table, Default::default())?;
    let result = mysql_server.shutdown().await;
    assert!(
        result
            .unwrap_err()
            .to_string()
            .contains("MySQL server is not started.")
    );

    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();
    let server_addr = mysql_server.bind_addr().unwrap();
    let server_port = server_addr.port();

    let mut join_handles = vec![];
    for _ in 0..2 {
        join_handles.push(tokio::spawn(async move {
            for _ in 0..1000 {
                match create_connection_default_db_name(server_port, false).await {
                    Ok(mut connection) => {
                        let result: u32 = connection
                            .query_first("SELECT uint32s FROM numbers LIMIT 1")
                            .await
                            .unwrap()
                            .unwrap();
                        assert_eq!(result, 0);
                        tokio::time::sleep(Duration::from_millis(10)).await;
                    }
                    Err(e) => return Err(e),
                }
            }
            Ok(())
        }))
    }

    tokio::time::sleep(Duration::from_millis(100)).await;
    mysql_server.shutdown().await.unwrap();

    for handle in join_handles.iter_mut() {
        let result = handle.await.unwrap();
        assert!(result.is_err());
        assert!(result.unwrap_err().is_fatal());
    }
    Ok(())
}

#[tokio::test]
async fn test_query_all_datatypes() -> Result<()> {
    common_telemetry::init_default_ut_logging();

    let server_tls = TlsOption::default();
    let client_tls = false;

    do_test_query_all_datatypes(server_tls, client_tls).await?;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_server_prefer_secure_client_plain() -> Result<()> {
    do_test_query_all_datatypes_with_secure_server(servers::tls::TlsMode::Prefer, false, false)
        .await?;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_server_prefer_secure_client_plain_with_pkcs8_priv_key() -> Result<()> {
    do_test_query_all_datatypes_with_secure_server(servers::tls::TlsMode::Prefer, false, true)
        .await?;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_server_require_secure_client_secure() -> Result<()> {
    do_test_query_all_datatypes_with_secure_server(servers::tls::TlsMode::Require, true, false)
        .await?;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_server_require_secure_client_secure_with_pkcs8_priv_key() -> Result<()> {
    do_test_query_all_datatypes_with_secure_server(servers::tls::TlsMode::Require, true, true)
        .await?;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_server_required_secure_client_plain() -> Result<()> {
    let server_tls = TlsOption {
        mode: servers::tls::TlsMode::Require,
        cert_path: "tests/ssl/server.crt".to_owned(),
        key_path: "tests/ssl/server-rsa.key".to_owned(),
        ca_cert_path: String::new(),
        watch: false,
    };

    let client_tls = false;

    #[allow(unused)]
    let TestingData {
        column_schemas,
        mysql_columns_def,
        columns,
        mysql_text_output_rows,
    } = all_datatype_testing_data();
    let schema = Arc::new(Schema::new(column_schemas.clone()));
    let recordbatch = RecordBatch::new(schema, columns).unwrap();
    let table = MemTable::table("all_datatypes", recordbatch);

    let mut mysql_server = create_mysql_server(
        table,
        MysqlOpts {
            tls: server_tls,
            ..Default::default()
        },
    )?;

    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();
    let server_addr = mysql_server.bind_addr().unwrap();

    let r = create_connection(server_addr.port(), None, client_tls).await;
    assert!(r.is_err());
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_server_required_secure_client_plain_with_pkcs8_priv_key() -> Result<()> {
    let server_tls = TlsOption {
        mode: servers::tls::TlsMode::Require,
        cert_path: "tests/ssl/server.crt".to_owned(),
        key_path: "tests/ssl/server-pkcs8.key".to_owned(),
        ca_cert_path: String::new(),
        watch: false,
    };

    let client_tls = false;

    #[allow(unused)]
    let TestingData {
        column_schemas,
        mysql_columns_def,
        columns,
        mysql_text_output_rows,
    } = all_datatype_testing_data();
    let schema = Arc::new(Schema::new(column_schemas.clone()));
    let recordbatch = RecordBatch::new(schema, columns).unwrap();
    let table = MemTable::table("all_datatypes", recordbatch);

    let mut mysql_server = create_mysql_server(
        table,
        MysqlOpts {
            tls: server_tls,
            ..Default::default()
        },
    )?;

    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();
    let server_addr = mysql_server.bind_addr().unwrap();

    let r = create_connection_default_db_name(server_addr.port(), client_tls).await;
    assert!(r.is_err());
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_db_name() -> Result<()> {
    let server_tls = TlsOption::default();
    let client_tls = false;

    #[allow(unused)]
    let TestingData {
        column_schemas,
        mysql_columns_def,
        columns,
        mysql_text_output_rows,
    } = all_datatype_testing_data();
    let schema = Arc::new(Schema::new(column_schemas.clone()));
    let recordbatch = RecordBatch::new(schema, columns).unwrap();
    let table = MemTable::table("all_datatypes", recordbatch);

    let mut mysql_server = create_mysql_server(
        table,
        MysqlOpts {
            tls: server_tls,
            ..Default::default()
        },
    )?;

    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();
    let server_addr = mysql_server.bind_addr().unwrap();

    // None actually uses default database name
    let r = create_connection_default_db_name(server_addr.port(), client_tls).await;
    let _ = r.unwrap();

    let r = create_connection(server_addr.port(), Some("tomcat"), client_tls).await;
    assert!(r.is_err());
    Ok(())
}

async fn do_test_query_all_datatypes(server_tls: TlsOption, client_tls: bool) -> Result<()> {
    common_telemetry::init_default_ut_logging();
    let TestingData {
        column_schemas,
        mysql_columns_def,
        columns,
        mysql_text_output_rows,
    } = all_datatype_testing_data();
    let schema = Arc::new(Schema::new(column_schemas.clone()));
    let recordbatch = RecordBatch::new(schema, columns).unwrap();
    let table = MemTable::table("all_datatypes", recordbatch);

    let mut mysql_server = create_mysql_server(
        table,
        MysqlOpts {
            tls: server_tls,
            ..Default::default()
        },
    )?;

    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();
    let server_addr = mysql_server.bind_addr().unwrap();

    let mut connection = create_connection_default_db_name(server_addr.port(), client_tls)
        .await
        .unwrap();

    let mut result = connection
        .query_iter("SELECT * FROM all_datatypes LIMIT 3")
        .await
        .unwrap();
    let columns = result.columns().unwrap();
    assert_eq!(column_schemas.len(), columns.len());

    for (i, column) in columns.iter().enumerate() {
        assert_eq!(mysql_columns_def[i] as u8, column.column_type() as u8);
        assert_eq!(column_schemas[i].name, column.name_str());
    }

    let rows = result.collect::<MysqlTextRow>().await.unwrap();
    assert_eq!(3, rows.len());
    for (expected, actual) in mysql_text_output_rows.iter().take(3).zip(rows.iter()) {
        assert_eq!(expected, &actual.values);
    }
    Ok(())
}

#[tokio::test]
async fn test_mysql_text_protocol_out_of_range_timestamp_fails_closed() -> Result<()> {
    let timestamp = i64::MAX;
    assert!(
        Timestamp::new_second(timestamp)
            .to_chrono_datetime_with_timezone(None)
            .is_none(),
        "the target value must be outside Chrono's timestamp range"
    );

    let (table, schema) = timestamp_table("out_of_range_timestamp", Some(timestamp));
    assert!(
        schema.column_schemas()[1].is_nullable(),
        "the timestamp field must be nullable while Arrow marks this value valid"
    );

    let (result, health_check) =
        query_timestamp_with_mysql_text_protocol(table, "out_of_range_timestamp").await;
    assert_timestamp_overflow(result);
    assert_eq!(Some(1), health_check.unwrap());

    Ok(())
}

#[tokio::test]
async fn test_mysql_text_protocol_max_timestamp_with_session_timezone_fails_closed() -> Result<()> {
    let timestamp = Timestamp::MAX_SECOND.value();
    let (table, _) = timestamp_table("max_timestamp_with_session_timezone", Some(timestamp));
    let query_handler = Arc::new(SessionTimezoneQueryHandler {
        inner: create_testing_sql_query_handler(table),
    });
    let mut mysql_server =
        create_mysql_server_with_query_handler(query_handler, Default::default())?;
    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();

    let server_addr = mysql_server.bind_addr().unwrap();
    let mut connection = create_connection_default_db_name(server_addr.port(), false)
        .await
        .unwrap();
    connection
        .query_drop("SET time_zone = '+08:00'")
        .await
        .unwrap();

    let result = match connection
        .query_iter("SELECT id, ts FROM max_timestamp_with_session_timezone")
        .await
    {
        Ok(mut result) => result.collect::<MysqlTextRow>().await,
        Err(error) => Err(error),
    };
    assert_timestamp_overflow_with_value(result, timestamp);
    assert_eq!(Some(1), connection.query_first("SELECT 1").await.unwrap());

    mysql_server.shutdown().await.unwrap();
    Ok(())
}

#[tokio::test]
async fn test_mysql_binary_protocol_out_of_range_timestamp_fails_closed() -> Result<()> {
    let timestamp = i64::MAX;
    assert!(
        Timestamp::new_second(timestamp)
            .to_chrono_datetime_with_timezone(None)
            .is_none(),
        "the target value must be outside Chrono's timestamp range"
    );

    let (table, schema) = timestamp_table("out_of_range_timestamp", Some(timestamp));
    assert!(
        schema.column_schemas()[1].is_nullable(),
        "the timestamp field must be nullable while Arrow marks this value valid"
    );

    let (result, health_check) =
        query_timestamp_with_mysql_binary_protocol(table, "out_of_range_timestamp").await;
    assert_timestamp_overflow(result);
    assert_eq!(Some(1), health_check.unwrap());

    Ok(())
}

#[tokio::test]
async fn test_mysql_binary_protocol_chrono_representable_year_overflow_fails_closed() -> Result<()>
{
    let timestamp = Timestamp::new_second(
        NaiveDate::from_ymd_opt(100_000, 1, 1)
            .unwrap()
            .and_hms_opt(0, 0, 0)
            .unwrap()
            .and_utc()
            .timestamp(),
    );
    assert_eq!(
        100_000,
        timestamp
            .to_chrono_datetime_with_timezone(None)
            .unwrap()
            .year(),
        "the target value must be Chrono-representable"
    );

    let (table, schema) = timestamp_table(
        "chrono_representable_year_overflow",
        Some(timestamp.value()),
    );
    assert!(
        schema.column_schemas()[1].is_nullable(),
        "the timestamp field must be nullable while Arrow marks this value valid"
    );

    let (result, health_check) =
        query_timestamp_with_mysql_binary_protocol(table, "chrono_representable_year_overflow")
            .await;
    assert_timestamp_overflow_with_value(result, timestamp.value());
    assert_eq!(Some(1), health_check.unwrap());

    Ok(())
}

#[tokio::test]
async fn test_mysql_text_protocol_timestamp_controls() -> Result<()> {
    let (table, _) = timestamp_table("in_range_timestamp", Some(0));
    let (result, health_check) =
        query_timestamp_with_mysql_text_protocol(table, "in_range_timestamp").await;
    let rows = result.unwrap();
    assert_eq!(
        vec![
            Value::from(b"7".to_vec()),
            Value::from(b"1970-01-01 00:00:00".to_vec())
        ],
        rows[0].values
    );
    assert_eq!(Some(1), health_check.unwrap());

    let (table, _) = timestamp_table("null_timestamp", None);
    let (result, health_check) =
        query_timestamp_with_mysql_text_protocol(table, "null_timestamp").await;
    let rows = result.unwrap();
    assert_eq!(
        vec![Value::from(b"7".to_vec()), Value::Null],
        rows[0].values
    );
    assert_eq!(Some(1), health_check.unwrap());

    Ok(())
}

#[tokio::test]
async fn test_mysql_binary_protocol_timestamp_controls() -> Result<()> {
    let (table, _) = timestamp_table("in_range_timestamp", Some(0));
    let (result, health_check) =
        query_timestamp_with_mysql_binary_protocol(table, "in_range_timestamp").await;
    let rows = result.unwrap();
    assert_eq!(1, rows.len());
    assert_eq!(Some(&mysql_async::Value::Int(7)), rows[0].as_ref(0));
    assert_eq!(
        Some(&mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 0)),
        rows[0].as_ref(1)
    );
    assert_eq!(Some(1), health_check.unwrap());

    let (table, _) = timestamp_table("null_timestamp", None);
    let (result, health_check) =
        query_timestamp_with_mysql_binary_protocol(table, "null_timestamp").await;
    let rows = result.unwrap();
    assert_eq!(1, rows.len());
    assert_eq!(Some(&mysql_async::Value::Int(7)), rows[0].as_ref(0));
    assert_eq!(Some(&mysql_async::Value::NULL), rows[0].as_ref(1));
    assert_eq!(Some(1), health_check.unwrap());

    Ok(())
}

#[tokio::test]
async fn test_mysql_timestamp_precision_slot_isolation() -> Result<()> {
    let (result, health_check) = query_mysql_text_protocol(
        precision_timestamp_table("precision_timestamps"),
        "SELECT id, ts_a, ts_b FROM precision_timestamps".to_string(),
    )
    .await;
    let rows = result.unwrap();
    assert_eq!(
        vec![
            Value::from(b"7".to_vec()),
            Value::from(b"1970-01-01 00:00:00.123456".to_vec()),
            Value::Null,
        ],
        rows[0].values
    );
    assert_eq!(
        vec![
            Value::from(b"8".to_vec()),
            Value::Null,
            Value::from(b"1970-01-01 00:00:00.987654".to_vec()),
        ],
        rows[1].values
    );
    assert_eq!(Some(1), health_check.unwrap());

    let (result, health_check) = query_mysql_binary_protocol(
        precision_timestamp_table("precision_timestamps"),
        "SELECT id, ts_a, ts_b FROM precision_timestamps".to_string(),
    )
    .await;
    let rows = result.unwrap();
    assert_eq!(2, rows.len());
    assert_eq!(Some(&mysql_async::Value::Int(7)), rows[0].as_ref(0));
    assert_eq!(
        Some(&mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 123_456)),
        rows[0].as_ref(1)
    );
    assert_eq!(Some(&mysql_async::Value::NULL), rows[0].as_ref(2));
    assert_eq!(Some(&mysql_async::Value::Int(8)), rows[1].as_ref(0));
    assert_eq!(Some(&mysql_async::Value::NULL), rows[1].as_ref(1));
    assert_eq!(
        Some(&mysql_async::Value::Date(1970, 1, 1, 0, 0, 0, 987_654)),
        rows[1].as_ref(2)
    );
    assert_eq!(Some(1), health_check.unwrap());

    Ok(())
}

fn timestamp_table(table_name: &str, timestamp: Option<i64>) -> (TableRef, Arc<Schema>) {
    let schema = Arc::new(Schema::new(vec![
        ColumnSchema::new(
            "id",
            datatypes::prelude::ConcreteDataType::int32_datatype(),
            false,
        ),
        ColumnSchema::new(
            "ts",
            datatypes::prelude::ConcreteDataType::timestamp_second_datatype(),
            true,
        ),
    ]));
    let recordbatch = RecordBatch::new(
        schema.clone(),
        vec![
            Arc::new(Int32Vector::from_values(vec![7])) as VectorRef,
            Arc::new(TimestampSecondVector::from(vec![timestamp])) as VectorRef,
        ],
    )
    .unwrap();

    (MemTable::table(table_name, recordbatch), schema)
}

fn precision_timestamp_table(table_name: &str) -> TableRef {
    let schema = Arc::new(Schema::new(vec![
        ColumnSchema::new(
            "id",
            datatypes::prelude::ConcreteDataType::int32_datatype(),
            false,
        ),
        ColumnSchema::new(
            "ts_a",
            datatypes::prelude::ConcreteDataType::timestamp_microsecond_datatype(),
            true,
        ),
        ColumnSchema::new(
            "ts_b",
            datatypes::prelude::ConcreteDataType::timestamp_microsecond_datatype(),
            true,
        ),
    ]));
    let recordbatch = RecordBatch::new(
        schema,
        vec![
            Arc::new(Int32Vector::from_values(vec![7, 8])) as VectorRef,
            Arc::new(TimestampMicrosecondVector::from(vec![Some(123_456), None])) as VectorRef,
            Arc::new(TimestampMicrosecondVector::from(vec![None, Some(987_654)])) as VectorRef,
        ],
    )
    .unwrap();

    MemTable::table(table_name, recordbatch)
}

async fn query_timestamp_with_mysql_text_protocol(
    table: TableRef,
    table_name: &str,
) -> (
    mysql_async::Result<Vec<MysqlTextRow>>,
    mysql_async::Result<Option<u8>>,
) {
    query_mysql_text_protocol(table, format!("SELECT id, ts FROM {table_name}")).await
}

async fn query_mysql_text_protocol(
    table: TableRef,
    query: String,
) -> (
    mysql_async::Result<Vec<MysqlTextRow>>,
    mysql_async::Result<Option<u8>>,
) {
    let mut mysql_server = create_mysql_server(table, Default::default()).unwrap();
    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();

    let server_addr = mysql_server.bind_addr().unwrap();
    let mut connection = create_connection_default_db_name(server_addr.port(), false)
        .await
        .unwrap();
    let result = match connection.query_iter(query).await {
        Ok(mut result) => result.collect::<MysqlTextRow>().await,
        Err(error) => Err(error),
    };
    let health_check = connection.query_first("SELECT 1").await;

    mysql_server.shutdown().await.unwrap();
    (result, health_check)
}

async fn query_timestamp_with_mysql_binary_protocol(
    table: TableRef,
    table_name: &str,
) -> (
    mysql_async::Result<Vec<Row>>,
    mysql_async::Result<Option<u8>>,
) {
    query_mysql_binary_protocol(table, format!("SELECT id, ts FROM {table_name}")).await
}

async fn query_mysql_binary_protocol(
    table: TableRef,
    query: String,
) -> (
    mysql_async::Result<Vec<Row>>,
    mysql_async::Result<Option<u8>>,
) {
    let mut mysql_server = create_mysql_server(table, Default::default()).unwrap();
    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();

    let server_addr = mysql_server.bind_addr().unwrap();
    let mut connection = create_connection_default_db_name(server_addr.port(), false)
        .await
        .unwrap();
    let result = match connection.prep(query).await {
        Ok(statement) => connection.exec(statement, ()).await,
        Err(error) => Err(error),
    };
    let health_check = connection.query_first("SELECT 1").await;

    mysql_server.shutdown().await.unwrap();
    (result, health_check)
}

fn assert_timestamp_overflow<T>(result: mysql_async::Result<T>) {
    assert_timestamp_overflow_with_value(result, i64::MAX);
}

fn assert_timestamp_overflow_with_value<T>(result: mysql_async::Result<T>, timestamp: i64) {
    match result {
        Err(mysql_async::Error::Server(error)) => {
            assert_eq!(1210, error.code, "expected ER_WRONG_ARGUMENTS");
            assert!(
                error.message.contains("Timestamp overflow"),
                "expected timestamp overflow marker, got: {}",
                error.message
            );
            assert!(
                error.message.contains(&timestamp.to_string()),
                "expected raw timestamp value, got: {}",
                error.message
            );
            assert!(
                error.message.contains("Second"),
                "expected timestamp unit, got: {}",
                error.message
            );
        }
        Err(error) => panic!("expected a MySQL server timestamp overflow error, got: {error}"),
        Ok(_) => panic!("out-of-range timestamp must fail closed instead of returning a row"),
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_query_concurrently() -> Result<()> {
    common_telemetry::init_default_ut_logging();

    let table = MemTable::default_numbers_table();

    let mut mysql_server = create_mysql_server(table, Default::default())?;
    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();
    let server_addr = mysql_server.bind_addr().unwrap();
    let server_port = server_addr.port();

    let threads = 4;
    let expect_executed_queries_per_worker = 200;
    let queries = Arc::new(
        (0..100u32)
            .map(|expected| format!("SELECT uint32s FROM numbers WHERE uint32s = {expected}"))
            .collect::<Vec<_>>(),
    );
    let mut join_handles = vec![];
    for worker_id in 0..threads {
        let queries = queries.clone();
        join_handles.push(tokio::spawn(async move {
            let mut connection = create_connection_default_db_name(server_port, false)
                .await
                .unwrap();
            for i in 0..expect_executed_queries_per_worker {
                let expected: u32 = ((i + worker_id) % 100) as u32;
                let result: u32 = connection
                    .query_first(&queries[expected as usize])
                    .await
                    .unwrap()
                    .unwrap();
                assert_eq!(result, expected);

                let should_recreate_conn = expected == 1;
                if should_recreate_conn {
                    connection = create_connection_default_db_name(server_port, false)
                        .await
                        .unwrap();
                }
            }
            expect_executed_queries_per_worker
        }))
    }
    let mut total_pending_queries = threads * expect_executed_queries_per_worker;
    for handle in join_handles.iter_mut() {
        total_pending_queries -= handle.await.unwrap();
    }
    assert_eq!(0, total_pending_queries);
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_query_prepared() -> Result<()> {
    common_telemetry::init_default_ut_logging();
    let TestingData {
        column_schemas,
        mysql_columns_def: _,
        columns,
        mysql_text_output_rows: _,
    } = all_datatype_testing_data();
    let schema = Arc::new(Schema::new(column_schemas.clone()));
    let recordbatch = RecordBatch::new(schema, columns.clone()).unwrap();
    let table = MemTable::table("all_datatypes", recordbatch);

    let mut mysql_server = create_mysql_server(
        table,
        MysqlOpts {
            ..Default::default()
        },
    )?;

    let listening = "127.0.0.1:0".parse::<SocketAddr>().unwrap();
    mysql_server.start(listening).await.unwrap();
    let server_addr = mysql_server.bind_addr().unwrap();

    let mut connection = create_connection_default_db_name(server_addr.port(), false)
        .await
        .unwrap();

    test_prepare_all_type(column_schemas, columns, &mut connection).await;

    match connection
        .prep("SELECT `timestamp` FROM t WHERE `timestamp` > NOW() - INTERVAL '1 hour'")
        .await
    {
        Err(mysql_async::Error::Server(e)) => assert_eq!(
            "ERROR HY000 (1210): (InvalidArguments): Invalid prepare statement: Invalid SQL syntax: sql parser error: INTERVAL requires a unit after the literal value",
            e.to_string()
        ),
        _ => unreachable!(),
    }

    // Regression test for #8142: LIMIT ? should work in prepared statements.
    // The LIMIT placeholder should be inferred as Int64 so the MySQL prepare
    // response advertises the correct parameter count.
    {
        let stmt = connection
            .prep("SELECT uint32s FROM all_datatypes LIMIT ?")
            .await
            .unwrap();
        let rows: Vec<Row> = connection
            .exec(stmt, vec![mysql_async::Value::Int(1)])
            .await
            .unwrap();
        assert_eq!(rows.len(), 1, "LIMIT 1 should return 1 row");
    }

    // Also cover mixed placeholders: the WHERE placeholder is inferred from
    // the column type and the LIMIT placeholder is inferred from its context.
    {
        let stmt = connection
            .prep("SELECT uint32s FROM all_datatypes WHERE uint32s >= ? LIMIT ?")
            .await
            .unwrap();
        let rows: Vec<Row> = connection
            .exec(
                stmt,
                vec![mysql_async::Value::UInt(0), mysql_async::Value::UInt(1)],
            )
            .await
            .unwrap();
        assert_eq!(rows.len(), 1, "LIMIT 1 should return 1 row");
    }

    // Untyped placeholders should still be advertised in the MySQL prepare
    // response. This used to fail on the client side because the server
    // reported 0 parameters for `SELECT ?`.
    {
        let stmt = connection.prep("SELECT ?").await.unwrap();
        assert_eq!(stmt.num_params(), 1);

        let row: Option<String> = connection
            .exec_first(stmt, vec![mysql_async::Value::Bytes(b"can't".to_vec())])
            .await
            .unwrap();
        assert_eq!(row, Some("can't".to_string()));

        let stmt = connection.prep("SELECT ?").await.unwrap();
        let row: Option<String> = connection
            .exec_first(stmt, vec![mysql_async::Value::Bytes(b"a\\'b".to_vec())])
            .await
            .unwrap();
        assert_eq!(row, Some("a\\'b".to_string()));

        let stmt = connection.prep("SELECT ?").await.unwrap();
        let row: Option<Vec<u8>> = connection
            .exec_first(stmt, vec![mysql_async::Value::Bytes(vec![0xFF, 0xFE])])
            .await
            .unwrap();
        assert_eq!(row, Some(vec![0xFF, 0xFE]));

        let stmt = connection.prep("SELECT ?").await.unwrap();
        let row: Option<Option<String>> = connection
            .exec_first(stmt, vec![mysql_async::Value::NULL])
            .await
            .unwrap();
        assert_eq!(row, Some(None));
    }

    // Values inserted into the SQL text must not be processed again while
    // replacing later placeholders.
    {
        let stmt = connection.prep("SELECT ?, ?").await.unwrap();
        assert_eq!(stmt.num_params(), 2);

        let row: Option<(String, String)> = connection
            .exec_first(
                stmt,
                vec![
                    mysql_async::Value::Bytes(b"keep $2".to_vec()),
                    mysql_async::Value::Bytes(b"second".to_vec()),
                ],
            )
            .await
            .unwrap();
        assert_eq!(row, Some(("keep $2".to_string(), "second".to_string())));
    }

    // Non-placeholder question marks inside string literals must not affect
    // the advertised prepare parameter count.
    {
        let stmt = connection.prep("SELECT '?', ?").await.unwrap();
        assert_eq!(stmt.num_params(), 1);

        let row: Option<(String, String)> = connection
            .exec_first(stmt, vec![mysql_async::Value::Bytes(b"actual".to_vec())])
            .await
            .unwrap();
        assert_eq!(row, Some(("?".to_string(), "actual".to_string())));

        let stmt = connection.prep("SELECT '$1', ?").await.unwrap();
        assert_eq!(stmt.num_params(), 1);

        let row: Option<(String, String)> = connection
            .exec_first(stmt, vec![mysql_async::Value::Bytes(b"actual".to_vec())])
            .await
            .unwrap();
        assert_eq!(row, Some(("$1".to_string(), "actual".to_string())));

        let stmt = connection.prep("SELECT /* ? */ ? -- ?\n").await.unwrap();
        assert_eq!(stmt.num_params(), 1);

        let row: Option<String> = connection
            .exec_first(stmt, vec![mysql_async::Value::Bytes(b"commented".to_vec())])
            .await
            .unwrap();
        assert_eq!(row, Some("commented".to_string()));

        let stmt = connection.prep("SELECT '中文', ?").await.unwrap();
        assert_eq!(stmt.num_params(), 1);

        let row: Option<(String, String)> = connection
            .exec_first(stmt, vec![mysql_async::Value::Bytes(b"actual".to_vec())])
            .await
            .unwrap();
        assert_eq!(row, Some(("中文".to_string(), "actual".to_string())));
    }

    // Also cover mixed known and unknown placeholders. The projection
    // placeholder is untyped, while the WHERE placeholder is inferred from the
    // column type. The prepare response must advertise both parameters.
    {
        let stmt = connection
            .prep("SELECT ?, uint32s FROM all_datatypes WHERE uint32s >= ?")
            .await
            .unwrap();
        assert_eq!(stmt.num_params(), 2);

        let rows: Vec<Row> = connection
            .exec(
                stmt,
                vec![
                    mysql_async::Value::Bytes(b"unknown".to_vec()),
                    mysql_async::Value::UInt(0),
                ],
            )
            .await
            .unwrap();
        assert!(!rows.is_empty());
    }

    // LIMIT placeholders used to be a common case where DataFusion did not
    // infer a parameter type. The prepare response must still advertise the
    // parameter and execution must substitute it correctly.
    {
        let stmt = connection
            .prep("SELECT uint32s FROM all_datatypes ORDER BY uint32s LIMIT ?")
            .await
            .unwrap();
        assert_eq!(stmt.num_params(), 1);

        let rows: Vec<Row> = connection
            .exec(stmt, vec![mysql_async::Value::UInt(1)])
            .await
            .unwrap();
        assert_eq!(rows.len(), 1);
    }

    Ok(())
}

async fn test_prepare_all_type(
    column_schemas: Vec<ColumnSchema>,
    columns: Vec<VectorRef>,
    connection: &mut Conn,
) {
    let mut column_index = 0;
    for (stmt_id, schema) in (1..).zip(column_schemas) {
        let query = format!(
            "SELECT {} FROM all_datatypes WHERE {} = ?",
            schema.name, schema.name
        );
        let statement = connection.prep(query).await;
        let statement = statement.unwrap();
        assert_eq!(stmt_id, statement.id());

        let vector_ref = columns.get(column_index).unwrap();
        for vector_index in 0..vector_ref.len() {
            let v = vector_ref.get(vector_index);
            let v = if let Some(v) = prepare_convert_type(v) {
                v
            } else {
                continue;
            };

            let output: std::result::Result<Vec<Row>, mysql_async::Error> =
                connection.exec(statement.clone(), vec![v]).await;

            let rows = output.unwrap();
            assert!(!rows.is_empty());
        }
        column_index += 1;
    }
}

fn prepare_convert_type(item: Value) -> Option<mysql_async::Value> {
    let v = match item {
        Value::UInt8(u) => mysql_async::Value::UInt(u as u64),
        Value::UInt16(u) => mysql_async::Value::UInt(u as u64),
        Value::UInt32(u) => mysql_async::Value::UInt(u as u64),
        Value::UInt64(u) => mysql_async::Value::UInt(u),
        Value::Int8(i) => mysql_async::Value::Int(i as i64),
        Value::Int16(i) => mysql_async::Value::Int(i as i64),
        Value::Int32(i) => mysql_async::Value::Int(i as i64),
        Value::Int64(i) => mysql_async::Value::Int(i),
        Value::Float32(f) => mysql_async::Value::Float(f.into()),
        Value::Float64(f) => mysql_async::Value::Double(f.into()),
        Value::String(s) => mysql_async::Value::Bytes(s.as_utf8().as_bytes().to_vec()),
        Value::Binary(b) => mysql_async::Value::Bytes(b.to_vec()),
        _ => return None,
    };
    Some(v)
}

async fn create_connection_default_db_name(
    port: u16,
    ssl: bool,
) -> mysql_async::Result<mysql_async::Conn> {
    create_connection(port, Some(DEFAULT_SCHEMA_NAME), ssl).await
}

async fn create_connection(
    port: u16,
    db_name: Option<&str>,
    ssl: bool,
) -> mysql_async::Result<mysql_async::Conn> {
    let mut opts = mysql_async::OptsBuilder::default()
        .ip_or_hostname("127.0.0.1")
        .tcp_port(port)
        .prefer_socket(false)
        .wait_timeout(Some(1000))
        .user(Some("greptime".to_string()))
        .pass(Some("greptime".to_string()));

    if let Some(db_name) = db_name {
        opts = opts.db_name(Some(db_name.to_string()));
    }

    if ssl {
        let ssl_opts = SslOpts::default()
            .with_danger_skip_domain_validation(true)
            .with_danger_accept_invalid_certs(true);
        opts = opts.ssl_opts(ssl_opts)
    }

    mysql_async::Conn::new(opts).await
}

async fn do_test_query_all_datatypes_with_secure_server(
    server_tls_mode: servers::tls::TlsMode,
    client_tls: bool,
    is_pkcs8_priv_key: bool,
) -> Result<()> {
    let server_tls = TlsOption {
        mode: server_tls_mode,
        cert_path: "tests/ssl/server.crt".to_owned(),
        key_path: {
            if is_pkcs8_priv_key {
                "tests/ssl/server-pkcs8.key".to_owned()
            } else {
                "tests/ssl/server-rsa.key".to_owned()
            }
        },
        ca_cert_path: String::new(),
        watch: false,
    };

    do_test_query_all_datatypes(server_tls, client_tls).await
}
