// 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.

mod column_schema;
pub mod constraint;
pub mod ext;

use std::collections::HashMap;
use std::sync::Arc;
use std::{fmt, mem};

use arrow::datatypes::{Field, Schema as ArrowSchema};
use datafusion_common::DFSchemaRef;
use serde::{Deserialize, Deserializer, Serialize};
use snafu::{ResultExt, ensure};

use crate::error::{self, DuplicateColumnSnafu, Error, ProjectArrowSchemaSnafu, Result};
use crate::prelude::ConcreteDataType;
pub use crate::schema::column_schema::{
    COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE, COLUMN_FULLTEXT_OPT_KEY_ANALYZER,
    COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE,
    COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY,
    COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY,
    COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COLUMN_VECTOR_INDEX_OPT_KEY_CONNECTIVITY,
    COLUMN_VECTOR_INDEX_OPT_KEY_ENGINE, COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_ADD,
    COLUMN_VECTOR_INDEX_OPT_KEY_EXPANSION_SEARCH, COLUMN_VECTOR_INDEX_OPT_KEY_METRIC, COMMENT_KEY,
    ColumnExtType, ColumnSchema, FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions,
    INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType,
    TIME_INDEX_KEY, VECTOR_INDEX_KEY, VectorDistanceMetric, VectorIndexEngineType,
    VectorIndexOptions,
};
pub use crate::schema::constraint::ColumnDefaultConstraint;

/// Key used to store version number of the schema in metadata.
pub const VERSION_KEY: &str = "greptime:version";
/// Key used to store actual column type in field metadata.
pub const TYPE_KEY: &str = "greptime:type";

/// A common schema, should be immutable.
#[derive(Clone, PartialEq, Eq, Serialize)]
pub struct Schema {
    column_schemas: Vec<ColumnSchema>,
    #[serde(skip_serializing)]
    name_to_index: HashMap<String, usize>,
    #[serde(skip_serializing)]
    arrow_schema: Arc<ArrowSchema>,
    /// Index of the timestamp key column.
    ///
    /// Timestamp key column is the column holds the timestamp and forms part of
    /// the primary key. None means there is no timestamp key column.
    #[serde(skip_serializing)]
    timestamp_index: Option<usize>,
    /// Version of the schema.
    ///
    /// Initial value is zero. The version should bump after altering schema.
    version: u32,
}

impl<'de> Deserialize<'de> for Schema {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        use serde::de::Error;

        #[derive(Deserialize)]
        struct RawSchema {
            column_schemas: Vec<ColumnSchema>,
            version: u32,
        }
        let raw = RawSchema::deserialize(deserializer)?;

        SchemaBuilder::try_from(raw.column_schemas)
            .and_then(|x| x.version(raw.version).build())
            .map_err(D::Error::custom)
    }
}

impl fmt::Debug for Schema {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Schema")
            .field("column_schemas", &self.column_schemas)
            .field("name_to_index", &self.name_to_index)
            .field("timestamp_index", &self.timestamp_index)
            .field("version", &self.version)
            .finish()
    }
}

impl Schema {
    /// Initial version of the schema.
    pub const INITIAL_VERSION: u32 = 0;

    /// Create a schema from a vector of [ColumnSchema].
    ///
    /// # Panics
    /// Panics when ColumnSchema's `default_constraint` can't be serialized into json.
    pub fn new(column_schemas: Vec<ColumnSchema>) -> Schema {
        // Builder won't fail in this case
        Self::new_with_version(column_schemas, Self::INITIAL_VERSION)
    }

    pub fn new_with_version(column_schemas: Vec<ColumnSchema>, version: u32) -> Schema {
        SchemaBuilder::try_from(column_schemas)
            .unwrap()
            .version(version)
            .build()
            .unwrap()
    }

    /// Try to Create a schema from a vector of [ColumnSchema].
    pub fn try_new(column_schemas: Vec<ColumnSchema>) -> Result<Schema> {
        SchemaBuilder::try_from(column_schemas)?.build()
    }

    pub fn arrow_schema(&self) -> &Arc<ArrowSchema> {
        &self.arrow_schema
    }

    pub fn column_schemas(&self) -> &[ColumnSchema] {
        &self.column_schemas
    }

    pub fn column_schema_by_name(&self, name: &str) -> Option<&ColumnSchema> {
        self.name_to_index
            .get(name)
            .map(|index| &self.column_schemas[*index])
    }

    /// Retrieve the column's name by index
    /// # Panics
    /// This method **may** panic if the index is out of range of column schemas.
    pub fn column_name_by_index(&self, idx: usize) -> &str {
        &self.column_schemas[idx].name
    }

    pub fn column_index_by_name(&self, name: &str) -> Option<usize> {
        self.name_to_index.get(name).copied()
    }

    pub fn contains_column(&self, name: &str) -> bool {
        self.name_to_index.contains_key(name)
    }

    pub fn num_columns(&self) -> usize {
        self.column_schemas.len()
    }

    pub fn is_empty(&self) -> bool {
        self.column_schemas.is_empty()
    }

    /// Returns index of the timestamp key column.
    pub fn timestamp_index(&self) -> Option<usize> {
        self.timestamp_index
    }

    pub fn timestamp_column(&self) -> Option<&ColumnSchema> {
        self.timestamp_index.map(|idx| &self.column_schemas[idx])
    }

    pub fn version(&self) -> u32 {
        self.version
    }

    pub fn metadata(&self) -> &HashMap<String, String> {
        &self.arrow_schema.metadata
    }

    /// Returns the estimated memory footprint of this schema.
    pub fn estimated_size(&self) -> usize {
        mem::size_of_val(self)
            + mem::size_of::<ColumnSchema>() * self.column_schemas.capacity()
            + self
                .column_schemas
                .iter()
                .map(|column_schema| {
                    column_schema.estimated_size() - mem::size_of::<ColumnSchema>()
                })
                .sum::<usize>()
            + mem::size_of::<(String, usize)>() * self.name_to_index.capacity()
            + self
                .name_to_index
                .keys()
                .map(|name| name.capacity())
                .sum::<usize>()
            + arrow_schema_size(self.arrow_schema.as_ref())
    }

    /// Generate a new projected schema
    ///
    /// # Panic
    ///
    /// If the index out ouf bound
    pub fn try_project(&self, indices: &[usize]) -> Result<Self> {
        let mut column_schemas = Vec::with_capacity(indices.len());
        let mut timestamp_index = None;
        for index in indices {
            if let Some(ts_index) = self.timestamp_index
                && ts_index == *index
            {
                timestamp_index = Some(column_schemas.len());
            }
            column_schemas.push(self.column_schemas[*index].clone());
        }
        let arrow_schema = self
            .arrow_schema
            .project(indices)
            .context(ProjectArrowSchemaSnafu)?;
        let name_to_index = column_schemas
            .iter()
            .enumerate()
            .map(|(pos, column_schema)| (column_schema.name.clone(), pos))
            .collect();

        Ok(Self {
            column_schemas,
            name_to_index,
            arrow_schema: Arc::new(arrow_schema),
            timestamp_index,
            version: self.version,
        })
    }
}

fn arrow_schema_size(schema: &ArrowSchema) -> usize {
    mem::size_of_val(schema)
        + schema.fields.size()
        + mem::size_of::<(String, String)>() * schema.metadata.capacity()
        + schema
            .metadata
            .iter()
            .map(|(key, value)| key.capacity() + value.capacity())
            .sum::<usize>()
}

#[derive(Default)]
pub struct SchemaBuilder {
    column_schemas: Vec<ColumnSchema>,
    name_to_index: HashMap<String, usize>,
    fields: Vec<Field>,
    timestamp_index: Option<usize>,
    version: u32,
    metadata: HashMap<String, String>,
}

impl TryFrom<Vec<ColumnSchema>> for SchemaBuilder {
    type Error = Error;

    fn try_from(column_schemas: Vec<ColumnSchema>) -> Result<SchemaBuilder> {
        SchemaBuilder::try_from_columns(column_schemas)
    }
}

impl SchemaBuilder {
    pub fn try_from_columns(column_schemas: Vec<ColumnSchema>) -> Result<Self> {
        let FieldsAndIndices {
            fields,
            name_to_index,
            timestamp_index,
        } = collect_fields(&column_schemas)?;

        Ok(Self {
            column_schemas,
            name_to_index,
            fields,
            timestamp_index,
            ..Default::default()
        })
    }

    pub fn version(mut self, version: u32) -> Self {
        self.version = version;
        self
    }

    /// Add key value pair to metadata.
    ///
    /// Old metadata with same key would be overwritten.
    pub fn add_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        let _ = self.metadata.insert(key.into(), value.into());
        self
    }

    pub fn build(mut self) -> Result<Schema> {
        if let Some(timestamp_index) = self.timestamp_index {
            validate_timestamp_index(&self.column_schemas, timestamp_index)?;
        }

        self.metadata
            .insert(VERSION_KEY.to_string(), self.version.to_string());

        let arrow_schema = ArrowSchema::new(self.fields).with_metadata(self.metadata);

        Ok(Schema {
            column_schemas: self.column_schemas,
            name_to_index: self.name_to_index,
            arrow_schema: Arc::new(arrow_schema),
            timestamp_index: self.timestamp_index,
            version: self.version,
        })
    }
}

struct FieldsAndIndices {
    fields: Vec<Field>,
    name_to_index: HashMap<String, usize>,
    timestamp_index: Option<usize>,
}

fn collect_fields(column_schemas: &[ColumnSchema]) -> Result<FieldsAndIndices> {
    let mut fields = Vec::with_capacity(column_schemas.len());
    let mut name_to_index = HashMap::with_capacity(column_schemas.len());
    let mut timestamp_index = None;
    for (index, column_schema) in column_schemas.iter().enumerate() {
        if column_schema.is_time_index() && timestamp_index.is_none() {
            timestamp_index = Some(index);
        }
        let mut field = Field::try_from(column_schema)?;

        // Column with type Json or Vector performs the same as binary column in Arrow, so we need to mark it
        let extype = match column_schema.data_type {
            ConcreteDataType::Json(_) => Some(ColumnExtType::Json),
            ConcreteDataType::Vector(d) => Some(ColumnExtType::Vector(d.dim)),
            _ => None,
        };
        if let Some(extype) = extype {
            field
                .metadata_mut()
                .insert(TYPE_KEY.to_string(), extype.to_string());
        }
        fields.push(field);
        ensure!(
            name_to_index
                .insert(column_schema.name.clone(), index)
                .is_none(),
            DuplicateColumnSnafu {
                column: &column_schema.name,
            }
        );
    }

    Ok(FieldsAndIndices {
        fields,
        name_to_index,
        timestamp_index,
    })
}

fn validate_timestamp_index(column_schemas: &[ColumnSchema], timestamp_index: usize) -> Result<()> {
    ensure!(
        timestamp_index < column_schemas.len(),
        error::InvalidTimestampIndexSnafu {
            index: timestamp_index,
        }
    );

    let column_schema = &column_schemas[timestamp_index];
    ensure!(
        column_schema.data_type.is_timestamp(),
        error::InvalidTimestampIndexSnafu {
            index: timestamp_index,
        }
    );
    ensure!(
        column_schema.is_time_index(),
        error::InvalidTimestampIndexSnafu {
            index: timestamp_index,
        }
    );

    Ok(())
}

pub type SchemaRef = Arc<Schema>;

impl TryFrom<Arc<ArrowSchema>> for Schema {
    type Error = Error;

    fn try_from(arrow_schema: Arc<ArrowSchema>) -> Result<Schema> {
        let mut column_schemas = Vec::with_capacity(arrow_schema.fields.len());
        let mut name_to_index = HashMap::with_capacity(arrow_schema.fields.len());
        for field in &arrow_schema.fields {
            let column_schema = ColumnSchema::try_from(field.as_ref())?;
            let _ = name_to_index.insert(field.name().clone(), column_schemas.len());
            column_schemas.push(column_schema);
        }

        let mut timestamp_index = None;
        for (index, column_schema) in column_schemas.iter().enumerate() {
            if column_schema.is_time_index() {
                validate_timestamp_index(&column_schemas, index)?;
                timestamp_index = Some(index);
                break;
            }
        }

        let version = try_parse_version(&arrow_schema.metadata, VERSION_KEY)?;

        Ok(Self {
            column_schemas,
            name_to_index,
            arrow_schema,
            timestamp_index,
            version,
        })
    }
}

impl TryFrom<ArrowSchema> for Schema {
    type Error = Error;

    fn try_from(arrow_schema: ArrowSchema) -> Result<Schema> {
        let arrow_schema = Arc::new(arrow_schema);

        Schema::try_from(arrow_schema)
    }
}

impl TryFrom<DFSchemaRef> for Schema {
    type Error = Error;

    fn try_from(value: DFSchemaRef) -> Result<Self> {
        value.inner().clone().try_into()
    }
}

fn try_parse_version(metadata: &HashMap<String, String>, key: &str) -> Result<u32> {
    if let Some(value) = metadata.get(key) {
        let version = value
            .parse()
            .context(error::ParseSchemaVersionSnafu { value })?;

        Ok(version)
    } else {
        Ok(Schema::INITIAL_VERSION)
    }
}

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

    #[test]
    fn test_build_empty_schema() {
        let schema = SchemaBuilder::default().build().unwrap();
        assert_eq!(0, schema.num_columns());
        assert!(schema.is_empty());
    }

    #[test]
    fn test_schema_no_timestamp() {
        let column_schemas = vec![
            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), false),
            ColumnSchema::new("col2", ConcreteDataType::float64_datatype(), true),
        ];
        let schema = Schema::new(column_schemas.clone());

        assert_eq!(2, schema.num_columns());
        assert!(!schema.is_empty());
        assert!(schema.timestamp_index().is_none());
        assert!(schema.timestamp_column().is_none());
        assert_eq!(Schema::INITIAL_VERSION, schema.version());

        for column_schema in &column_schemas {
            let found = schema.column_schema_by_name(&column_schema.name).unwrap();
            assert_eq!(column_schema, found);
        }
        assert!(schema.column_schema_by_name("col3").is_none());

        let new_schema = Schema::try_from(schema.arrow_schema().clone()).unwrap();

        assert_eq!(schema, new_schema);
        assert_eq!(column_schemas, schema.column_schemas());
    }

    #[test]
    fn test_schema_duplicate_column() {
        let column_schemas = vec![
            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), false),
            ColumnSchema::new("col1", ConcreteDataType::float64_datatype(), true),
        ];
        let err = Schema::try_new(column_schemas).unwrap_err();

        assert!(
            matches!(err, Error::DuplicateColumn { .. }),
            "expect DuplicateColumn, found {}",
            err
        );
    }

    #[test]
    fn test_metadata() {
        let column_schemas = vec![ColumnSchema::new(
            "col1",
            ConcreteDataType::int32_datatype(),
            false,
        )];
        let schema = SchemaBuilder::try_from(column_schemas)
            .unwrap()
            .add_metadata("k1", "v1")
            .build()
            .unwrap();

        assert_eq!("v1", schema.metadata().get("k1").unwrap());
    }

    #[test]
    fn test_schema_with_timestamp() {
        let column_schemas = vec![
            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true),
            ColumnSchema::new(
                "ts",
                ConcreteDataType::timestamp_millisecond_datatype(),
                false,
            )
            .with_time_index(true),
        ];
        let schema = SchemaBuilder::try_from(column_schemas.clone())
            .unwrap()
            .version(123)
            .build()
            .unwrap();

        assert_eq!(1, schema.timestamp_index().unwrap());
        assert_eq!(&column_schemas[1], schema.timestamp_column().unwrap());
        assert_eq!(123, schema.version());

        let new_schema = Schema::try_from(schema.arrow_schema().clone()).unwrap();
        assert_eq!(1, schema.timestamp_index().unwrap());
        assert_eq!(schema, new_schema);
    }

    #[test]
    fn test_schema_wrong_timestamp() {
        let column_schemas = vec![
            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true)
                .with_time_index(true),
            ColumnSchema::new("col2", ConcreteDataType::float64_datatype(), false),
        ];
        assert!(
            SchemaBuilder::try_from(column_schemas)
                .unwrap()
                .build()
                .is_err()
        );

        let column_schemas = vec![
            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true),
            ColumnSchema::new("col2", ConcreteDataType::float64_datatype(), false)
                .with_time_index(true),
        ];

        assert!(
            SchemaBuilder::try_from(column_schemas)
                .unwrap()
                .build()
                .is_err()
        );
    }
}
