// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

use std::collections::{BTreeMap, BTreeSet};
use std::sync::LazyLock;

use itertools::Itertools;
use regex::Regex;

use mysql_async::prelude::{FromRow, Queryable};
use mysql_async::{FromRowError, Row};

use mz_repr::adt::char::CharLength;
use mz_repr::adt::numeric::{NumericMaxScale, NUMERIC_DATUM_MAX_PRECISION};
use mz_repr::adt::timestamp::TimestampPrecision;
use mz_repr::adt::varchar::VarCharMaxLength;
use mz_repr::{ColumnType, ScalarType};

use crate::desc::{
    MySqlColumnDesc, MySqlColumnMeta, MySqlColumnMetaEnum, MySqlKeyDesc, MySqlTableDesc,
};
use crate::{MySqlError, UnsupportedDataType};

/// Helper for querying information_schema.columns
// NOTE: The order of these names *must* match the order of fields of the [`InfoSchema`] struct.
const INFO_SCHEMA_COLS: &[&str] = &[
    "column_name",
    "data_type",
    "column_type",
    "is_nullable",
    "numeric_precision",
    "numeric_scale",
    "datetime_precision",
    "character_maximum_length",
];

// NOTE: The order of these fields *must* match the order of names of the [`INFO_SCHEMA_COLS`] list.
#[derive(Debug, Clone)]
pub struct InfoSchema {
    column_name: String,
    data_type: String,
    column_type: String,
    is_nullable: String,
    numeric_precision: Option<i64>,
    numeric_scale: Option<i64>,
    datetime_precision: Option<i64>,
    character_maximum_length: Option<i64>,
}

impl FromRow for InfoSchema {
    fn from_row_opt(row: Row) -> Result<Self, FromRowError> {
        let actual = row.columns_ref().iter().map(|c| c.name_ref());
        let expected = INFO_SCHEMA_COLS.iter().map(|c| c.as_bytes());
        itertools::assert_equal(actual, expected);
        let (a, b, c, d, e, f, g, h) = FromRow::from_row_opt(row)?;
        Ok(Self {
            column_name: a,
            data_type: b,
            column_type: c,
            is_nullable: d,
            numeric_precision: e,
            numeric_scale: f,
            datetime_precision: g,
            character_maximum_length: h,
        })
    }
}

/// A representation of the raw schema info for a table from MySQL
#[derive(Debug, Clone)]
pub struct MySqlTableSchema {
    pub schema_name: String,
    pub name: String,
    pub columns: Vec<InfoSchema>,
    pub keys: BTreeSet<MySqlKeyDesc>,
}

impl MySqlTableSchema {
    pub fn table_ref<'a>(&'a self) -> QualifiedTableRef<'a> {
        QualifiedTableRef {
            schema_name: &self.schema_name,
            table_name: &self.name,
        }
    }

    /// Convert the raw table schema to our MySqlTableDesc representation
    /// using any provided text_columns and ignore_columns
    pub fn to_desc(
        self,
        text_columns: Option<&BTreeSet<&str>>,
        ignore_columns: Option<&BTreeSet<&str>>,
    ) -> Result<MySqlTableDesc, MySqlError> {
        // Verify there are no duplicates in text_columns and ignore_columns
        match (&text_columns, &ignore_columns) {
            (Some(text_cols), Some(ignore_cols)) => {
                let intersection: Vec<_> = text_cols.intersection(ignore_cols).collect();
                if !intersection.is_empty() {
                    Err(MySqlError::DuplicatedColumnNames {
                        qualified_table_name: format!("{:?}.{:?}", self.schema_name, self.name),
                        columns: intersection.iter().map(|s| (*s).to_string()).collect(),
                    })?;
                }
            }
            _ => (),
        };

        let mut columns = Vec::with_capacity(self.columns.len());
        let mut error_cols = vec![];
        for info in self.columns {
            // If this column is designated as a text column and of a supported text-column type
            // treat it as a string and skip type parsing.
            if let Some(text_columns) = &text_columns {
                if text_columns.contains(&info.column_name.as_str()) {
                    match parse_as_text_column(&info, &self.schema_name, &self.name) {
                        Err(err) => error_cols.push(err),
                        Ok((scalar_type, meta)) => columns.push(MySqlColumnDesc {
                            name: info.column_name,
                            column_type: Some(ColumnType {
                                scalar_type,
                                nullable: &info.is_nullable == "YES",
                            }),
                            meta,
                        }),
                    }
                    continue;
                }
            }

            // If this column is ignored, use None for the column type to signal that it should be.
            if let Some(ignore_cols) = &ignore_columns {
                if ignore_cols.contains(&info.column_name.as_str()) {
                    columns.push(MySqlColumnDesc {
                        name: info.column_name,
                        column_type: None,
                        meta: None,
                    });
                    continue;
                }
            }

            // Collect the parsed data types or errors for later reporting.
            match parse_data_type(&info, &self.schema_name, &self.name) {
                Err(err) => error_cols.push(err),
                Ok(scalar_type) => columns.push(MySqlColumnDesc {
                    name: info.column_name,
                    column_type: Some(ColumnType {
                        scalar_type,
                        nullable: &info.is_nullable == "YES",
                    }),
                    meta: None,
                }),
            }
        }
        if error_cols.len() > 0 {
            Err(MySqlError::UnsupportedDataTypes {
                columns: error_cols,
            })?;
        }

        Ok(MySqlTableDesc {
            schema_name: self.schema_name,
            name: self.name,
            columns,
            keys: self.keys,
        })
    }
}

/// Request for table schemas from MySQL
pub enum SchemaRequest<'a> {
    /// Request schemas for all tables in the database
    All,
    /// Request schemas for all tables in the specified schemas/databases
    Schemas(Vec<&'a str>),
    /// Request schemas for all specified tables, specified as (schema_name, table_name)
    Tables(Vec<(&'a str, &'a str)>),
}

/// A reference to a table in a schema/database
#[derive(Debug, Hash, PartialEq, Eq, Clone, Ord, PartialOrd)]
pub struct QualifiedTableRef<'a> {
    pub schema_name: &'a str,
    pub table_name: &'a str,
}

/// Retrieve the tables and column descriptions for tables in the given schemas.
pub async fn schema_info<'a, Q>(
    conn: &mut Q,
    schema_request: &SchemaRequest<'a>,
) -> Result<Vec<MySqlTableSchema>, MySqlError>
where
    Q: Queryable,
{
    let table_rows: Vec<(String, String)> = match schema_request {
        SchemaRequest::All => {
            // Get all tables in non-system schemas.
            // TODO(roshan): Many users create user-defined tables in the `mysql` system schema, since mysql doesn't
            // prevent this. We may want to consider adding a warning for this in the docs, since that schema
            // contains dozens of built-in system tables that we need to filter out.
            let table_q = "SELECT table_name, table_schema
                FROM information_schema.tables
                WHERE table_type = 'BASE TABLE'
                AND table_schema NOT IN ('information_schema', 'performance_schema', 'mysql', 'sys')";
            conn.exec(table_q, ()).await?
        }
        SchemaRequest::Schemas(schemas) => {
            // Get all tables of type 'Base Table' in specified schemas
            if schemas.is_empty() {
                return Ok(vec![]);
            }
            let table_q = format!(
                "SELECT table_name, table_schema
                FROM information_schema.tables
                WHERE table_type = 'BASE TABLE'
                AND table_schema IN ({})",
                schemas.iter().map(|_| "?").join(", ")
            );
            conn.exec(table_q, schemas).await?
        }
        SchemaRequest::Tables(tables) => {
            // Get all specified tables
            if tables.is_empty() {
                return Ok(vec![]);
            }
            let table_q = format!(
                "SELECT table_name, table_schema
                FROM information_schema.tables
                WHERE table_type = 'BASE TABLE'
                AND (table_schema, table_name) IN ({})",
                tables.iter().map(|_| "(?, ?)").join(", ")
            );
            conn.exec(
                table_q,
                tables
                    .iter()
                    .flat_map(|(s, t)| [*s, *t])
                    .collect::<Vec<_>>(),
            )
            .await?
        }
    };

    let mut tables = vec![];
    for (table_name, schema_name) in table_rows {
        // NOTE: It's important that we order by ordinal_position ASC since we rely on this as
        // the ordering in which columns are returned in a row.
        let column_q = format!(
            "SELECT {}
             FROM information_schema.columns
             WHERE table_name = ? AND table_schema = ?
             ORDER BY ordinal_position ASC",
            INFO_SCHEMA_COLS
                .iter()
                .map(|c| format!("{c} AS {c}"))
                .join(", ")
        );
        let column_rows = conn
            .exec::<InfoSchema, _, _>(column_q, (&table_name, &schema_name))
            .await?;

        // Query for primary key and unique constraints that do not contain expressions / functional key parts.
        // When a constraint contains expressions, the column_name field is NULL.
        let index_rows = conn
            .exec::<(String, String), _, _>(
                "SELECT
                    index_name,
                    column_name
                FROM information_schema.statistics AS outt
                WHERE
                    table_schema NOT IN ('information_schema', 'performance_schema', 'mysql', 'sys')
                    AND NOT EXISTS (
                        SELECT 1
                        FROM information_schema.statistics AS inn
                        WHERE outt.index_name = inn.index_name AND inn.column_name IS NULL
                    )
                    AND non_unique = 0
                    AND table_name = ?
                    AND table_schema = ?
                ORDER BY index_name, seq_in_index
            ",
                (&table_name, &schema_name),
            )
            .await?;

        let mut indices = BTreeMap::new();
        for (index_name, column) in index_rows {
            indices
                .entry(index_name)
                .or_insert_with(Vec::new)
                .push(column);
        }
        let mut keys = BTreeSet::new();
        while let Some((index_name, columns)) = indices.pop_first() {
            keys.insert(MySqlKeyDesc {
                is_primary: &index_name == "PRIMARY",
                name: index_name,
                columns,
            });
        }

        tables.push(MySqlTableSchema {
            schema_name,
            name: table_name,
            columns: column_rows,
            keys,
        });
    }

    Ok(tables)
}

fn parse_data_type(
    info: &InfoSchema,
    schema_name: &str,
    table_name: &str,
) -> Result<ScalarType, UnsupportedDataType> {
    let unsigned = info.column_type.contains("unsigned");

    match info.data_type.as_str() {
        "tinyint" | "smallint" => {
            if unsigned {
                Ok(ScalarType::UInt16)
            } else {
                Ok(ScalarType::Int16)
            }
        }
        "mediumint" | "int" => {
            if unsigned {
                Ok(ScalarType::UInt32)
            } else {
                Ok(ScalarType::Int32)
            }
        }
        "bigint" => {
            if unsigned {
                Ok(ScalarType::UInt64)
            } else {
                Ok(ScalarType::Int64)
            }
        }
        "float" => Ok(ScalarType::Float32),
        "double" => Ok(ScalarType::Float64),
        "date" => Ok(ScalarType::Date),
        "datetime" | "timestamp" => Ok(ScalarType::Timestamp {
            // both mysql and our scalar type use a max six-digit fractional-second precision
            // this is bounds-checked in the TryFrom impl
            precision: info
                .datetime_precision
                .map(TimestampPrecision::try_from)
                .transpose()
                .map_err(|_| UnsupportedDataType {
                    column_type: info.column_type.clone(),
                    qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
                    column_name: info.column_name.clone(),
                    intended_type: None,
                })?,
        }),
        "time" => Ok(ScalarType::Time),
        "decimal" | "numeric" => {
            // validate the precision is within the bounds of our numeric type
            // here since we don't use this precision on the ScalarType itself
            // whereas the scale will be bounds-checked in the TryFrom impl
            if info.numeric_precision.unwrap_or_default() > NUMERIC_DATUM_MAX_PRECISION.into() {
                Err(UnsupportedDataType {
                    column_type: info.column_type.clone(),
                    qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
                    column_name: info.column_name.clone(),
                    intended_type: None,
                })?
            }
            Ok(ScalarType::Numeric {
                max_scale: info
                    .numeric_scale
                    .map(NumericMaxScale::try_from)
                    .transpose()
                    .map_err(|_| UnsupportedDataType {
                        column_type: info.column_type.clone(),
                        qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
                        column_name: info.column_name.clone(),
                        intended_type: None,
                    })?,
            })
        }
        "char" => Ok(ScalarType::Char {
            length: info
                .character_maximum_length
                .and_then(|f| Some(CharLength::try_from(f)))
                .transpose()
                .map_err(|_| UnsupportedDataType {
                    column_type: info.column_type.clone(),
                    qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
                    column_name: info.column_name.clone(),
                    intended_type: None,
                })?,
        }),
        "varchar" => Ok(ScalarType::VarChar {
            max_length: info
                .character_maximum_length
                .and_then(|f| Some(VarCharMaxLength::try_from(f)))
                .transpose()
                .map_err(|_| UnsupportedDataType {
                    column_type: info.column_type.clone(),
                    qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
                    column_name: info.column_name.clone(),
                    intended_type: None,
                })?,
        }),
        "text" | "tinytext" | "mediumtext" | "longtext" => Ok(ScalarType::String),
        "binary" | "varbinary" | "tinyblob" | "blob" | "mediumblob" | "longblob" => {
            Ok(ScalarType::Bytes)
        }
        "json" => Ok(ScalarType::Jsonb),
        _ => Err(UnsupportedDataType {
            column_type: info.column_type.clone(),
            qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
            column_name: info.column_name.clone(),
            intended_type: None,
        }),
    }
}

/// Parse the specified column as a TEXT COLUMN. We only support the set of types that are
/// represented as an encoded-string in both the mysql-common binary query response and binlog
/// event representation, OR types that we've added explicit casting support for.
fn parse_as_text_column(
    info: &InfoSchema,
    schema_name: &str,
    table_name: &str,
) -> Result<(ScalarType, Option<MySqlColumnMeta>), UnsupportedDataType> {
    match info.data_type.as_str() {
        "year" => Ok((ScalarType::String, Some(MySqlColumnMeta::Year))),
        "json" => Ok((ScalarType::String, Some(MySqlColumnMeta::Json))),
        "enum" => Ok((
            ScalarType::String,
            Some(MySqlColumnMeta::Enum(MySqlColumnMetaEnum {
                values: enum_vals_from_column_type(info.column_type.as_str()).map_err(|_| {
                    UnsupportedDataType {
                        column_type: info.column_type.clone(),
                        qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
                        column_name: info.column_name.clone(),
                        intended_type: Some("text".to_string()),
                    }
                })?,
            })),
        )),
        "date" => Ok((ScalarType::String, Some(MySqlColumnMeta::Date))),
        "datetime" | "timestamp" => Ok((
            ScalarType::String,
            Some(MySqlColumnMeta::Timestamp(
                info.datetime_precision
                    // Default precision is 0 in MySQL if not specified
                    .unwrap_or_default()
                    .try_into()
                    .map_err(|_| UnsupportedDataType {
                        column_type: info.column_type.clone(),
                        qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
                        column_name: info.column_name.clone(),
                        intended_type: Some("text".to_string()),
                    })?,
            )),
        )),
        _ => Err(UnsupportedDataType {
            column_type: info.column_type.clone(),
            qualified_table_name: format!("{:?}.{:?}", schema_name, table_name),
            column_name: info.column_name.clone(),
            intended_type: Some("text".to_string()),
        }),
    }
}

static ENUM_VAL_REGEX: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"'((?:[^']|'')*)'").expect("valid regex"));

/// Parse the enum values from a column_type value on an enum column, which is a string like
/// "enum('apple','banana','cher,ry','ora''nge')"
/// We need to handle the case where the enum value itself contains a comma or a
/// single quote (escaped with another quote), so we use a regex to do so
fn enum_vals_from_column_type(s: &str) -> Result<Vec<String>, anyhow::Error> {
    let vals_str = s
        .strip_prefix("enum(")
        .and_then(|s| s.strip_suffix(')'))
        .ok_or(anyhow::format_err!(
            "Unable to parse enum column type string"
        ))?;

    Ok(ENUM_VAL_REGEX
        .captures_iter(vals_str)
        .map(|s| s[1].replace("''", "'"))
        .collect())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[mz_ore::test]
    fn test_enum_value_parsing() {
        let vals =
            enum_vals_from_column_type("enum('apple','banana','cher,ry','ora''nge')").unwrap();
        assert_eq!(vals, vec!["apple", "banana", "cher,ry", "ora'nge"]);
    }
}