// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! E2E integration tests for append-only (no primary key) tables.
//!
//! Covers: unpartitioned, partitioned, bucket=-1, fixed bucket,
//! multiple commits, column projection, and bucket predicate filtering.

use arrow_array::{Array, Int32Array, RecordBatch, StringArray};
use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema};
use futures::TryStreamExt;
use paimon::catalog::Identifier;
use paimon::io::FileIOBuilder;
use paimon::spec::{DataType, IntType, Schema, TableSchema, VarCharType};
use paimon::table::Table;
use std::sync::Arc;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn memory_file_io() -> paimon::io::FileIO {
    FileIOBuilder::new("memory").build().unwrap()
}

async fn setup_dirs(file_io: &paimon::io::FileIO, table_path: &str) {
    file_io
        .mkdirs(&format!("{table_path}/snapshot/"))
        .await
        .unwrap();
    file_io
        .mkdirs(&format!("{table_path}/manifest/"))
        .await
        .unwrap();
}

fn make_table(file_io: &paimon::io::FileIO, table_path: &str, schema: TableSchema) -> Table {
    Table::new(
        file_io.clone(),
        Identifier::new("default", "test"),
        table_path.to_string(),
        schema,
        None,
    )
}

fn int_batch(ids: Vec<i32>, values: Vec<i32>) -> RecordBatch {
    let schema = Arc::new(ArrowSchema::new(vec![
        ArrowField::new("id", ArrowDataType::Int32, false),
        ArrowField::new("value", ArrowDataType::Int32, false),
    ]));
    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(Int32Array::from(ids)),
            Arc::new(Int32Array::from(values)),
        ],
    )
    .unwrap()
}

fn partitioned_batch(pts: Vec<&str>, ids: Vec<i32>) -> RecordBatch {
    let schema = Arc::new(ArrowSchema::new(vec![
        ArrowField::new("pt", ArrowDataType::Utf8, false),
        ArrowField::new("id", ArrowDataType::Int32, false),
    ]));
    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(StringArray::from(pts)),
            Arc::new(Int32Array::from(ids)),
        ],
    )
    .unwrap()
}

fn collect_int_col(batches: &[RecordBatch], col: &str) -> Vec<i32> {
    let mut vals: Vec<i32> = batches
        .iter()
        .flat_map(|b| {
            let idx = b.schema().index_of(col).unwrap();
            b.column(idx)
                .as_any()
                .downcast_ref::<Int32Array>()
                .unwrap()
                .values()
                .to_vec()
        })
        .collect();
    vals.sort();
    vals
}

fn collect_string_col(batches: &[RecordBatch], col: &str) -> Vec<String> {
    let mut vals: Vec<String> = batches
        .iter()
        .flat_map(|b| {
            let idx = b.schema().index_of(col).unwrap();
            let arr = b
                .column(idx)
                .as_any()
                .downcast_ref::<StringArray>()
                .unwrap();
            (0..arr.len())
                .map(|i| arr.value(i).to_string())
                .collect::<Vec<_>>()
        })
        .collect();
    vals.sort();
    vals
}

/// Write batches → commit → scan → read, return all batches.
async fn write_commit_read(table: &Table, batches: Vec<RecordBatch>) -> Vec<RecordBatch> {
    let wb = table.new_write_builder();
    let mut tw = wb.new_write().unwrap();
    for batch in &batches {
        tw.write_arrow_batch(batch).await.unwrap();
    }
    wb.new_commit()
        .commit(tw.prepare_commit().await.unwrap())
        .await
        .unwrap();

    let rb = table.new_read_builder();
    let plan = rb.new_scan().plan().await.unwrap();
    let read = rb.new_read().unwrap();
    read.to_arrow(plan.splits())
        .unwrap()
        .try_collect()
        .await
        .unwrap()
}

// ---------------------------------------------------------------------------
// Unpartitioned, bucket = -1 (default)
// ---------------------------------------------------------------------------

fn unpartitioned_schema() -> TableSchema {
    let schema = Schema::builder()
        .column("id", DataType::Int(IntType::new()))
        .column("value", DataType::Int(IntType::new()))
        .build()
        .unwrap();
    TableSchema::new(0, &schema)
}

#[tokio::test]
async fn test_unpartitioned_single_batch() {
    let file_io = memory_file_io();
    let path = "memory:/append_unpart_single";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, unpartitioned_schema());

    let result = write_commit_read(&table, vec![int_batch(vec![1, 2, 3], vec![10, 20, 30])]).await;

    assert_eq!(collect_int_col(&result, "id"), vec![1, 2, 3]);
    assert_eq!(collect_int_col(&result, "value"), vec![10, 20, 30]);
}

#[tokio::test]
async fn test_unpartitioned_multiple_batches() {
    let file_io = memory_file_io();
    let path = "memory:/append_unpart_multi";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, unpartitioned_schema());

    let result = write_commit_read(
        &table,
        vec![
            int_batch(vec![1, 2], vec![10, 20]),
            int_batch(vec![3, 4, 5], vec![30, 40, 50]),
        ],
    )
    .await;

    assert_eq!(collect_int_col(&result, "id"), vec![1, 2, 3, 4, 5]);
}

#[tokio::test]
async fn test_unpartitioned_two_commits() {
    let file_io = memory_file_io();
    let path = "memory:/append_unpart_two_commits";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, unpartitioned_schema());

    // First commit
    let wb = table.new_write_builder();
    let mut tw = wb.new_write().unwrap();
    tw.write_arrow_batch(&int_batch(vec![1, 2], vec![10, 20]))
        .await
        .unwrap();
    wb.new_commit()
        .commit(tw.prepare_commit().await.unwrap())
        .await
        .unwrap();

    // Second commit
    let mut tw2 = wb.new_write().unwrap();
    tw2.write_arrow_batch(&int_batch(vec![3, 4], vec![30, 40]))
        .await
        .unwrap();
    wb.new_commit()
        .commit(tw2.prepare_commit().await.unwrap())
        .await
        .unwrap();

    // Read all
    let rb = table.new_read_builder();
    let plan = rb.new_scan().plan().await.unwrap();
    let read = rb.new_read().unwrap();
    let result: Vec<RecordBatch> = read
        .to_arrow(plan.splits())
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    assert_eq!(collect_int_col(&result, "id"), vec![1, 2, 3, 4]);
}

#[tokio::test]
async fn test_unpartitioned_projection() {
    let file_io = memory_file_io();
    let path = "memory:/append_unpart_proj";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, unpartitioned_schema());

    // Write
    let wb = table.new_write_builder();
    let mut tw = wb.new_write().unwrap();
    tw.write_arrow_batch(&int_batch(vec![1, 2, 3], vec![10, 20, 30]))
        .await
        .unwrap();
    wb.new_commit()
        .commit(tw.prepare_commit().await.unwrap())
        .await
        .unwrap();

    // Read with projection
    let mut rb = table.new_read_builder();
    rb.with_projection(&["value"])
        .expect("Projection should succeed");
    let plan = rb.new_scan().plan().await.unwrap();
    let read = rb.new_read().unwrap();
    let result: Vec<RecordBatch> = read
        .to_arrow(plan.splits())
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    assert_eq!(result[0].schema().fields().len(), 1);
    assert_eq!(result[0].schema().field(0).name(), "value");
    assert_eq!(collect_int_col(&result, "value"), vec![10, 20, 30]);
}

// ---------------------------------------------------------------------------
// Unpartitioned, fixed bucket
// ---------------------------------------------------------------------------

fn fixed_bucket_schema(buckets: i32) -> TableSchema {
    let schema = Schema::builder()
        .column("id", DataType::Int(IntType::new()))
        .column("value", DataType::Int(IntType::new()))
        .option("bucket", buckets.to_string())
        .option("bucket-key", "id")
        .build()
        .unwrap();
    TableSchema::new(0, &schema)
}

#[tokio::test]
async fn test_fixed_bucket_write_read() {
    let file_io = memory_file_io();
    let path = "memory:/append_fixed_bucket";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, fixed_bucket_schema(4));

    let result = write_commit_read(
        &table,
        vec![int_batch(
            vec![1, 2, 3, 4, 5, 6, 7, 8],
            vec![10, 20, 30, 40, 50, 60, 70, 80],
        )],
    )
    .await;

    assert_eq!(collect_int_col(&result, "id"), vec![1, 2, 3, 4, 5, 6, 7, 8]);
}

#[tokio::test]
async fn test_fixed_bucket_scan_filters_by_bucket() {
    use paimon::spec::{Datum, PredicateBuilder};

    let file_io = memory_file_io();
    let path = "memory:/append_bucket_filter";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, fixed_bucket_schema(4));

    // Write enough data to spread across buckets
    let wb = table.new_write_builder();
    let mut tw = wb.new_write().unwrap();
    tw.write_arrow_batch(&int_batch(
        vec![1, 2, 3, 4, 5, 6, 7, 8],
        vec![10, 20, 30, 40, 50, 60, 70, 80],
    ))
    .await
    .unwrap();
    wb.new_commit()
        .commit(tw.prepare_commit().await.unwrap())
        .await
        .unwrap();

    // Full scan — should have multiple buckets
    let full_rb = table.new_read_builder();
    let full_plan = full_rb.new_scan().plan().await.unwrap();
    let all_buckets: std::collections::HashSet<i32> =
        full_plan.splits().iter().map(|s| s.bucket()).collect();

    if all_buckets.len() <= 1 {
        // All rows hashed to same bucket — can't test filtering
        return;
    }

    // Filter by id = 1 — should narrow to one bucket
    let pb = PredicateBuilder::new(table.schema().fields());
    let filter = pb.equal("id", Datum::Int(1)).unwrap();

    let mut rb = table.new_read_builder();
    rb.with_filter(filter);
    let plan = rb.new_scan().plan().await.unwrap();
    let filtered_buckets: std::collections::HashSet<i32> =
        plan.splits().iter().map(|s| s.bucket()).collect();

    assert_eq!(
        filtered_buckets.len(),
        1,
        "Bucket predicate should narrow to one bucket, got: {filtered_buckets:?}"
    );
    assert!(filtered_buckets.is_subset(&all_buckets));

    // Read and verify id=1 is in the result
    let read = rb.new_read().unwrap();
    let result: Vec<RecordBatch> = read
        .to_arrow(plan.splits())
        .unwrap()
        .try_collect()
        .await
        .unwrap();
    let ids = collect_int_col(&result, "id");
    assert!(ids.contains(&1));
}

// ---------------------------------------------------------------------------
// Partitioned, bucket = -1
// ---------------------------------------------------------------------------

fn partitioned_schema() -> TableSchema {
    let schema = Schema::builder()
        .column("pt", DataType::VarChar(VarCharType::string_type()))
        .column("id", DataType::Int(IntType::new()))
        .partition_keys(["pt"])
        .build()
        .unwrap();
    TableSchema::new(0, &schema)
}

#[tokio::test]
async fn test_partitioned_write_read() {
    let file_io = memory_file_io();
    let path = "memory:/append_partitioned";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, partitioned_schema());

    let result = write_commit_read(
        &table,
        vec![partitioned_batch(
            vec!["a", "b", "a", "b"],
            vec![1, 2, 3, 4],
        )],
    )
    .await;

    let total: usize = result.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total, 4);
    assert_eq!(collect_int_col(&result, "id"), vec![1, 2, 3, 4]);
    assert_eq!(collect_string_col(&result, "pt"), vec!["a", "a", "b", "b"]);
}

#[tokio::test]
async fn test_partitioned_two_commits() {
    let file_io = memory_file_io();
    let path = "memory:/append_part_two_commits";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, partitioned_schema());

    let wb = table.new_write_builder();

    // First commit: partition "a"
    let mut tw1 = wb.new_write().unwrap();
    tw1.write_arrow_batch(&partitioned_batch(vec!["a", "a"], vec![1, 2]))
        .await
        .unwrap();
    wb.new_commit()
        .commit(tw1.prepare_commit().await.unwrap())
        .await
        .unwrap();

    // Second commit: partition "b"
    let mut tw2 = wb.new_write().unwrap();
    tw2.write_arrow_batch(&partitioned_batch(vec!["b", "b"], vec![3, 4]))
        .await
        .unwrap();
    wb.new_commit()
        .commit(tw2.prepare_commit().await.unwrap())
        .await
        .unwrap();

    // Read all
    let rb = table.new_read_builder();
    let plan = rb.new_scan().plan().await.unwrap();
    let read = rb.new_read().unwrap();
    let result: Vec<RecordBatch> = read
        .to_arrow(plan.splits())
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    assert_eq!(collect_int_col(&result, "id"), vec![1, 2, 3, 4]);
    assert_eq!(collect_string_col(&result, "pt"), vec!["a", "a", "b", "b"]);
}

#[tokio::test]
async fn test_partitioned_scan_partition_filter() {
    use paimon::spec::{Datum, PredicateBuilder};

    let file_io = memory_file_io();
    let path = "memory:/append_part_filter";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, partitioned_schema());

    // Write data to two partitions
    let wb = table.new_write_builder();
    let mut tw = wb.new_write().unwrap();
    tw.write_arrow_batch(&partitioned_batch(
        vec!["a", "b", "a", "b"],
        vec![1, 2, 3, 4],
    ))
    .await
    .unwrap();
    wb.new_commit()
        .commit(tw.prepare_commit().await.unwrap())
        .await
        .unwrap();

    // Filter by pt = "a"
    let pb = PredicateBuilder::new(table.schema().fields());
    let filter = pb.equal("pt", Datum::String("a".into())).unwrap();

    let mut rb = table.new_read_builder();
    rb.with_filter(filter);
    let plan = rb.new_scan().plan().await.unwrap();

    // Only partition "a" splits should survive
    for split in plan.splits() {
        let pt = split.partition().get_string(0).unwrap().to_string();
        assert_eq!(pt, "a");
    }

    let read = rb.new_read().unwrap();
    let result: Vec<RecordBatch> = read
        .to_arrow(plan.splits())
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    assert_eq!(collect_int_col(&result, "id"), vec![1, 3]);
    assert_eq!(collect_string_col(&result, "pt"), vec!["a", "a"]);
}

// ---------------------------------------------------------------------------
// Partitioned + fixed bucket
// ---------------------------------------------------------------------------

fn partitioned_bucket_schema(buckets: i32) -> TableSchema {
    let schema = Schema::builder()
        .column("pt", DataType::VarChar(VarCharType::string_type()))
        .column("id", DataType::Int(IntType::new()))
        .column("value", DataType::Int(IntType::new()))
        .partition_keys(["pt"])
        .option("bucket", buckets.to_string())
        .option("bucket-key", "id")
        .build()
        .unwrap();
    TableSchema::new(0, &schema)
}

fn partitioned_value_batch(pts: Vec<&str>, ids: Vec<i32>, values: Vec<i32>) -> RecordBatch {
    let schema = Arc::new(ArrowSchema::new(vec![
        ArrowField::new("pt", ArrowDataType::Utf8, false),
        ArrowField::new("id", ArrowDataType::Int32, false),
        ArrowField::new("value", ArrowDataType::Int32, false),
    ]));
    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(StringArray::from(pts)),
            Arc::new(Int32Array::from(ids)),
            Arc::new(Int32Array::from(values)),
        ],
    )
    .unwrap()
}

#[tokio::test]
async fn test_partitioned_fixed_bucket_write_read() {
    let file_io = memory_file_io();
    let path = "memory:/append_part_bucket";
    setup_dirs(&file_io, path).await;
    let table = make_table(&file_io, path, partitioned_bucket_schema(2));

    let wb = table.new_write_builder();
    let mut tw = wb.new_write().unwrap();
    tw.write_arrow_batch(&partitioned_value_batch(
        vec!["a", "a", "b", "b"],
        vec![1, 2, 3, 4],
        vec![10, 20, 30, 40],
    ))
    .await
    .unwrap();
    wb.new_commit()
        .commit(tw.prepare_commit().await.unwrap())
        .await
        .unwrap();

    let rb = table.new_read_builder();
    let plan = rb.new_scan().plan().await.unwrap();
    let read = rb.new_read().unwrap();
    let result: Vec<RecordBatch> = read
        .to_arrow(plan.splits())
        .unwrap()
        .try_collect()
        .await
        .unwrap();

    assert_eq!(collect_int_col(&result, "id"), vec![1, 2, 3, 4]);
    assert_eq!(collect_int_col(&result, "value"), vec![10, 20, 30, 40]);
}

#[tokio::test]
async fn test_reject_fixed_bucket_without_bucket_key() {
    let schema = Schema::builder()
        .column("id", DataType::Int(IntType::new()))
        .column("value", DataType::Int(IntType::new()))
        .option("bucket", "4")
        .build()
        .unwrap();
    let table_schema = TableSchema::new(0, &schema);

    let file_io = memory_file_io();
    let path = "memory:/append_reject_no_bucket_key";
    let table = make_table(&file_io, path, table_schema);

    let result = table.new_write_builder().new_write();
    assert!(result.is_err());
    let err = result.err().unwrap();
    assert!(
        matches!(&err, paimon::Error::Unsupported { message } if message.contains("bucket-key")),
        "Expected Unsupported error for missing bucket-key, got: {err:?}"
    );
}
