// 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::{BTreeMap, HashMap};

use api::helper::encode_json_value;
use api::v1::Rows;
use api::v1::helper::row;
use api::v1::value::ValueData;
use common_base::readable_size::ReadableSize;
use common_error::ext::{ErrorExt, WhateverResult};
use common_error::status_code::StatusCode;
use common_recordbatch::RecordBatches;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion_common::ScalarValue;
use datafusion_expr::{col, lit};
use datatypes::arrow::array::AsArray;
use datatypes::arrow::datatypes::{Float64Type, TimestampMillisecondType};
use datatypes::json::value::JsonValue;
use datatypes::prelude::ConcreteDataType;
use datatypes::types::json_type::{JsonNativeType, JsonObjectType};
use futures::TryStreamExt;
use serde_json::json;
use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner};
use store_api::region_request::RegionRequest;
use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution};

use crate::config::MitoConfig;
use crate::error::Error;
use crate::read::read_columns::{ReadColumn, ReadColumns};
use crate::read::scan_region::Scanner;
use crate::test_util;
use crate::test_util::{CreateRequestBuilder, TestEnv};

#[tokio::test]
async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResult<()> {
    // Create a region with a JSON2 field whose physical Parquet representation is a nested struct.
    // The scan below will only ask for field_0.a.x.

    let request = CreateRequestBuilder::new()
        .field_datatype(ConcreteDataType::json2(JsonNativeType::Object(
            JsonObjectType::from([
                (
                    "a".to_string(),
                    JsonNativeType::Object(JsonObjectType::from([
                        ("x".to_string(), JsonNativeType::i64()),
                        ("y".to_string(), JsonNativeType::String),
                    ])),
                ),
                ("b".to_string(), JsonNativeType::String),
            ]),
        )))
        .build();
    let schema = test_util::rows_schema(&request);

    let mut env = TestEnv::new().await;
    let engine = env.create_engine(MitoConfig::default()).await;
    let region_id = RegionId::new(1024, 0);
    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await?;

    // Write full JSON objects, then flush them so the scanner has an Parquet file where nested
    // projection can be pushed down.
    let rows = Rows {
        schema,
        rows: vec![
            row(vec![
                ValueData::StringValue("tag-1".to_string()),
                ValueData::JsonValue(encode_json_value(JsonValue::from(json!({
                    "a": { "x": 10, "y": "ignored-a" },
                    "b": "ignored-b"
                })))),
                ValueData::TimestampMillisecondValue(1000),
            ]),
            row(vec![
                ValueData::StringValue("tag-2".to_string()),
                ValueData::JsonValue(encode_json_value(JsonValue::from(json!({
                    "a": { "x": 20, "y": "ignored-c" },
                    "b": "ignored-d"
                })))),
                ValueData::TimestampMillisecondValue(2000),
            ]),
        ],
    };
    test_util::put_rows(&engine, region_id, rows).await;
    test_util::flush_region(&engine, region_id, None).await;

    // Without a type hint, the scanner reads the whole JSON2 root column.
    let request = ScanRequest {
        projection: Some(vec![1, 0]),
        ..Default::default()
    };
    let scanner = engine.scanner(region_id, request).await?;
    let Scanner::Seq(seq_scan) = &scanner else {
        unreachable!();
    };
    assert_eq!(
        seq_scan.input().read_cols,
        ReadColumns::from_deduped_column_ids([1, 0])
    );

    let stream = scanner.scan().await?;
    let batches = RecordBatches::try_collect(stream).await?;
    let expected = r#"
+------------------------------------------+-------+
| field_0                                  | tag_0 |
+------------------------------------------+-------+
| {a: {x: 10, y: ignored-a}, b: ignored-b} | tag-1 |
| {a: {x: 20, y: ignored-c}, b: ignored-d} | tag-2 |
+------------------------------------------+-------+
"#;
    assert_eq!(batches.pretty_print()?, expected.trim());

    // Simulate a query expression like json_get(field_0, 'a.x'): the logical projection still
    // returns the JSON2 root column, while json_type_hint tells scan input construction which
    // nested physical path is needed.
    let json_type_hint = HashMap::from([(
        "field_0".to_string(),
        JsonNativeType::Object(JsonObjectType::from([(
            "a".to_string(),
            JsonNativeType::Object(JsonObjectType::from([(
                "x".to_string(),
                JsonNativeType::i64(),
            )])),
        )])),
    )]);
    let request = ScanRequest {
        projection: Some(vec![1, 0]),
        json_type_hint: json_type_hint.clone(),
        ..Default::default()
    };
    let scanner = engine.scanner(region_id, request).await?;
    let Scanner::Seq(seq_scan) = &scanner else {
        unreachable!();
    };
    // Verify the scan input only asks storage for field_0.a.x instead of the
    // whole JSON2 struct. tag_0 is still read as a normal root column.
    assert_eq!(
        seq_scan.input().read_cols,
        ReadColumns {
            cols: vec![
                ReadColumn::new(
                    1,
                    vec![vec![
                        "field_0".to_string(),
                        "a".to_string(),
                        "x".to_string()
                    ]]
                ),
                ReadColumn::new(0, vec![]),
            ]
        }
    );

    // The scanner should still return a valid RecordBatch in the requested logical projection.
    // Fields outside the pushed-down path are pruned from the returned JSON2 type and value.

    let stream = scanner.scan().await?;
    let batches = RecordBatches::try_collect(stream).await?;
    let expected = r#"
+--------------+-------+
| field_0      | tag_0 |
+--------------+-------+
| {a: {x: 10}} | tag-1 |
| {a: {x: 20}} | tag-2 |
+--------------+-------+
"#;
    assert_eq!(batches.pretty_print()?, expected.trim());

    // A type hint only narrows a projected JSON2 root; it must not add an unprojected root.
    let request = ScanRequest {
        projection: Some(vec![0]),
        json_type_hint,
        ..Default::default()
    };
    let scanner = engine.scanner(region_id, request).await?;
    let Scanner::Seq(seq_scan) = &scanner else {
        unreachable!();
    };
    assert_eq!(
        seq_scan.input().read_cols,
        ReadColumns::from_deduped_column_ids([0])
    );

    let stream = scanner.scan().await?;
    let batches = RecordBatches::try_collect(stream).await?;
    let expected = r#"
+-------+
| tag_0 |
+-------+
| tag-1 |
| tag-2 |
+-------+
"#;
    assert_eq!(batches.pretty_print()?, expected.trim());
    Ok(())
}

#[tokio::test]
async fn test_incremental_query_stale_error() {
    let mut env = TestEnv::with_prefix("test_incremental_query_stale_error").await;
    let engine = env.create_engine(MitoConfig::default()).await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let rows = Rows {
        schema: column_schemas.clone(),
        rows: test_util::build_rows(0, 3),
    };
    test_util::put_rows(&engine, region_id, rows).await;
    test_util::flush_region(&engine, region_id, None).await;

    let err = engine
        .scanner(
            region_id,
            ScanRequest {
                memtable_min_sequence: Some(0),
                ..Default::default()
            },
        )
        .await
        .err()
        .expect("expect stale incremental error");

    let min_readable_seq = match &err {
        Error::IncrementalQueryStale {
            region_id: err_region_id,
            given_seq,
            min_readable_seq,
            ..
        } => {
            assert_eq!(*err_region_id, region_id);
            assert_eq!(*given_seq, 0);
            assert!(*min_readable_seq > 0);
            *min_readable_seq
        }
        _ => panic!("unexpected err: {err}"),
    };
    assert_eq!(StatusCode::RequestOutdated, err.status_code());
    let err_msg = err.to_string();
    assert!(err_msg.contains("STALE_CURSOR"));
    assert!(err_msg.contains(&region_id.to_string()));
    assert!(err_msg.contains("given_seq: 0"));
    assert!(err_msg.contains(&format!("min_readable_seq: {min_readable_seq}")));

    let incremental_rows = Rows {
        schema: column_schemas,
        rows: test_util::build_rows(3, 5),
    };
    test_util::put_rows(&engine, region_id, incremental_rows).await;

    let scanner = engine
        .scanner(
            region_id,
            ScanRequest {
                memtable_min_sequence: Some(min_readable_seq),
                skip_sst_files: true,
                ..Default::default()
            },
        )
        .await
        .unwrap();
    let stream = scanner.scan().await.unwrap();
    let batches = RecordBatches::try_collect(stream).await.unwrap();
    assert_eq!(
        batches.pretty_print().unwrap(),
        "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 3     | 3.0     | 1970-01-01T00:00:03 |
| 4     | 4.0     | 1970-01-01T00:00:04 |
+-------+---------+---------------------+"
    );
}

#[tokio::test]
async fn test_scan_with_min_sst_sequence() {
    test_scan_with_min_sst_sequence_with_format(false).await;
    test_scan_with_min_sst_sequence_with_format(true).await;
}

#[tokio::test]
async fn test_full_snapshot_upper_bound_returns_outdated_after_late_flush() {
    let mut env =
        TestEnv::with_prefix("test_full_snapshot_upper_bound_returns_outdated_after_late_flush")
            .await;
    let engine = env.create_engine(MitoConfig::default()).await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let first_rows = Rows {
        schema: column_schemas.clone(),
        rows: test_util::build_rows(0, 3),
    };
    test_util::put_rows(&engine, region_id, first_rows).await;
    test_util::flush_region(&engine, region_id, None).await;

    let snapshot_upper_bound = engine.get_committed_sequence(region_id).await.unwrap();

    let second_rows = Rows {
        schema: column_schemas,
        rows: test_util::build_rows(3, 5),
    };
    test_util::put_rows(&engine, region_id, second_rows).await;
    test_util::flush_region(&engine, region_id, None).await;

    let err = engine
        .scanner(
            region_id,
            ScanRequest {
                memtable_max_sequence: Some(snapshot_upper_bound),
                ..Default::default()
            },
        )
        .await
        .err()
        .expect("expect stale snapshot fence error");

    assert_eq!(StatusCode::RequestOutdated, err.status_code());
    let err_msg = err.to_string();
    assert!(err_msg.contains("STALE_SNAPSHOT_FENCE"));
    assert!(err_msg.contains(&region_id.to_string()));
    assert!(err_msg.contains(&format!("given_seq: {snapshot_upper_bound}")));
}

#[tokio::test]
async fn test_snapshot_bound_query_binds_memtable_upper_bound_at_scan_open() {
    let mut env =
        TestEnv::with_prefix("test_snapshot_bound_query_binds_memtable_upper_bound_at_scan_open")
            .await;
    let engine = env.create_engine(MitoConfig::default()).await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let first_rows = Rows {
        schema: column_schemas.clone(),
        rows: test_util::build_rows(0, 3),
    };
    test_util::put_rows(&engine, region_id, first_rows).await;

    let expected_snapshot = engine.get_committed_sequence(region_id).await.unwrap();
    let scanner = engine
        .scanner(
            region_id,
            ScanRequest {
                snapshot_on_scan: true,
                ..Default::default()
            },
        )
        .await
        .unwrap();
    assert_eq!(scanner.snapshot_sequence(), Some(expected_snapshot));

    let second_rows = Rows {
        schema: column_schemas,
        rows: test_util::build_rows(3, 5),
    };
    test_util::put_rows(&engine, region_id, second_rows).await;

    let stream = scanner.scan().await.unwrap();
    let batches = RecordBatches::try_collect(stream).await.unwrap();
    assert_eq!(
        batches.pretty_print().unwrap(),
        "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 0     | 0.0     | 1970-01-01T00:00:00 |
| 1     | 1.0     | 1970-01-01T00:00:01 |
| 2     | 2.0     | 1970-01-01T00:00:02 |
+-------+---------+---------------------+"
    );
}

#[tokio::test]
async fn test_snapshot_bound_query_keeps_open_snapshot_after_late_flush() {
    let mut env =
        TestEnv::with_prefix("test_snapshot_bound_query_keeps_open_snapshot_after_late_flush")
            .await;
    let engine = env.create_engine(MitoConfig::default()).await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let first_rows = Rows {
        schema: column_schemas.clone(),
        rows: test_util::build_rows(0, 3),
    };
    test_util::put_rows(&engine, region_id, first_rows).await;

    let expected_snapshot = engine.get_committed_sequence(region_id).await.unwrap();
    let scanner = engine
        .scanner(
            region_id,
            ScanRequest {
                snapshot_on_scan: true,
                ..Default::default()
            },
        )
        .await
        .unwrap();
    assert_eq!(scanner.snapshot_sequence(), Some(expected_snapshot));

    let second_rows = Rows {
        schema: column_schemas,
        rows: test_util::build_rows(3, 5),
    };
    test_util::put_rows(&engine, region_id, second_rows).await;
    test_util::flush_region(&engine, region_id, None).await;

    assert_eq!(scanner.snapshot_sequence(), Some(expected_snapshot));
}

#[tokio::test]
async fn test_snapshot_bound_query_keeps_correct_result_after_late_flush() {
    let mut env =
        TestEnv::with_prefix("test_snapshot_bound_query_keeps_correct_result_after_late_flush")
            .await;
    let engine = env.create_engine(MitoConfig::default()).await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let first_rows = Rows {
        schema: column_schemas.clone(),
        rows: test_util::build_rows(0, 3),
    };
    test_util::put_rows(&engine, region_id, first_rows).await;

    let expected_snapshot = engine.get_committed_sequence(region_id).await.unwrap();
    let scanner = engine
        .scanner(
            region_id,
            ScanRequest {
                snapshot_on_scan: true,
                ..Default::default()
            },
        )
        .await
        .unwrap();
    assert_eq!(scanner.snapshot_sequence(), Some(expected_snapshot));

    let second_rows = Rows {
        schema: column_schemas,
        rows: test_util::build_rows(3, 5),
    };
    test_util::put_rows(&engine, region_id, second_rows).await;
    test_util::flush_region(&engine, region_id, None).await;

    assert_eq!(scanner.snapshot_sequence(), Some(expected_snapshot));

    let stream = scanner.scan().await.unwrap();
    let batches = RecordBatches::try_collect(stream).await.unwrap();
    assert_eq!(
        batches.pretty_print().unwrap(),
        "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 0     | 0.0     | 1970-01-01T00:00:00 |
| 1     | 1.0     | 1970-01-01T00:00:01 |
| 2     | 2.0     | 1970-01-01T00:00:02 |
+-------+---------+---------------------+"
    );
}

async fn test_scan_with_min_sst_sequence_with_format(flat_format: bool) {
    let mut env = TestEnv::with_prefix("test_scan_with_min_sst_sequence").await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let put_rows = async |start, end| {
        let rows = Rows {
            schema: column_schemas.clone(),
            rows: test_util::build_rows(start, end),
        };
        test_util::put_rows(&engine, region_id, rows).await;
        test_util::flush_region(&engine, region_id, None).await;
    };
    // generates 3 SST files
    put_rows(0, 3).await;
    put_rows(3, 6).await;
    put_rows(6, 9).await;

    let scan_engine = async |file_min_sequence, expected_files, expected_data| {
        let request = ScanRequest {
            sst_min_sequence: file_min_sequence,
            ..Default::default()
        };
        let scanner = engine.scanner(region_id, request).await.unwrap();
        assert_eq!(scanner.num_files(), expected_files);

        let stream = scanner.scan().await.unwrap();
        let batches = RecordBatches::try_collect(stream).await.unwrap();
        assert_eq!(batches.pretty_print().unwrap(), expected_data);
    };

    // scans with no sst minimal sequence limit
    scan_engine(
        None,
        3,
        "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 0     | 0.0     | 1970-01-01T00:00:00 |
| 1     | 1.0     | 1970-01-01T00:00:01 |
| 2     | 2.0     | 1970-01-01T00:00:02 |
| 3     | 3.0     | 1970-01-01T00:00:03 |
| 4     | 4.0     | 1970-01-01T00:00:04 |
| 5     | 5.0     | 1970-01-01T00:00:05 |
| 6     | 6.0     | 1970-01-01T00:00:06 |
| 7     | 7.0     | 1970-01-01T00:00:07 |
| 8     | 8.0     | 1970-01-01T00:00:08 |
+-------+---------+---------------------+",
    )
    .await;

    // scans with sst minimal sequence > 2
    scan_engine(
        Some(2),
        3,
        "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 0     | 0.0     | 1970-01-01T00:00:00 |
| 1     | 1.0     | 1970-01-01T00:00:01 |
| 2     | 2.0     | 1970-01-01T00:00:02 |
| 3     | 3.0     | 1970-01-01T00:00:03 |
| 4     | 4.0     | 1970-01-01T00:00:04 |
| 5     | 5.0     | 1970-01-01T00:00:05 |
| 6     | 6.0     | 1970-01-01T00:00:06 |
| 7     | 7.0     | 1970-01-01T00:00:07 |
| 8     | 8.0     | 1970-01-01T00:00:08 |
+-------+---------+---------------------+",
    )
    .await;

    // scans with sst minimal sequence > 3
    scan_engine(
        Some(3),
        2,
        "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 3     | 3.0     | 1970-01-01T00:00:03 |
| 4     | 4.0     | 1970-01-01T00:00:04 |
| 5     | 5.0     | 1970-01-01T00:00:05 |
| 6     | 6.0     | 1970-01-01T00:00:06 |
| 7     | 7.0     | 1970-01-01T00:00:07 |
| 8     | 8.0     | 1970-01-01T00:00:08 |
+-------+---------+---------------------+",
    )
    .await;

    // scans with sst minimal sequence > 7
    scan_engine(
        Some(7),
        1,
        "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 6     | 6.0     | 1970-01-01T00:00:06 |
| 7     | 7.0     | 1970-01-01T00:00:07 |
| 8     | 8.0     | 1970-01-01T00:00:08 |
+-------+---------+---------------------+",
    )
    .await;

    // scans with sst minimal sequence > 9 (no sst files will be selected to scan)
    scan_engine(
        Some(9),
        0,
        "\
++
++",
    )
    .await;
}

#[tokio::test]
async fn test_max_concurrent_scan_files() {
    test_max_concurrent_scan_files_with_format(false).await;
    test_max_concurrent_scan_files_with_format(true).await;
}

async fn test_max_concurrent_scan_files_with_format(flat_format: bool) {
    let mut env = TestEnv::with_prefix("test_max_concurrent_scan_files").await;
    let config = MitoConfig {
        default_flat_format: flat_format,
        max_concurrent_scan_files: 2,
        ..Default::default()
    };
    let engine = env.create_engine(config).await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let put_and_flush = async |start, end| {
        let rows = Rows {
            schema: column_schemas.clone(),
            rows: test_util::build_rows(start, end),
        };
        test_util::put_rows(&engine, region_id, rows).await;
        test_util::flush_region(&engine, region_id, None).await;
    };

    // Write overlapping files.
    put_and_flush(0, 4).await;
    put_and_flush(3, 7).await;
    put_and_flush(6, 9).await;

    let request = ScanRequest::default();
    let scanner = engine.scanner(region_id, request).await.unwrap();
    let Scanner::Seq(scanner) = scanner else {
        panic!("Scanner should be seq scan");
    };
    let error = scanner.check_scan_limit().unwrap_err();
    assert_eq!(StatusCode::RateLimited, error.status_code());

    let request = ScanRequest {
        distribution: Some(TimeSeriesDistribution::PerSeries),
        ..Default::default()
    };
    let scanner = engine.scanner(region_id, request).await.unwrap();
    let Scanner::Series(scanner) = scanner else {
        panic!("Scanner should be series scan");
    };
    let error = scanner.check_scan_limit().unwrap_err();
    assert_eq!(StatusCode::RateLimited, error.status_code());
}

#[tokio::test]
async fn test_series_scan() {
    test_series_scan_with_format(false).await;
    test_series_scan_with_format(true).await;
}

async fn test_series_scan_with_format(flat_format: bool) {
    let mut env = TestEnv::with_prefix("test_series_scan").await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: flat_format,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new()
        .insert_option("compaction.type", "twcs")
        .insert_option("compaction.twcs.time_window", "1h")
        .build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    let put_flush_rows = async |start, end| {
        let rows = Rows {
            schema: column_schemas.clone(),
            rows: test_util::build_rows(start, end),
        };
        test_util::put_rows(&engine, region_id, rows).await;
        test_util::flush_region(&engine, region_id, None).await;
    };
    // generates 3 SST files
    put_flush_rows(0, 3).await;
    put_flush_rows(2, 6).await;
    put_flush_rows(3600, 3603).await;
    // Put to memtable.
    let rows = Rows {
        schema: column_schemas.clone(),
        rows: test_util::build_rows(7200, 7203),
    };
    test_util::put_rows(&engine, region_id, rows).await;

    let request = ScanRequest {
        distribution: Some(TimeSeriesDistribution::PerSeries),
        preserve_pk_dictionary_encoding: true,
        ..Default::default()
    };
    let scanner = engine.scanner(region_id, request).await.unwrap();
    let Scanner::Series(mut scanner) = scanner else {
        panic!("Scanner should be series scan");
    };
    // 3 partition ranges for 3 time window.
    assert_eq!(
        3,
        scanner.properties().partitions[0].len(),
        "unexpected ranges: {:?}",
        scanner.properties().partitions
    );
    let raw_ranges: Vec<_> = scanner
        .properties()
        .partitions
        .iter()
        .flatten()
        .cloned()
        .collect();
    let mut new_ranges = Vec::with_capacity(3);
    for range in raw_ranges {
        new_ranges.push(vec![range]);
    }
    scanner
        .prepare(PrepareRequest {
            ranges: Some(new_ranges),
            ..Default::default()
        })
        .unwrap();

    let actual_rows = collect_partition_rows_round_robin(&scanner, 3).await;

    let mut expected_rows = Vec::new();
    for value in [0_i64, 1, 2, 3, 4, 5, 3600, 3601, 3602, 7200, 7201, 7202] {
        expected_rows.push((value.to_string(), value as f64, value * 1000));
    }
    expected_rows.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2)));

    assert_eq!(expected_rows, actual_rows);
}

/// Scans all partitions in round-robin fashion and returns rows sorted by (tag, ts).
/// Also asserts that each series appears in only one partition.
async fn collect_partition_rows_round_robin(
    scanner: &dyn RegionScanner,
    num_partitions: usize,
) -> Vec<(String, f64, i64)> {
    let metrics_set = ExecutionPlanMetricsSet::default();

    let mut partition_batches = vec![vec![]; num_partitions];
    let mut streams: Vec<_> = (0..num_partitions)
        .map(|partition| {
            let stream = scanner
                .scan_partition(&Default::default(), &metrics_set, partition)
                .unwrap();
            Some(stream)
        })
        .collect();
    let mut num_done = 0;
    let mut schema = None;
    // Pull streams in round-robin fashion to get the consistent output from the sender.
    while num_done < num_partitions {
        if schema.is_none() {
            schema = Some(streams[0].as_ref().unwrap().schema().clone());
        }
        for i in 0..num_partitions {
            let Some(mut stream) = streams[i].take() else {
                continue;
            };
            let Some(rb) = stream.try_next().await.unwrap() else {
                num_done += 1;
                continue;
            };
            partition_batches[i].push(rb);
            streams[i] = Some(stream);
        }
    }

    let schema = schema.unwrap();
    collect_and_assert_partition_rows(schema, partition_batches)
}

/// Collects rows sorted by (tag, ts) from partition batches.
/// Also asserts that each series appears in only one partition.
fn collect_and_assert_partition_rows(
    schema: datatypes::schema::SchemaRef,
    partition_batches: Vec<Vec<common_recordbatch::RecordBatch>>,
) -> Vec<(String, f64, i64)> {
    let mut series_to_partition = BTreeMap::new();
    let mut actual_rows = Vec::new();

    for (partition, batches) in partition_batches.into_iter().enumerate() {
        let batches = RecordBatches::try_new(schema.clone(), batches).unwrap();
        let mut partition_series = Vec::new();

        for batch in batches.iter() {
            let tags = batch.column_by_name("tag_0").unwrap();
            let fields = batch
                .column_by_name("field_0")
                .unwrap()
                .as_primitive::<Float64Type>();
            let ts = batch
                .column_by_name("ts")
                .unwrap()
                .as_primitive::<TimestampMillisecondType>();

            for row in 0..batch.num_rows() {
                let tag = datatypes::arrow_array::string_array_value_at_index(tags, row)
                    .unwrap()
                    .to_string();
                let field = fields.value(row);
                let ts = ts.value(row);
                partition_series.push(tag.clone());
                actual_rows.push((tag, field, ts));
            }
        }

        partition_series.sort();
        partition_series.dedup();
        for tag in partition_series {
            let prev = series_to_partition.insert(tag.clone(), partition);
            assert_eq!(
                None, prev,
                "series {tag} appears in multiple partitions: {prev:?} and {partition}"
            );
        }
    }

    actual_rows.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2)));
    actual_rows
}

/// Tests series scan with multiple partition ranges (each with multiple overlapping sources)
/// and small semaphore permits (controlled by num_partitions).
#[tokio::test]
async fn test_series_scan_flat_small_permits() {
    let mut env = TestEnv::with_prefix("test_series_scan_small_permits").await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: true,
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new()
        .insert_option("compaction.type", "twcs")
        .insert_option("compaction.twcs.time_window", "1h")
        .build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    // Create overlapping SSTs in each time window so partition ranges have multiple sources.
    let put_flush_rows = async |start, end| {
        let rows = Rows {
            schema: column_schemas.clone(),
            rows: test_util::build_rows(start, end),
        };
        test_util::put_rows(&engine, region_id, rows).await;
        test_util::flush_region(&engine, region_id, None).await;
    };
    // Window 0 (0s-999s): 3 overlapping SSTs
    put_flush_rows(0, 3).await;
    put_flush_rows(1, 5).await;
    put_flush_rows(3, 7).await;
    // Window 1 (3600s-4599s): 2 overlapping SSTs
    put_flush_rows(3600, 3603).await;
    put_flush_rows(3601, 3605).await;
    // Window 2 (7200s-8199s): 2 overlapping SSTs
    put_flush_rows(7200, 7203).await;
    put_flush_rows(7201, 7204).await;

    let mut expected_rows = Vec::new();
    for value in [
        0_i64, 1, 2, 3, 4, 5, 6, 3600, 3601, 3602, 3603, 3604, 7200, 7201, 7202, 7203,
    ] {
        expected_rows.push((value.to_string(), value as f64, value * 1000));
    }
    expected_rows.sort_by(|a, b| a.0.cmp(&b.0).then(a.2.cmp(&b.2)));

    // Test with different semaphore sizes (num_partitions controls Semaphore::new(num_partitions)).
    for num_partitions in [1, 2] {
        let request = ScanRequest {
            distribution: Some(TimeSeriesDistribution::PerSeries),
            ..Default::default()
        };
        let scanner = engine.scanner(region_id, request).await.unwrap();
        let Scanner::Series(mut scanner) = scanner else {
            panic!("Scanner should be series scan");
        };

        // Collect all partition ranges and redistribute into `num_partitions` partitions.
        let raw_ranges: Vec<_> = scanner
            .properties()
            .partitions
            .iter()
            .flatten()
            .cloned()
            .collect();
        assert!(
            raw_ranges.len() >= 3,
            "expected at least 3 partition ranges, got {}",
            raw_ranges.len()
        );

        let mut new_ranges = vec![vec![]; num_partitions];
        for (i, range) in raw_ranges.into_iter().enumerate() {
            new_ranges[i % num_partitions].push(range);
        }
        scanner
            .prepare(PrepareRequest {
                ranges: Some(new_ranges),
                ..Default::default()
            })
            .unwrap();

        let actual_rows = collect_partition_rows_round_robin(&scanner, num_partitions).await;
        assert_eq!(
            expected_rows, actual_rows,
            "mismatch with num_partitions={num_partitions}"
        );
    }
}

// Regression test: `ts = a OR ts = b` extracts to a `TimestampRange` that
// `GenericRange::or` widens into `[min(a, b), max(a, b) + 1)`. Two such
// predicates with different `a` values can both extract to ranges that cover
// the same partition while selecting different (or no) rows. The previous
// cover check would strip both predicates from the cache key, letting the
// second scan return the first scan's cached row.
#[tokio::test]
async fn test_range_cache_separates_or_equality_time_filters() {
    let mut env = TestEnv::new().await;
    let engine = env
        .create_engine(MitoConfig {
            default_flat_format: true,
            // Explicitly enable the range result cache: the bug only reproduces
            // when the second scan can replay the first scan's cached batches.
            range_result_cache_size: ReadableSize::mb(64),
            ..Default::default()
        })
        .await;

    let region_id = RegionId::new(1, 1);
    let request = CreateRequestBuilder::new().build();
    let column_schemas = test_util::rows_schema(&request);

    engine
        .handle_request(region_id, RegionRequest::Create(request))
        .await
        .unwrap();

    // Partition rows ts=5..10 (5000ms..9000ms), flushed to SST.
    test_util::put_rows(
        &engine,
        region_id,
        Rows {
            schema: column_schemas.clone(),
            rows: test_util::build_rows(5, 10),
        },
    )
    .await;
    test_util::flush_region(&engine, region_id, None).await;

    let ts_lit = |ms: i64| lit(ScalarValue::TimestampMillisecond(Some(ms), None));
    let tag_filter = || col("tag_0").gt_eq(lit(ScalarValue::Utf8(Some("0".to_string()))));

    // First scan: (ts = 5000) OR (ts = 100000) -- extracts to `[5000, 100001)`,
    // which covers the partition `[5000, 9000]`. Selects ts=5.
    let stream = engine
        .scan_to_stream(
            region_id,
            ScanRequest {
                filters: vec![
                    tag_filter(),
                    col("ts").eq(ts_lit(5000)).or(col("ts").eq(ts_lit(100000))),
                ],
                ..Default::default()
            },
        )
        .await
        .unwrap();
    let batches = RecordBatches::try_collect(stream).await.unwrap();
    let expected_first = "\
+-------+---------+---------------------+
| tag_0 | field_0 | ts                  |
+-------+---------+---------------------+
| 5     | 5.0     | 1970-01-01T00:00:05 |
+-------+---------+---------------------+";
    assert_eq!(expected_first, batches.pretty_print().unwrap());

    // Second scan: (ts = 3000) OR (ts = 100000) -- extracts to `[3000, 100001)`,
    // which also covers the partition. Selects nothing. With the buggy cover
    // check both scans built the same cache key (tag filter only), so this scan
    // would replay the first scan's cached row and incorrectly return ts=5.
    let stream = engine
        .scan_to_stream(
            region_id,
            ScanRequest {
                filters: vec![
                    tag_filter(),
                    col("ts").eq(ts_lit(3000)).or(col("ts").eq(ts_lit(100000))),
                ],
                ..Default::default()
            },
        )
        .await
        .unwrap();
    let batches = RecordBatches::try_collect(stream).await.unwrap();
    let row_count: usize = batches.iter().map(|b| b.num_rows()).sum();
    assert_eq!(
        0,
        row_count,
        "expected empty result, got: {}",
        batches.pretty_print().unwrap()
    );
}
