// 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::fmt::Write as FmtWrite;
use std::io;
use std::time::Duration;

use arrow::array::{Array, AsArray};
use arrow::datatypes::{
    Date32Type, Decimal128Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type,
    Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalYearMonthType, UInt8Type,
    UInt16Type, UInt32Type, UInt64Type,
};
use arrow_schema::{DataType, IntervalUnit};
use chrono::{Datelike, NaiveDateTime};
use common_decimal::Decimal128;
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
use common_query::{Output, OutputData};
use common_recordbatch::{
    RecordBatch, SendableRecordBatchStream, map_dictionary_to_values_data_type,
};
use common_telemetry::{debug, error};
use common_time::{Date, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
use datafusion_common::ScalarValue;
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::SchemaRef;
use datatypes::types::jsonb_to_string;
use futures::StreamExt;
use opensrv_mysql::{
    Column, ColumnFlags, ColumnType, ErrorKind, OkResponse, QueryResultWriter, RowWriter,
    ToMysqlValue,
};
use session::SessionRef;
use session::context::QueryContextRef;
use snafu::prelude::*;
use tokio::io::AsyncWrite;

use crate::error::{
    self, ConvertSqlValueSnafu, DataFusionSnafu, InternalSnafu, NotSupportedSnafu, Result,
    TimestampOverflowSnafu,
};
use crate::metrics::*;

const MYSQL_DATETIME_MIN_YEAR: i32 = 1000;
const MYSQL_DATETIME_MAX_YEAR: i32 = 9999;

/// Try to write multiple output to the writer if possible.
pub async fn write_output<W: AsyncWrite + Send + Sync + Unpin>(
    mut writer: QueryResultWriter<'_, W>,
    query_context: QueryContextRef,
    session: SessionRef,
    outputs: Vec<Result<Output>>,
) -> Result<()> {
    if let Some(warning) = query_context.warning() {
        session.add_warning(warning);
    }

    enum Response {
        ResultSet {
            columns: Vec<Column>,
            stream: SendableRecordBatchStream,
        },
        AffectedRows(usize),
    }

    let mut responses = Vec::with_capacity(outputs.len());
    for output in outputs {
        match output {
            Ok(x) => {
                let output = match x.data {
                    OutputData::Stream(stream) => either::Left(stream),
                    OutputData::RecordBatches(record_batches) => {
                        either::Left(record_batches.as_stream())
                    }
                    OutputData::AffectedRows(rows) => either::Right(rows),
                };
                responses.push(match output {
                    either::Left(stream) => {
                        let schema = stream.schema();
                        let columns = match create_mysql_column_def(&schema) {
                            Ok(columns) => columns,
                            Err(e) => {
                                MysqlResultWriter::write_query_error(
                                    e,
                                    writer,
                                    query_context.clone(),
                                )
                                .await?;
                                return Ok(());
                            }
                        };
                        Response::ResultSet { columns, stream }
                    }
                    either::Right(rows) => Response::AffectedRows(rows),
                });
            }
            Err(e) => {
                MysqlResultWriter::write_query_error(e, writer, query_context.clone()).await?;
                return Ok(());
            }
        }
    }

    for response in &mut responses {
        writer = match response {
            Response::ResultSet { columns, stream } => {
                let mut row_writer = writer.start(columns).await?;
                while let Some(record_batch) = stream.next().await {
                    match record_batch {
                        Ok(record_batch) => {
                            if let Err(e) = MysqlResultWriter::write_recordbatch(
                                &mut row_writer,
                                record_batch,
                                query_context.clone(),
                            )
                            .await
                            {
                                let (kind, err) = handle_err(e, query_context);
                                row_writer.finish_error(kind, &err.as_bytes()).await?;
                                return Ok(());
                            }
                        }
                        Err(e) => {
                            let (kind, err) = handle_err(e, query_context);
                            row_writer.finish_error(kind, &err.as_bytes()).await?;
                            return Ok(());
                        }
                    }
                }
                row_writer.finish_one().await?
            }
            Response::AffectedRows(rows) => {
                MysqlResultWriter::write_affected_rows(writer, *rows, &session).await?
            }
        }
    }

    writer.no_more_results().await?;
    Ok(())
}

/// Handle GreptimeDB error, convert it to MySQL error
pub fn handle_err(e: impl ErrorExt, query_ctx: QueryContextRef) -> (ErrorKind, String) {
    let status_code = e.status_code();
    let kind = mysql_error_kind(&status_code);

    if status_code.should_log_error() {
        let root_error = e.root_cause().unwrap_or(&e);
        error!(e; "Failed to handle mysql query, code: {}, error: {}, db: {}", status_code, root_error.to_string(), query_ctx.get_db_string());
    } else {
        debug!(
            "Failed to handle mysql query, code: {}, db: {}, error: {:?}",
            status_code,
            query_ctx.get_db_string(),
            e
        );
    };
    let msg = e.output_msg();
    // Inline the status code to output message for MySQL
    let err_msg = format!("({status_code}): {msg}");

    (kind, err_msg)
}

struct MysqlResultWriter;

struct PrecisionTimestamp<'a> {
    formatted: &'a str,
    datetime: chrono::NaiveDateTime,
}

struct StagedTimestamp {
    datetime: Option<NaiveDateTime>,
    formatted: String,
}

impl StagedTimestamp {
    fn new() -> Self {
        Self {
            datetime: None,
            formatted: String::with_capacity(32),
        }
    }
}

impl<'a> ToMysqlValue for PrecisionTimestamp<'a> {
    fn to_mysql_text<W: std::io::Write>(&self, w: &mut W) -> io::Result<()> {
        self.formatted.to_mysql_text(w)
    }

    fn to_mysql_bin<W: std::io::Write>(&self, w: &mut W, c: &Column) -> io::Result<()> {
        self.datetime.to_mysql_bin(w, c)
    }
}

impl MysqlResultWriter {
    async fn write_affected_rows<'a, W: AsyncWrite + Unpin>(
        w: QueryResultWriter<'a, W>,
        rows: usize,
        session: &SessionRef,
    ) -> io::Result<QueryResultWriter<'a, W>> {
        let warnings = session.warnings_count() as u16;

        let next_writer = w
            .complete_one(OkResponse {
                affected_rows: rows as u64,
                warnings,
                ..Default::default()
            })
            .await?;
        Ok(next_writer)
    }

    async fn write_recordbatch<W: AsyncWrite + Unpin>(
        row_writer: &mut RowWriter<'_, '_, W>,
        record_batch: RecordBatch,
        query_context: QueryContextRef,
    ) -> Result<()> {
        let schema = record_batch.schema.clone();
        let record_batch = record_batch.into_df_record_batch();
        let mut timestamp_slots = vec![None; record_batch.num_columns()];
        let mut staged_timestamps = record_batch
            .columns()
            .iter()
            .enumerate()
            .filter(|(_, column)| matches!(column.data_type(), DataType::Timestamp(_, _)))
            .enumerate()
            .map(|(slot, (column_index, column))| {
                timestamp_slots[column_index] = Some(slot);
                (column, StagedTimestamp::new())
            })
            .collect::<Vec<_>>();
        for i in 0..record_batch.num_rows() {
            for (column, staged_timestamp) in &mut staged_timestamps {
                let column = *column;
                staged_timestamp.datetime = None;
                staged_timestamp.formatted.clear();
                if !column.is_null(i) {
                    let timestamp = datatypes::arrow_array::timestamp_array_value(column, i);
                    let datetime = timestamp
                        .to_chrono_datetime_with_timezone(Some(&query_context.timezone()))
                        .with_context(|| TimestampOverflowSnafu {
                            error: format!(
                                "timestamp {} overflow with unit {}",
                                timestamp.value(),
                                timestamp.unit()
                            ),
                        })?;
                    let year = datetime.year();
                    if !(MYSQL_DATETIME_MIN_YEAR..=MYSQL_DATETIME_MAX_YEAR).contains(&year) {
                        return TimestampOverflowSnafu {
                            error: format!(
                                "timestamp {} with unit {} has local year {}, outside MySQL DATETIME range {}..={}",
                                timestamp.value(),
                                timestamp.unit(),
                                year,
                                MYSQL_DATETIME_MIN_YEAR,
                                MYSQL_DATETIME_MAX_YEAR,
                            ),
                        }
                        .fail();
                    }
                    write!(
                        &mut staged_timestamp.formatted,
                        "{}",
                        datetime.format("%Y-%m-%d %H:%M:%S%.f")
                    )
                    .map_err(|_| {
                        InternalSnafu {
                            err_msg: "timestamp formatting failed",
                        }
                        .build()
                    })?;
                    staged_timestamp.datetime = Some(datetime);
                }
            }

            for (j, column) in record_batch.columns().iter().enumerate() {
                if column.is_null(i) {
                    row_writer.write_col(None::<u8>)?;
                    continue;
                }

                match column.data_type() {
                    DataType::Null => {
                        row_writer.write_col(None::<u8>)?;
                    }
                    DataType::Boolean => {
                        let array = column.as_boolean();
                        row_writer.write_col(array.value(i) as i8)?;
                    }
                    DataType::UInt8 => {
                        let array = column.as_primitive::<UInt8Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::UInt16 => {
                        let array = column.as_primitive::<UInt16Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::UInt32 => {
                        let array = column.as_primitive::<UInt32Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::UInt64 => {
                        let array = column.as_primitive::<UInt64Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::Int8 => {
                        let array = column.as_primitive::<Int8Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::Int16 => {
                        let array = column.as_primitive::<Int16Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::Int32 => {
                        let array = column.as_primitive::<Int32Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::Int64 => {
                        let array = column.as_primitive::<Int64Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::Float32 => {
                        let array = column.as_primitive::<Float32Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::Float64 => {
                        let array = column.as_primitive::<Float64Type>();
                        row_writer.write_col(array.value(i))?;
                    }
                    DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => {
                        let v = datatypes::arrow_array::string_array_value(column, i);
                        row_writer.write_col(v)?;
                    }
                    DataType::Binary | DataType::BinaryView | DataType::LargeBinary => {
                        let v = datatypes::arrow_array::binary_array_value(column, i);
                        if let ConcreteDataType::Json(_) = &schema.column_schemas()[j].data_type {
                            let s = jsonb_to_string(v).context(ConvertSqlValueSnafu)?;
                            row_writer.write_col(s)?;
                        } else {
                            row_writer.write_col(v)?;
                        }
                    }
                    DataType::Date32 => {
                        let array = column.as_primitive::<Date32Type>();
                        let v = Date::new(array.value(i));
                        row_writer.write_col(v.to_chrono_date())?;
                    }
                    DataType::Timestamp(_, _) => {
                        let slot = timestamp_slots
                            .get(j)
                            .context(InternalSnafu {
                                err_msg: "timestamp column index is invalid",
                            })?
                            .as_ref()
                            .context(InternalSnafu {
                                err_msg: "timestamp column has no staging slot",
                            })?;
                        let (_, staged_timestamp) =
                            staged_timestamps.get(*slot).context(InternalSnafu {
                                err_msg: "timestamp staging slot is missing",
                            })?;
                        let datetime = staged_timestamp.datetime.context(InternalSnafu {
                            err_msg: "timestamp staging value is missing",
                        })?;
                        row_writer.write_col(PrecisionTimestamp {
                            formatted: staged_timestamp.formatted.as_str(),
                            datetime,
                        })?;
                    }
                    DataType::Interval(interval_unit) => match interval_unit {
                        IntervalUnit::YearMonth => {
                            let array = column.as_primitive::<IntervalYearMonthType>();
                            let v: IntervalYearMonth = array.value(i).into();
                            row_writer.write_col(v.to_iso8601_string())?;
                        }
                        IntervalUnit::DayTime => {
                            let array = column.as_primitive::<IntervalDayTimeType>();
                            let v: IntervalDayTime = array.value(i).into();
                            row_writer.write_col(v.to_iso8601_string())?;
                        }
                        IntervalUnit::MonthDayNano => {
                            let array = column.as_primitive::<IntervalMonthDayNanoType>();
                            let v: IntervalMonthDayNano = array.value(i).into();
                            row_writer.write_col(v.to_iso8601_string())?;
                        }
                    },
                    DataType::Duration(_) => {
                        let v: Duration =
                            datatypes::arrow_array::duration_array_value(column, i).into();
                        row_writer.write_col(v)?;
                    }
                    DataType::List(_) | DataType::Struct(_) => {
                        let v = ScalarValue::try_from_array(column, i).context(DataFusionSnafu)?;
                        row_writer.write_col(v.to_string())?;
                    }
                    DataType::Time32(_) | DataType::Time64(_) => {
                        let time = datatypes::arrow_array::time_array_value(column, i);
                        let v = time.to_timezone_aware_string(Some(&query_context.timezone()));
                        row_writer.write_col(v)?;
                    }
                    DataType::Decimal128(precision, scale) => {
                        let array = column.as_primitive::<Decimal128Type>();
                        let v = Decimal128::new(array.value(i), *precision, *scale);
                        row_writer.write_col(v.to_string())?;
                    }
                    _ => {
                        return NotSupportedSnafu {
                            feat: format!("convert {} to MySQL value", column.data_type()),
                        }
                        .fail();
                    }
                }
            }
            row_writer.end_row().await?;
        }
        Ok(())
    }

    async fn write_query_error<'a, W: AsyncWrite + Unpin>(
        error: impl ErrorExt,
        w: QueryResultWriter<'a, W>,
        query_context: QueryContextRef,
    ) -> io::Result<()> {
        METRIC_ERROR_COUNTER
            .with_label_values(&[METRIC_ERROR_COUNTER_LABEL_MYSQL])
            .inc();

        let (kind, err) = handle_err(error, query_context);
        debug!("Write query error, kind: {:?}, err: {}", kind, err);
        w.error(kind, err.as_bytes()).await?;
        Ok(())
    }
}

pub fn create_mysql_column(data_type: &ConcreteDataType, column_name: &str) -> Result<Column> {
    let data_type = &map_dictionary_to_values_data_type(data_type);
    let column_type = match data_type {
        ConcreteDataType::Null(_) => Ok(ColumnType::MYSQL_TYPE_NULL),
        ConcreteDataType::Boolean(_) | ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => {
            Ok(ColumnType::MYSQL_TYPE_TINY)
        }
        ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => {
            Ok(ColumnType::MYSQL_TYPE_SHORT)
        }
        ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => Ok(ColumnType::MYSQL_TYPE_LONG),
        ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => {
            Ok(ColumnType::MYSQL_TYPE_LONGLONG)
        }
        ConcreteDataType::Float32(_) => Ok(ColumnType::MYSQL_TYPE_FLOAT),
        ConcreteDataType::Float64(_) => Ok(ColumnType::MYSQL_TYPE_DOUBLE),
        ConcreteDataType::Binary(_) | ConcreteDataType::String(_) => {
            Ok(ColumnType::MYSQL_TYPE_VARCHAR)
        }
        ConcreteDataType::Timestamp(_) => Ok(ColumnType::MYSQL_TYPE_TIMESTAMP),
        ConcreteDataType::Time(_) => Ok(ColumnType::MYSQL_TYPE_TIME),
        ConcreteDataType::Date(_) => Ok(ColumnType::MYSQL_TYPE_DATE),
        ConcreteDataType::Interval(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
        ConcreteDataType::Duration(_) => Ok(ColumnType::MYSQL_TYPE_TIME),
        ConcreteDataType::Decimal128(_) => Ok(ColumnType::MYSQL_TYPE_DECIMAL),
        ConcreteDataType::Json(_) => Ok(ColumnType::MYSQL_TYPE_JSON),
        ConcreteDataType::Vector(_) => Ok(ColumnType::MYSQL_TYPE_BLOB),
        ConcreteDataType::List(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
        ConcreteDataType::Struct(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
        _ => error::UnsupportedDataTypeSnafu {
            data_type,
            reason: "not implemented",
        }
        .fail(),
    };
    let mut colflags = ColumnFlags::empty();
    match data_type {
        ConcreteDataType::UInt16(_)
        | ConcreteDataType::UInt8(_)
        | ConcreteDataType::UInt32(_)
        | ConcreteDataType::UInt64(_) => colflags |= ColumnFlags::UNSIGNED_FLAG,
        _ => {}
    };
    column_type.map(|column_type| Column {
        column: column_name.to_string(),
        coltype: column_type,
        // TODO(LFC): Currently "table" and "colflags" are not relevant in MySQL server
        //   implementation, will revisit them again in the future.
        table: String::default(),
        collen: 0, // 0 means "use default".
        colflags,
    })
}

/// Creates MySQL columns definition from our column schema.
pub fn create_mysql_column_def(schema: &SchemaRef) -> Result<Vec<Column>> {
    schema
        .column_schemas()
        .iter()
        .map(|column_schema| create_mysql_column(&column_schema.data_type, &column_schema.name))
        .collect()
}

fn mysql_error_kind(status_code: &StatusCode) -> ErrorKind {
    match status_code {
        StatusCode::Success => ErrorKind::ER_YES,
        StatusCode::Unknown | StatusCode::External => ErrorKind::ER_UNKNOWN_ERROR,
        StatusCode::Unsupported => ErrorKind::ER_NOT_SUPPORTED_YET,
        StatusCode::Cancelled | StatusCode::DeadlineExceeded => ErrorKind::ER_QUERY_INTERRUPTED,
        StatusCode::RuntimeResourcesExhausted => ErrorKind::ER_OUT_OF_RESOURCES,
        StatusCode::InvalidSyntax => ErrorKind::ER_SYNTAX_ERROR,
        StatusCode::RegionAlreadyExists | StatusCode::TableAlreadyExists => {
            ErrorKind::ER_TABLE_EXISTS_ERROR
        }
        StatusCode::RegionNotFound | StatusCode::TableNotFound => ErrorKind::ER_NO_SUCH_TABLE,
        StatusCode::RegionReadonly => ErrorKind::ER_READ_ONLY_MODE,
        StatusCode::DatabaseNotFound => ErrorKind::ER_WRONG_DB_NAME,
        StatusCode::UserNotFound => ErrorKind::ER_NO_SUCH_USER,
        StatusCode::UnsupportedPasswordType => ErrorKind::ER_PASSWORD_FORMAT,
        StatusCode::PermissionDenied | StatusCode::AccessDenied => {
            ErrorKind::ER_ACCESS_DENIED_ERROR
        }
        StatusCode::UserPasswordMismatch => ErrorKind::ER_DBACCESS_DENIED_ERROR,
        StatusCode::InvalidAuthHeader | StatusCode::AuthHeaderNotFound => {
            ErrorKind::ER_NOT_SUPPORTED_AUTH_MODE
        }
        StatusCode::Unexpected
        | StatusCode::Internal
        | StatusCode::IllegalState
        | StatusCode::PlanQuery
        | StatusCode::EngineExecuteQuery
        | StatusCode::RegionNotReady
        | StatusCode::RegionBusy
        | StatusCode::TableUnavailable
        | StatusCode::StorageUnavailable
        | StatusCode::RequestOutdated => ErrorKind::ER_INTERNAL_ERROR,
        StatusCode::InvalidArguments => ErrorKind::ER_WRONG_ARGUMENTS,
        StatusCode::TableColumnNotFound => ErrorKind::ER_BAD_FIELD_ERROR,
        StatusCode::TableColumnExists => ErrorKind::ER_DUP_FIELDNAME,
        StatusCode::DatabaseAlreadyExists => ErrorKind::ER_DB_CREATE_EXISTS,
        StatusCode::RateLimited => ErrorKind::ER_TOO_MANY_CONCURRENT_TRXS,
        StatusCode::FlowAlreadyExists => ErrorKind::ER_TABLE_EXISTS_ERROR,
        StatusCode::FlowNotFound => ErrorKind::ER_NO_SUCH_TABLE,
        StatusCode::TriggerAlreadyExists => ErrorKind::ER_TABLE_EXISTS_ERROR,
        StatusCode::TriggerNotFound => ErrorKind::ER_NO_SUCH_TABLE,
        StatusCode::Suspended => ErrorKind::ER_SERVER_SHUTDOWN,
    }
}
