//! Partitioned writes: `create_partitioned_table` + identity-partition `append`
//! tag data files with the partition value (single-partition per commit) and
//! round-trip; a batch spanning partitions is rejected rather than mis-tagged.

use std::sync::Arc;

use anyhow::Result;
use skade::arrow_array::{Int64Array, RecordBatch, StringArray};
use skade::arrow_schema::{DataType, Field, Schema};

fn schema() -> Schema {
    Schema::new(vec![
        Field::new("repo", DataType::Utf8, false),
        Field::new("id", DataType::Int64, false),
    ])
}

fn batch(repo: &str, ids: Vec<i64>) -> Result<RecordBatch> {
    let repos: Vec<&str> = vec![repo; ids.len()];
    Ok(RecordBatch::try_new(
        Arc::new(schema()),
        vec![Arc::new(StringArray::from(repos)), Arc::new(Int64Array::from(ids))],
    )?)
}

#[tokio::test]
async fn partitioned_append_round_trips() -> Result<()> {
    let tmp = tempfile::tempdir()?;
    let wh = skade::open(tmp.path().join("lake")).await?;
    let mut t = wh.create_partitioned_table("facts", &schema(), &["repo"]).await?;

    // One repo per commit — the single-partition shape append commits.
    t.append(&[batch("holger", vec![1, 2, 3])?]).await?;
    t.append(&[batch("znippy", vec![4, 5])?]).await?;

    let all = t.read().await?;
    let total: usize = all.iter().map(|b| b.num_rows()).sum();
    assert_eq!(total, 5, "all rows across both partitions must read back");
    assert_eq!(t.count().await?, 5);
    Ok(())
}

#[tokio::test]
async fn multi_partition_batch_is_rejected() -> Result<()> {
    let tmp = tempfile::tempdir()?;
    let wh = skade::open(tmp.path().join("lake")).await?;
    let mut t = wh.create_partitioned_table("facts", &schema(), &["repo"]).await?;

    // A batch spanning two partition values must error, not silently mis-tag.
    let mixed = RecordBatch::try_new(
        Arc::new(schema()),
        vec![
            Arc::new(StringArray::from(vec!["holger", "znippy"])),
            Arc::new(Int64Array::from(vec![1i64, 2])),
        ],
    )?;
    let err = t.append(&[mixed]).await.expect_err("multi-partition append must fail");
    assert!(err.to_string().contains("multiple"), "unexpected error: {err}");
    Ok(())
}

#[tokio::test]
async fn empty_partition_cols_is_unpartitioned() -> Result<()> {
    let tmp = tempfile::tempdir()?;
    let wh = skade::open(tmp.path().join("lake")).await?;
    // Empty partition_cols == create_table: a mixed batch is then fine.
    let mut t = wh.create_partitioned_table("facts", &schema(), &[]).await?;
    t.append(&[batch("holger", vec![1])?]).await?;
    t.append(&[RecordBatch::try_new(
        Arc::new(schema()),
        vec![
            Arc::new(StringArray::from(vec!["a", "b"])),
            Arc::new(Int64Array::from(vec![2i64, 3])),
        ],
    )?])
    .await?;
    assert_eq!(t.count().await?, 3);
    Ok(())
}

#[tokio::test]
async fn ingest_parallel_partitioned_round_trips() -> Result<()> {
    let tmp = tempfile::tempdir()?;
    let wh = skade::open(tmp.path().join("lake")).await?;
    let mut t = wh.create_partitioned_table("facts", &schema(), &["repo"]).await?;

    // 8 groups → 8 parquet files, all single-partition (repo = holger),
    // encoded across cores, committed every 3 files.
    let groups: Vec<Vec<RecordBatch>> =
        (0..8).map(|k| vec![batch("holger", (k * 10..k * 10 + 10).collect()).unwrap()]).collect();
    let stats = t.ingest_parallel(groups, 3).await?;
    assert_eq!(stats.rows, 80);
    assert_eq!(t.count().await?, 80, "all parallel-encoded rows read back");
    Ok(())
}

#[tokio::test]
async fn ingest_parallel_unpartitioned_round_trips() -> Result<()> {
    let tmp = tempfile::tempdir()?;
    let wh = skade::open(tmp.path().join("lake")).await?;
    let mut t = wh.create_table("facts", &schema()).await?;
    let groups: Vec<Vec<RecordBatch>> =
        (0..5).map(|k| vec![batch("any", (k * 100..k * 100 + 100).collect()).unwrap()]).collect();
    let stats = t.ingest_parallel(groups, 2).await?;
    assert_eq!(stats.rows, 500);
    assert_eq!(t.count().await?, 500);
    Ok(())
}
