// 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::hash::Hasher;
use std::sync::Arc;

use datatypes::arrow::array::{Array, ArrayRef, BinaryBuilder, UInt64Array};
use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::arrow_array::{is_string_null_at, string_array_value_at_index};
use fxhash::FxHasher;
use mito_codec::row_converter::SparsePrimaryKeyCodec;
use snafu::ResultExt;
use store_api::storage::ColumnId;
use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;

use crate::error::{EncodePrimaryKeySnafu, Result, UnexpectedRequestSnafu};

/// Info about a tag column for TSID computation and sparse primary key encoding.
#[allow(dead_code)]
pub struct TagColumnInfo {
    /// Column name (used for label-name hash).
    pub name: String,
    /// Column index in the RecordBatch.
    pub index: usize,
    /// Column ID in the physical region.
    pub column_id: ColumnId,
}

/// Computes the TSID for each row in a [RecordBatch].
///
/// The TSID is a stable hash of the set of labels (tags) present in each row.
/// It accounts for both the names and values of all non-null tag columns.
///
/// # Logic
/// - If a row has no nulls across all `sorted_tag_columns`, it uses a precomputed hash of all label names.
/// - If a row has nulls, it dynamically computes a hash of the names of labels that are present (non-null).
/// - In both cases, it then hashes the values of those present labels in the order specified by `sorted_tag_columns`.
pub fn compute_tsid_array(
    batch: &RecordBatch,
    sorted_tag_columns: &[TagColumnInfo],
    tag_arrays: &[&ArrayRef],
) -> UInt64Array {
    let num_rows = batch.num_rows();

    let label_name_hash = {
        let mut hasher = FxHasher::default();
        for tag_col in sorted_tag_columns {
            hasher.write(tag_col.name.as_bytes());
            hasher.write_u8(0xff);
        }
        hasher.finish()
    };

    let mut tsid_values = Vec::with_capacity(num_rows);
    for row in 0..num_rows {
        let has_null = tag_arrays.iter().any(|arr| is_string_null_at(arr, row));

        let tsid = if !has_null {
            let mut hasher = FxHasher::default();
            hasher.write_u64(label_name_hash);
            for arr in tag_arrays {
                if let Some(value) = string_array_value_at_index(arr, row) {
                    hasher.write(value.as_bytes());
                    hasher.write_u8(0xff);
                }
            }
            hasher.finish()
        } else {
            let mut name_hasher = FxHasher::default();
            for (tc, arr) in sorted_tag_columns.iter().zip(tag_arrays.iter()) {
                if !is_string_null_at(arr, row) {
                    name_hasher.write(tc.name.as_bytes());
                    name_hasher.write_u8(0xff);
                }
            }
            let row_label_hash = name_hasher.finish();

            let mut val_hasher = FxHasher::default();
            val_hasher.write_u64(row_label_hash);
            for arr in tag_arrays {
                if let Some(value) = string_array_value_at_index(arr, row) {
                    val_hasher.write(value.as_bytes());
                    val_hasher.write_u8(0xff);
                }
            }
            val_hasher.finish()
        };

        tsid_values.push(tsid);
    }

    UInt64Array::from(tsid_values)
}

fn build_tag_arrays<'a>(
    batch: &'a RecordBatch,
    sorted_tag_columns: &[TagColumnInfo],
) -> Result<Vec<&'a ArrayRef>> {
    sorted_tag_columns
        .iter()
        .map(|tc| {
            let array = batch.column(tc.index);
            match array.data_type() {
                DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(array),
                DataType::Dictionary(key_type, value_type)
                    if key_type.is_integer() && value_type.is_string() =>
                {
                    Ok(array)
                }
                data_type => UnexpectedRequestSnafu {
                    reason: format!(
                        "Tag column '{}' must be a string, given: {data_type}",
                        tc.name
                    ),
                }
                .fail(),
            }
        })
        .collect()
}

/// Modifies a [RecordBatch] to include a sparse primary key column.
///
/// This function transforms the input `batch` into a new `RecordBatch` where the first column
/// is the generated primary key (named [PRIMARY_KEY_COLUMN_NAME]), followed by columns
/// indicated by `extra_column_indices`.
///
/// The primary key uses a "sparse" encoding, which compactly represents the row's identity
/// by only including non-null tag values. The encoding, handled by [SparsePrimaryKeyCodec],
/// consists of:
/// 1. The `table_id`.
/// 2. A `tsid` (Time Series ID), which is a hash of the present tags.
/// 3. The actual non-null tag values paired with their `column_id`.
///
/// # Parameters
/// - `batch`: The source [RecordBatch].
/// - `table_id`: The ID of the table.
/// - `sorted_tag_columns`: Metadata for tag columns, used for both TSID computation and PK encoding.
/// - `extra_column_indices`: Indices of columns from the original batch to keep in the output
///   (typically the timestamp and value fields).
pub fn modify_batch_sparse(
    batch: RecordBatch,
    table_id: u32,
    sorted_tag_columns: &[TagColumnInfo],
    extra_column_indices: &[usize],
) -> Result<RecordBatch> {
    let num_rows = batch.num_rows();
    let codec = SparsePrimaryKeyCodec::schemaless();
    let tag_arrays = build_tag_arrays(&batch, sorted_tag_columns)?;
    let tsid_array = compute_tsid_array(&batch, sorted_tag_columns, &tag_arrays);

    let mut pk_builder = BinaryBuilder::with_capacity(num_rows, 0);
    let mut buffer = Vec::new();
    for row in 0..num_rows {
        buffer.clear();
        codec
            .encode_internal(table_id, tsid_array.value(row), &mut buffer)
            .context(EncodePrimaryKeySnafu)?;

        let tags = sorted_tag_columns
            .iter()
            .zip(tag_arrays.iter())
            .filter_map(|(tc, arr)| {
                string_array_value_at_index(arr, row).map(|value| (tc.column_id, value.as_bytes()))
            });
        codec
            .encode_raw_tag_value(tags, &mut buffer)
            .context(EncodePrimaryKeySnafu)?;

        pk_builder.append_value(&buffer);
    }

    let pk_array = pk_builder.finish();

    let mut fields = vec![Arc::new(Field::new(
        PRIMARY_KEY_COLUMN_NAME,
        DataType::Binary,
        false,
    ))];
    let mut columns: Vec<Arc<dyn Array>> = vec![Arc::new(pk_array)];

    for &idx in extra_column_indices {
        fields.push(batch.schema().fields()[idx].clone());
        columns.push(batch.column(idx).clone());
    }

    let new_schema = Arc::new(ArrowSchema::new(fields));
    RecordBatch::try_new(new_schema, columns).map_err(|e| {
        UnexpectedRequestSnafu {
            reason: format!("Failed to build modified sparse RecordBatch: {e}"),
        }
        .build()
    })
}

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

    use api::v1::value::ValueData;
    use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, SemanticType, Value};
    use datatypes::arrow::array::{BinaryArray, Int64Array, StringArray};
    use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
    use datatypes::arrow::record_batch::RecordBatch;
    use store_api::codec::PrimaryKeyEncoding;
    use store_api::metadata::ColumnMetadata;
    use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;

    use super::*;
    use crate::row_modifier::{RowModifier, RowsIter, TableIdInput};

    fn build_sparse_test_batch() -> RecordBatch {
        let schema = Arc::new(ArrowSchema::new(vec![
            Field::new("greptime_timestamp", DataType::Int64, false),
            Field::new("greptime_value", DataType::Float64, true),
            Field::new("namespace", DataType::Utf8, true),
            Field::new("host", DataType::Utf8, true),
        ]));
        RecordBatch::try_new(
            schema,
            vec![
                Arc::new(Int64Array::from(vec![1000])),
                Arc::new(datatypes::arrow::array::Float64Array::from(vec![42.0])),
                Arc::new(StringArray::from(vec!["greptimedb"])),
                Arc::new(StringArray::from(vec!["127.0.0.1"])),
            ],
        )
        .unwrap()
    }

    fn sparse_tag_columns() -> Vec<TagColumnInfo> {
        vec![
            TagColumnInfo {
                name: "host".to_string(),
                index: 3,
                column_id: 3,
            },
            TagColumnInfo {
                name: "namespace".to_string(),
                index: 2,
                column_id: 2,
            },
        ]
    }

    #[test]
    fn test_compute_tsid_basic() {
        let schema = Arc::new(ArrowSchema::new(vec![
            Field::new("namespace", DataType::Utf8, true),
            Field::new("host", DataType::Utf8, true),
        ]));
        let batch = RecordBatch::try_new(
            schema,
            vec![
                Arc::new(StringArray::from(vec!["greptimedb"])),
                Arc::new(StringArray::from(vec!["127.0.0.1"])),
            ],
        )
        .unwrap();

        let tag_columns: Vec<TagColumnInfo> = vec![
            TagColumnInfo {
                name: "host".to_string(),
                index: 1,
                column_id: 2,
            },
            TagColumnInfo {
                name: "namespace".to_string(),
                index: 0,
                column_id: 1,
            },
        ];
        let tag_arrays = build_tag_arrays(&batch, &tag_columns).unwrap();
        let tsid_array = compute_tsid_array(&batch, &tag_columns, &tag_arrays);

        assert_eq!(tsid_array.value(0), 2721566936019240841);
    }

    #[test]
    fn test_compute_tsid_with_nulls() {
        let schema = Arc::new(ArrowSchema::new(vec![
            Field::new("a", DataType::Utf8, true),
            Field::new("b", DataType::Utf8, true),
        ]));
        let batch_no_null = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(vec!["A"])),
                Arc::new(StringArray::from(vec!["B"])),
            ],
        )
        .unwrap();
        let tag_cols_2: Vec<TagColumnInfo> = vec![
            TagColumnInfo {
                name: "a".to_string(),
                index: 0,
                column_id: 1,
            },
            TagColumnInfo {
                name: "b".to_string(),
                index: 1,
                column_id: 2,
            },
        ];
        let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2).unwrap();
        let tsid_no_null = compute_tsid_array(&batch_no_null, &tag_cols_2, &tag_arrays_2);

        let schema3 = Arc::new(ArrowSchema::new(vec![
            Field::new("a", DataType::Utf8, true),
            Field::new("b", DataType::Utf8, true),
            Field::new("c", DataType::Utf8, true),
        ]));
        let batch_with_null = RecordBatch::try_new(
            schema3,
            vec![
                Arc::new(StringArray::from(vec!["A"])),
                Arc::new(StringArray::from(vec!["B"])),
                Arc::new(StringArray::from(vec![None as Option<&str>])),
            ],
        )
        .unwrap();
        let tag_cols_3: Vec<TagColumnInfo> = vec![
            TagColumnInfo {
                name: "a".to_string(),
                index: 0,
                column_id: 1,
            },
            TagColumnInfo {
                name: "b".to_string(),
                index: 1,
                column_id: 2,
            },
            TagColumnInfo {
                name: "c".to_string(),
                index: 2,
                column_id: 3,
            },
        ];
        let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3).unwrap();
        let tsid_with_null = compute_tsid_array(&batch_with_null, &tag_cols_3, &tag_arrays_3);

        assert_eq!(tsid_no_null.value(0), tsid_with_null.value(0));
    }

    #[test]
    fn test_modify_batch_sparse() {
        let batch = build_sparse_test_batch();
        let tag_columns = sparse_tag_columns();
        let non_tag_indices = vec![0, 1];
        let table_id: u32 = 1025;

        let modified =
            modify_batch_sparse(batch, table_id, &tag_columns, &non_tag_indices).unwrap();

        assert_eq!(modified.num_columns(), 3);
        assert_eq!(modified.schema().field(0).name(), PRIMARY_KEY_COLUMN_NAME);
        assert_eq!(modified.schema().field(1).name(), "greptime_timestamp");
        assert_eq!(modified.schema().field(2).name(), "greptime_value");
    }

    #[test]
    fn test_modify_batch_sparse_matches_row_modifier() {
        let batch = build_sparse_test_batch();
        let tag_columns = sparse_tag_columns();
        let non_tag_indices = vec![0, 1];
        let table_id: u32 = 1025;
        let modified =
            modify_batch_sparse(batch, table_id, &tag_columns, &non_tag_indices).unwrap();

        let make_info = |name: &str, column_id: ColumnId| ColumnMetadata {
            column_schema: datatypes::schema::ColumnSchema::new(
                name.to_string(),
                datatypes::prelude::ConcreteDataType::string_datatype(),
                false,
            ),
            semantic_type: SemanticType::Tag,
            column_id,
        };
        let name_to_column_id: HashMap<String, ColumnMetadata> = [
            (
                "greptime_timestamp".to_string(),
                make_info("greptime_timestamp", 0),
            ),
            ("greptime_value".to_string(), make_info("greptime_value", 1)),
            ("namespace".to_string(), make_info("namespace", 2)),
            ("host".to_string(), make_info("host", 3)),
        ]
        .into_iter()
        .collect();

        let rows = Rows {
            schema: vec![
                ColumnSchema {
                    column_name: "greptime_timestamp".to_string(),
                    datatype: ColumnDataType::TimestampMillisecond as i32,
                    semantic_type: SemanticType::Timestamp as i32,
                    ..Default::default()
                },
                ColumnSchema {
                    column_name: "greptime_value".to_string(),
                    datatype: ColumnDataType::Float64 as i32,
                    semantic_type: SemanticType::Field as i32,
                    ..Default::default()
                },
                ColumnSchema {
                    column_name: "namespace".to_string(),
                    datatype: ColumnDataType::String as i32,
                    semantic_type: SemanticType::Tag as i32,
                    ..Default::default()
                },
                ColumnSchema {
                    column_name: "host".to_string(),
                    datatype: ColumnDataType::String as i32,
                    semantic_type: SemanticType::Tag as i32,
                    ..Default::default()
                },
            ],
            rows: vec![Row {
                values: vec![
                    Value {
                        value_data: Some(ValueData::TimestampMillisecondValue(1000)),
                    },
                    Value {
                        value_data: Some(ValueData::F64Value(42.0)),
                    },
                    Value {
                        value_data: Some(ValueData::StringValue("greptimedb".to_string())),
                    },
                    Value {
                        value_data: Some(ValueData::StringValue("127.0.0.1".to_string())),
                    },
                ],
            }],
        };

        let row_iter = RowsIter::new(rows, &name_to_column_id);
        let rows = RowModifier::default()
            .modify_rows(
                row_iter,
                TableIdInput::Single(table_id),
                PrimaryKeyEncoding::Sparse,
            )
            .unwrap();
        let ValueData::BinaryValue(expected_pk) =
            rows.rows[0].values[0].value_data.clone().unwrap()
        else {
            panic!("expected binary primary key");
        };

        let actual_array = modified
            .column(0)
            .as_any()
            .downcast_ref::<BinaryArray>()
            .unwrap();
        assert_eq!(actual_array.value(0), expected_pk.as_slice());
    }

    #[test]
    fn label_replace_with_utf8view_labels_does_not_panic() {
        // Tag (label) columns may arrive as any string representation (`Utf8`,
        // `LargeUtf8`, `Utf8View`, or dictionary-encoded). `modify_batch_sparse`
        // must not assume the tag columns are plain `StringArray`s.
        let tag_arrays: Vec<ArrayRef> = vec![
            Arc::new(StringArray::from(vec!["greptimedb"])),
            Arc::new(datatypes::arrow::array::LargeStringArray::from(vec![
                "greptimedb",
            ])),
            Arc::new(datatypes::arrow::array::StringViewArray::from(vec![
                "greptimedb",
            ])),
            Arc::new(datatypes::arrow::array::DictionaryArray::<
                datatypes::arrow::datatypes::UInt32Type,
            >::new(
                datatypes::arrow::array::UInt32Array::from(vec![0]),
                Arc::new(StringArray::from(vec!["greptimedb"])),
            )),
        ];
        let tag_columns = vec![TagColumnInfo {
            name: "namespace".to_string(),
            index: 2,
            column_id: 2,
        }];
        let non_tag_indices = vec![0, 1];

        let primary_keys = tag_arrays
            .into_iter()
            .map(|tag_array| {
                let schema = Arc::new(ArrowSchema::new(vec![
                    Field::new("greptime_timestamp", DataType::Int64, false),
                    Field::new("greptime_value", DataType::Float64, true),
                    Field::new("namespace", tag_array.data_type().clone(), true),
                ]));
                let batch = RecordBatch::try_new(
                    schema,
                    vec![
                        Arc::new(Int64Array::from(vec![1000])),
                        Arc::new(datatypes::arrow::array::Float64Array::from(vec![42.0])),
                        tag_array,
                    ],
                )
                .unwrap();
                let modified =
                    modify_batch_sparse(batch, 1025, &tag_columns, &non_tag_indices).unwrap();
                modified
                    .column(0)
                    .as_any()
                    .downcast_ref::<BinaryArray>()
                    .unwrap()
                    .value(0)
                    .to_vec()
            })
            .collect::<Vec<_>>();

        assert!(primary_keys.windows(2).all(|keys| keys[0] == keys[1]));
    }

    #[test]
    fn test_compute_tsid_utf8view_with_nulls_matches_utf8() {
        // Regression test for Utf8View tag/label columns: the whole
        // `build_tag_arrays` -> `compute_tsid_array` -> sparse primary key path must handle
        // `Utf8View` columns exactly like `Utf8`, including rows with null tags, instead of
        // panicking or producing different TSIDs.
        let namespace_values: Vec<Option<&str>> =
            vec![Some("ns-a"), Some("ns-b"), None, Some("ns-c")];
        let host_values: Vec<Option<&str>> = vec![Some("host-1"), None, Some("host-2"), None];

        let build_batch = |data_type: DataType| {
            let (namespace, host): (ArrayRef, ArrayRef) = match &data_type {
                DataType::Utf8 => (
                    Arc::new(StringArray::from(namespace_values.clone())),
                    Arc::new(StringArray::from(host_values.clone())),
                ),
                DataType::Utf8View => (
                    Arc::new(datatypes::arrow::array::StringViewArray::from(
                        namespace_values.clone(),
                    )),
                    Arc::new(datatypes::arrow::array::StringViewArray::from(
                        host_values.clone(),
                    )),
                ),
                _ => unreachable!(),
            };
            let schema = Arc::new(ArrowSchema::new(vec![
                Field::new("greptime_timestamp", DataType::Int64, false),
                Field::new("greptime_value", DataType::Float64, true),
                Field::new("namespace", data_type.clone(), true),
                Field::new("host", data_type, true),
            ]));
            RecordBatch::try_new(
                schema,
                vec![
                    Arc::new(Int64Array::from(vec![1000, 1001, 1002, 1003])),
                    Arc::new(datatypes::arrow::array::Float64Array::from(vec![
                        1.0, 2.0, 3.0, 4.0,
                    ])),
                    namespace,
                    host,
                ],
            )
            .unwrap()
        };

        let tag_columns = vec![
            TagColumnInfo {
                name: "host".to_string(),
                index: 3,
                column_id: 3,
            },
            TagColumnInfo {
                name: "namespace".to_string(),
                index: 2,
                column_id: 2,
            },
        ];

        let utf8_batch = build_batch(DataType::Utf8);
        let utf8view_batch = build_batch(DataType::Utf8View);

        // Both representations must be accepted by `build_tag_arrays`.
        let utf8_tag_arrays = build_tag_arrays(&utf8_batch, &tag_columns).unwrap();
        let utf8view_tag_arrays = build_tag_arrays(&utf8view_batch, &tag_columns).unwrap();

        // TSIDs must be identical across representations, including rows with null tags.
        let utf8_tsids = compute_tsid_array(&utf8_batch, &tag_columns, &utf8_tag_arrays);
        let utf8view_tsids =
            compute_tsid_array(&utf8view_batch, &tag_columns, &utf8view_tag_arrays);
        assert_eq!(utf8_tsids, utf8view_tsids);

        // The full sparse primary key (TSID + tag values) must also be identical.
        let non_tag_indices = vec![0, 1];
        let modified_utf8 =
            modify_batch_sparse(utf8_batch, 1025, &tag_columns, &non_tag_indices).unwrap();
        let modified_utf8view =
            modify_batch_sparse(utf8view_batch, 1025, &tag_columns, &non_tag_indices).unwrap();
        let utf8_pks = modified_utf8
            .column(0)
            .as_any()
            .downcast_ref::<BinaryArray>()
            .unwrap();
        let utf8view_pks = modified_utf8view
            .column(0)
            .as_any()
            .downcast_ref::<BinaryArray>()
            .unwrap();
        for row in 0..utf8_pks.len() {
            assert_eq!(utf8_pks.value(row), utf8view_pks.value(row));
        }
    }
}
