// 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::collections::HashMap;

use datatypes::data_type::ConcreteDataType;
use datatypes::value::Value;
use derive_builder::Builder;
use rand::Rng;
use rand::seq::SliceRandom;
use snafu::{ResultExt, ensure};

use super::Generator;
use crate::context::TableContextRef;
use crate::error::{self, Error, Result};
use crate::fake::{MappedGenerator, WordGenerator, random_capitalize_map};
use crate::generator::{ColumnOptionGenerator, ConcreteDataTypeGenerator, Random};
use crate::ir::create_expr::{
    ColumnOption, CreateDatabaseExprBuilder, CreateTableExprBuilder, PartitionDef,
};
use crate::ir::partition_expr::SimplePartitions;
use crate::ir::{
    Column, ColumnTypeGenerator, CreateDatabaseExpr, CreateTableExpr, Ident,
    PartibleColumnTypeGenerator, StringColumnTypeGenerator, TsColumnTypeGenerator,
    column_options_generator, generate_columns, generate_partition_bounds,
    partible_column_options_generator, primary_key_options_generator, ts_column_options_generator,
};

#[derive(Builder)]
#[builder(default, pattern = "owned")]
pub struct CreateTableExprGenerator<R: Rng + 'static> {
    columns: usize,
    #[builder(setter(into))]
    engine: String,
    partition: usize,
    partition_column: bool,
    if_not_exists: bool,
    #[builder(setter(into))]
    name: Ident,
    #[builder(setter(into))]
    with_clause: HashMap<String, String>,
    name_generator: Box<dyn Random<Ident, R>>,
    ts_column_type_generator: ConcreteDataTypeGenerator<R>,
    column_type_generator: ConcreteDataTypeGenerator<R>,
    partible_column_type_generator: ConcreteDataTypeGenerator<R>,
    partible_column_options_generator: ColumnOptionGenerator<R>,
    column_options_generator: ColumnOptionGenerator<R>,
    ts_column_options_generator: ColumnOptionGenerator<R>,
}

const DEFAULT_ENGINE: &str = "mito";

impl<R: Rng + 'static> Default for CreateTableExprGenerator<R> {
    fn default() -> Self {
        Self {
            columns: 0,
            engine: DEFAULT_ENGINE.to_string(),
            if_not_exists: false,
            partition: 0,
            partition_column: false,
            name: Ident::new(""),
            with_clause: HashMap::default(),
            name_generator: Box::new(MappedGenerator::new(WordGenerator, random_capitalize_map)),
            ts_column_type_generator: Box::new(TsColumnTypeGenerator),
            column_type_generator: Box::new(ColumnTypeGenerator),
            partible_column_type_generator: Box::new(PartibleColumnTypeGenerator),
            partible_column_options_generator: Box::new(partible_column_options_generator),
            column_options_generator: Box::new(column_options_generator),
            ts_column_options_generator: Box::new(ts_column_options_generator),
        }
    }
}

impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateTableExprGenerator<R> {
    type Error = Error;

    /// Generates the [CreateTableExpr].
    fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
        ensure!(
            self.columns != 0,
            error::UnexpectedSnafu {
                violated: "The columns must larger than zero"
            }
        );

        let mut builder = CreateTableExprBuilder::default();
        let mut columns = Vec::with_capacity(self.columns);
        let mut primary_keys = vec![];
        let need_partible_column = self.partition > 1 || self.partition_column;
        let mut column_names = self.name_generator.choose(rng, self.columns);

        if self.columns == 1 {
            // Generates the ts column.
            // Safety: columns must large than 0.
            let name = column_names.pop().unwrap();
            let column = generate_columns(
                rng,
                vec![name.clone()],
                self.ts_column_type_generator.as_ref(),
                self.ts_column_options_generator.as_ref(),
            )
            .remove(0);
            columns.push(column);
        } else {
            // Generates the partible column.
            if need_partible_column {
                // Safety: columns must large than 0.
                let name = column_names.pop().unwrap();
                let column = generate_columns(
                    rng,
                    vec![name.clone()],
                    self.partible_column_type_generator.as_ref(),
                    self.partible_column_options_generator.as_ref(),
                )
                .remove(0);

                if self.partition > 1 {
                    // Generates partition bounds.
                    let partition_def = generate_partition_def(
                        self.partition,
                        column.column_type.clone(),
                        name.clone(),
                    );
                    builder.partition(partition_def);
                }
                columns.push(column);
            }
            // Generates the ts column.
            // Safety: columns must large than 1.
            let name = column_names.pop().unwrap();
            columns.extend(generate_columns(
                rng,
                vec![name],
                self.ts_column_type_generator.as_ref(),
                self.ts_column_options_generator.as_ref(),
            ));
            // Generates rest columns
            columns.extend(generate_columns(
                rng,
                column_names,
                self.column_type_generator.as_ref(),
                self.column_options_generator.as_ref(),
            ));
        }

        for (idx, column) in columns.iter().enumerate() {
            if column.is_primary_key() {
                primary_keys.push(idx);
            }
        }
        // Shuffles the primary keys.
        primary_keys.shuffle(rng);

        builder.columns(columns);
        builder.primary_keys(primary_keys);
        builder.engine(self.engine.clone());
        builder.if_not_exists(self.if_not_exists);
        if self.name.is_empty() {
            builder.table_name(self.name_generator.generate(rng));
        } else {
            builder.table_name(self.name.clone());
        }
        if !self.with_clause.is_empty() {
            let mut options = HashMap::new();
            for (key, value) in &self.with_clause {
                options.insert(key.clone(), Value::from(value.clone()));
            }
            builder.options(options);
        }
        builder.build().context(error::BuildCreateTableExprSnafu)
    }
}

pub fn generate_partition_def(
    partitions: usize,
    column_type: ConcreteDataType,
    column_name: Ident,
) -> PartitionDef {
    assert!(partitions > 1, "partitions must be greater than 1");
    let bounds = generate_partition_bounds(&column_type, partitions - 1);
    let partitions = SimplePartitions::new(column_name.clone(), bounds);
    let partition_exprs = partitions.generate().unwrap();

    PartitionDef {
        columns: vec![column_name.clone()],
        exprs: partition_exprs,
    }
}

fn metric_partition_column() -> Column {
    Column {
        name: Ident::new("host"),
        column_type: ConcreteDataType::string_datatype(),
        options: vec![ColumnOption::PrimaryKey],
    }
}

pub fn generate_metric_partition_def(partitions: usize) -> PartitionDef {
    assert!(partitions > 1, "partitions must be greater than 1");
    let partition_column = metric_partition_column();
    let bounds = generate_partition_bounds(&partition_column.column_type, partitions - 1);
    let partitions = SimplePartitions::new(partition_column.name.clone(), bounds);
    PartitionDef {
        columns: vec![partitions.column_name.clone()],
        exprs: partitions.generate().unwrap(),
    }
}

/// Generate a physical table with 2 columns: ts of TimestampType::Millisecond as time index and val of Float64Type.
#[derive(Builder)]
#[builder(pattern = "owned")]
pub struct CreatePhysicalTableExprGenerator<R: Rng + 'static> {
    #[builder(default = "Box::new(WordGenerator)")]
    name_generator: Box<dyn Random<Ident, R>>,
    #[builder(default = "false")]
    if_not_exists: bool,
    #[builder(default = "0")]
    partition: usize,
    #[builder(default = "false")]
    partition_column: bool,
    #[builder(default, setter(into))]
    with_clause: HashMap<String, String>,
}

impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreatePhysicalTableExprGenerator<R> {
    type Error = Error;

    fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
        let mut options = HashMap::with_capacity(self.with_clause.len() + 1);
        options.insert("physical_metric_table".to_string(), Value::from(""));
        for (key, value) in &self.with_clause {
            options.insert(key.clone(), Value::from(value.clone()));
        }

        let mut columns = vec![
            Column {
                name: Ident::new("ts"),
                column_type: ConcreteDataType::timestamp_millisecond_datatype(),
                options: vec![ColumnOption::TimeIndex],
            },
            Column {
                name: Ident::new("val"),
                column_type: ConcreteDataType::float64_datatype(),
                options: vec![],
            },
        ];

        let mut partition = None;
        let mut primary_keys = vec![];
        if self.partition > 1 || self.partition_column {
            columns.push(metric_partition_column());
            primary_keys.push(columns.len() - 1);
        }
        if self.partition > 1 {
            partition = Some(generate_metric_partition_def(self.partition));
        }

        Ok(CreateTableExpr {
            table_name: self.name_generator.generate(rng),
            columns,
            if_not_exists: self.if_not_exists,
            partition,
            engine: "metric".to_string(),
            options,
            primary_keys,
        })
    }
}

/// Generate a logical table based on an existing physical table.
#[derive(Builder)]
#[builder(pattern = "owned")]
pub struct CreateLogicalTableExprGenerator<R: Rng + 'static> {
    physical_table_ctx: TableContextRef,
    labels: usize,
    if_not_exists: bool,
    #[builder(default = "true")]
    include_partition_column: bool,
    #[builder(default = "Box::new(WordGenerator)")]
    name_generator: Box<dyn Random<Ident, R>>,
}

impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateLogicalTableExprGenerator<R> {
    type Error = Error;

    fn generate(&self, rng: &mut R) -> Result<CreateTableExpr> {
        // Currently we mock the usage of GreptimeDB as Prometheus' backend, the physical table must have ts and val.
        ensure!(
            self.physical_table_ctx.columns.len() >= 2,
            error::UnexpectedSnafu {
                violated: "The physical table must have at least two columns"
            }
        );

        // Generates the logical table columns based on the physical table.
        let logical_table_name = self
            .physical_table_ctx
            .generate_unique_table_name(rng, self.name_generator.as_ref());
        let mut physical_columns = self.physical_table_ctx.columns.clone();
        if !self.include_partition_column
            && let Some(partition_def) = &self.physical_table_ctx.partition
        {
            physical_columns.retain(|column| !partition_def.columns.contains(&column.name));
        }

        let mut logical_table = CreateTableExpr {
            table_name: logical_table_name,
            columns: physical_columns,
            if_not_exists: self.if_not_exists,
            partition: None,
            engine: "metric".to_string(),
            options: [(
                "on_physical_table".to_string(),
                self.physical_table_ctx.name.value.clone().into(),
            )]
            .into(),
            primary_keys: vec![],
        };

        let column_names = self.name_generator.choose(rng, self.labels);
        logical_table.columns.extend(generate_columns(
            rng,
            column_names,
            &StringColumnTypeGenerator,
            Box::new(primary_key_options_generator),
        ));

        // Currently only the `primary key` option is kept in physical table,
        // so we only keep the `primary key` option in the logical table for fuzz test.
        let mut primary_keys = vec![];
        for (idx, column) in logical_table.columns.iter().enumerate() {
            if column.is_primary_key() {
                primary_keys.push(idx);
            }
        }
        primary_keys.shuffle(rng);
        logical_table.primary_keys = primary_keys;

        Ok(logical_table)
    }
}

#[derive(Builder)]
#[builder(default, pattern = "owned")]
pub struct CreateDatabaseExprGenerator<R: Rng + 'static> {
    #[builder(setter(into))]
    database_name: String,
    name_generator: Box<dyn Random<Ident, R>>,
    if_not_exists: bool,
}

impl<R: Rng + 'static> Default for CreateDatabaseExprGenerator<R> {
    fn default() -> Self {
        Self {
            database_name: String::new(),
            name_generator: Box::new(MappedGenerator::new(WordGenerator, random_capitalize_map)),
            if_not_exists: false,
        }
    }
}

impl<R: Rng + 'static> Generator<CreateDatabaseExpr, R> for CreateDatabaseExprGenerator<R> {
    type Error = Error;

    fn generate(&self, rng: &mut R) -> Result<CreateDatabaseExpr> {
        let mut builder = CreateDatabaseExprBuilder::default();
        builder.if_not_exists(self.if_not_exists);
        if self.database_name.is_empty() {
            builder.database_name(self.name_generator.generate(rng));
        } else {
            builder.database_name(self.database_name.clone());
        }
        builder.build().context(error::BuildCreateDatabaseExprSnafu)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use datatypes::data_type::ConcreteDataType;
    use datatypes::value::Value;
    use rand::SeedableRng;

    use super::*;
    use crate::context::TableContext;
    use crate::ir::PARTIBLE_DATA_TYPES;

    #[test]
    fn test_float64() {
        let value = Value::from(0.047318541668048164);
        assert_eq!("0.047318541668048164", value.to_string());
        let value: f64 = "0.047318541668048164".parse().unwrap();
        assert_eq!("0.047318541668048164", value.to_string());
    }

    #[test]
    fn test_create_table_expr_generator() {
        let mut rng = rand::rng();

        let expr = CreateTableExprGeneratorBuilder::default()
            .columns(10)
            .partition(3)
            .if_not_exists(true)
            .engine("mito2")
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();
        assert_eq!(expr.engine, "mito2");
        assert!(expr.if_not_exists);
        assert_eq!(expr.columns.len(), 10);
        assert_eq!(expr.partition.unwrap().exprs.len(), 3);

        let expr = CreateTableExprGeneratorBuilder::default()
            .columns(10)
            .partition(1)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();
        assert_eq!(expr.columns.len(), 10);
        assert!(expr.partition.is_none());

        let expr = CreateTableExprGeneratorBuilder::default()
            .columns(10)
            .partition(1)
            .partition_column(true)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();
        assert_eq!(expr.columns.len(), 10);
        assert!(expr.partition.is_none());
        assert!(PARTIBLE_DATA_TYPES.contains(&expr.columns[0].column_type));
    }

    #[test]
    fn test_create_table_expr_generator_deterministic() {
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
        let expr = CreateTableExprGeneratorBuilder::default()
            .columns(10)
            .partition(3)
            .if_not_exists(true)
            .engine("mito2")
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();

        let serialized = serde_json::to_string(&expr).unwrap();
        let expected = r#"{"table_name":{"value":"quasi","quote_style":null},"columns":[{"name":{"value":"mOLEsTIAs","quote_style":null},"column_type":{"Float64":{}},"options":["PrimaryKey","Null"]},{"name":{"value":"CUMQUe","quote_style":null},"column_type":{"Timestamp":{"Second":null}},"options":["TimeIndex"]},{"name":{"value":"NaTus","quote_style":null},"column_type":{"Int64":{}},"options":[]},{"name":{"value":"EXPeDITA","quote_style":null},"column_type":{"Float64":{}},"options":[]},{"name":{"value":"ImPEDiT","quote_style":null},"column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.56425774}}]},{"name":{"value":"ADIpisci","quote_style":null},"column_type":{"Float32":{}},"options":["PrimaryKey"]},{"name":{"value":"deBITIs","quote_style":null},"column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.31315368}}]},{"name":{"value":"toTaM","quote_style":null},"column_type":{"Int32":{}},"options":["NotNull"]},{"name":{"value":"QuI","quote_style":null},"column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.39941502}}]},{"name":{"value":"INVeNtOre","quote_style":null},"column_type":{"Boolean":null},"options":["PrimaryKey"]}],"if_not_exists":true,"partition":{"columns":[{"value":"mOLEsTIAs","quote_style":null}],"exprs":[{"lhs":{"Column":"mOLEsTIAs"},"op":"Lt","rhs":{"Value":{"Float64":5.992310449541053e+307}}},{"lhs":{"Expr":{"lhs":{"Column":"mOLEsTIAs"},"op":"GtEq","rhs":{"Value":{"Float64":5.992310449541053e+307}}}},"op":"And","rhs":{"Expr":{"lhs":{"Column":"mOLEsTIAs"},"op":"Lt","rhs":{"Value":{"Float64":1.1984620899082105e+308}}}}},{"lhs":{"Column":"mOLEsTIAs"},"op":"GtEq","rhs":{"Value":{"Float64":1.1984620899082105e+308}}}]},"engine":"mito2","options":{},"primary_keys":[0,5,9]}"#;
        assert_eq!(expected, serialized);
    }

    #[test]
    fn test_create_logical_table_expr_generator() {
        let mut rng = rand::rng();

        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
            .if_not_exists(false)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();
        assert_eq!(physical_table_expr.engine, "metric");
        assert_eq!(physical_table_expr.columns.len(), 2);

        let physical_ts = physical_table_expr.columns.iter().position(|column| {
            column
                .options
                .iter()
                .any(|option| option == &ColumnOption::TimeIndex)
        });
        let physical_ts_name = physical_table_expr.columns[physical_ts.unwrap()]
            .name
            .value
            .clone();

        let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));

        let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
            .physical_table_ctx(physical_table_ctx)
            .labels(5)
            .if_not_exists(false)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();
        let logical_ts = logical_table_expr.columns.iter().position(|column| {
            column
                .options
                .iter()
                .any(|option| option == &ColumnOption::TimeIndex)
        });
        let logical_ts_name = logical_table_expr.columns[logical_ts.unwrap()]
            .name
            .value
            .clone();

        assert_eq!(logical_table_expr.engine, "metric");
        assert_eq!(logical_table_expr.columns.len(), 7);
        assert_eq!(logical_ts_name, physical_ts_name);
        assert!(logical_table_expr.columns.iter().all(|column| {
            column.column_type != ConcreteDataType::string_datatype()
                || column
                    .options
                    .iter()
                    .any(|option| option == &ColumnOption::PrimaryKey)
        }));
    }

    #[test]
    fn test_create_physical_table_expr_generator_with_partition() {
        let mut rng = rand::rng();
        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
            .partition(3)
            .if_not_exists(false)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();

        assert_eq!(physical_table_expr.engine, "metric");
        assert!(physical_table_expr.partition.is_some());
        assert_eq!(physical_table_expr.partition.unwrap().exprs.len(), 3);
    }

    #[test]
    fn test_create_physical_table_expr_generator_with_partition_column() {
        let mut rng = rand::rng();
        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
            .partition(1)
            .partition_column(true)
            .if_not_exists(false)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();

        assert_eq!(physical_table_expr.engine, "metric");
        assert!(physical_table_expr.partition.is_none());
        assert_eq!(physical_table_expr.columns.len(), 3);
        assert_eq!(physical_table_expr.columns[2].name, Ident::new("host"));
        assert_eq!(physical_table_expr.primary_keys, vec![2]);
    }

    #[test]
    fn test_create_logical_table_expr_generator_without_partition_column() {
        let mut rng = rand::rng();
        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
            .partition(3)
            .if_not_exists(false)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();
        let partition_columns = physical_table_expr
            .partition
            .as_ref()
            .unwrap()
            .columns
            .clone();
        let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));

        let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
            .physical_table_ctx(physical_table_ctx)
            .labels(3)
            .include_partition_column(false)
            .if_not_exists(false)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();

        assert!(
            logical_table_expr
                .columns
                .iter()
                .all(|column| !partition_columns.contains(&column.name))
        );
    }

    #[test]
    fn test_create_logical_table_expr_generator_deterministic() {
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
        let physical_table_expr = CreatePhysicalTableExprGeneratorBuilder::default()
            .if_not_exists(false)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();
        let physical_table_serialized = serde_json::to_string(&physical_table_expr).unwrap();
        let physical_table_expected = r#"{"table_name":{"value":"expedita","quote_style":null},"columns":[{"name":{"value":"ts","quote_style":null},"column_type":{"Timestamp":{"Millisecond":null}},"options":["TimeIndex"]},{"name":{"value":"val","quote_style":null},"column_type":{"Float64":{}},"options":[]}],"if_not_exists":false,"partition":null,"engine":"metric","options":{"physical_metric_table":{"String":""}},"primary_keys":[]}"#;
        assert_eq!(physical_table_expected, physical_table_serialized);

        let physical_table_ctx = Arc::new(TableContext::from(&physical_table_expr));

        let logical_table_expr = CreateLogicalTableExprGeneratorBuilder::default()
            .physical_table_ctx(physical_table_ctx)
            .labels(5)
            .if_not_exists(false)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();

        let logical_table_serialized = serde_json::to_string(&logical_table_expr).unwrap();
        let logical_table_expected = r#"{"table_name":{"value":"impedit","quote_style":null},"columns":[{"name":{"value":"ts","quote_style":null},"column_type":{"Timestamp":{"Millisecond":null}},"options":["TimeIndex"]},{"name":{"value":"val","quote_style":null},"column_type":{"Float64":{}},"options":[]},{"name":{"value":"totam","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]},{"name":{"value":"cumque","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]},{"name":{"value":"natus","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]},{"name":{"value":"molestias","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]},{"name":{"value":"qui","quote_style":null},"column_type":{"String":{"size_type":"Utf8"}},"options":["PrimaryKey"]}],"if_not_exists":false,"partition":null,"engine":"metric","options":{"on_physical_table":{"String":"expedita"}},"primary_keys":[4,2,3,6,5]}"#;
        assert_eq!(logical_table_expected, logical_table_serialized);
    }

    #[test]
    fn test_create_database_expr_generator() {
        let mut rng = rand::rng();

        let expr = CreateDatabaseExprGeneratorBuilder::default()
            .if_not_exists(true)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();
        assert!(expr.if_not_exists);
    }

    #[test]
    fn test_create_database_expr_generator_deterministic() {
        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
        let expr = CreateDatabaseExprGeneratorBuilder::default()
            .if_not_exists(true)
            .build()
            .unwrap()
            .generate(&mut rng)
            .unwrap();

        let serialized = serde_json::to_string(&expr).unwrap();
        let expected =
            r#"{"database_name":{"value":"EXPediTA","quote_style":null},"if_not_exists":true}"#;
        assert_eq!(expected, serialized);
    }
}
