use std::sync::Arc;

use arrow::array::{ArrayRef, RecordBatch, StringArray, StringViewArray};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
use datafusion::execution::memory_pool::MemoryPool;
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::functions_aggregate::count::count_udaf;
use datafusion::physical_expr::aggregate::AggregateExprBuilder;
use datafusion::physical_expr::expressions::{Column, Literal};
use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion::physical_plan::displayable;
use datafusion::physical_plan::repartition::RepartitionExec;
use datafusion::physical_plan::test::TestMemoryExec;
use datafusion::physical_plan::{ExecutionPlan, Partitioning, collect};
use datafusion::prelude::{SessionConfig, SessionContext};

// --- Knobs ---
const USE_UTF8VIEW: bool = true;
const NUM_ROWS: usize = 2_000_000;
const DATA_PARTITIONS: usize = 100;
const HASH_PARTITIONS: usize = 100;
const POOL_SIZE: usize = 1024 * 1024 * 1024; // 1GB

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    env_logger::init();

    let (schema, batch) = make_data();
    println!(
        "input batch: {} rows, {} bytes",
        batch.num_rows(),
        batch.get_array_memory_size()
    );

    let plan = build_plan(schema, batch)?;
    println!("plan:\n{}", displayable(plan.as_ref()).indent(true));

    let (ctx, pool) = make_session_ctx()?;

    println!("pool reserved before collect: {} bytes", pool.reserved());
    let results = collect(plan, ctx.task_ctx()).await?;
    println!("pool reserved after collect:  {} bytes", pool.reserved());

    let total_rows: usize = results.iter().map(|b| b.num_rows()).sum();
    let total_bytes: usize = results.iter().map(|b| b.get_array_memory_size()).sum();
    println!(
        "{} batches, {total_rows} rows, {total_bytes} bytes in output",
        results.len()
    );

    Ok(())
}

fn make_data() -> (SchemaRef, RecordBatch) {
    let values: Vec<String> = (0..NUM_ROWS).map(|i| format!("value_{i}")).collect();
    let string_refs: Vec<&str> = values.iter().map(|s| s.as_str()).collect();

    let (dt, col): (DataType, ArrayRef) = if USE_UTF8VIEW {
        (
            DataType::Utf8View,
            Arc::new(StringViewArray::from(string_refs)),
        )
    } else {
        (DataType::Utf8, Arc::new(StringArray::from(string_refs)))
    };

    let schema = Arc::new(Schema::new(vec![Field::new("a", dt, false)]));
    let batch = RecordBatch::try_new(Arc::clone(&schema), vec![col]).unwrap();
    (schema, batch)
}

fn build_plan(
    schema: SchemaRef,
    batch: RecordBatch,
) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
    let rows_per_partition = NUM_ROWS / DATA_PARTITIONS;
    let partitions: Vec<Vec<RecordBatch>> = (0..DATA_PARTITIONS)
        .map(|i| vec![batch.slice(i * rows_per_partition, rows_per_partition)])
        .collect();
    let plan = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?;

    // GROUP BY "a", COUNT(*)
    let group_by =
        PhysicalGroupBy::new_single(vec![(Arc::new(Column::new("a", 0)), "a".to_string())]);
    let count_expr = Arc::new(
        AggregateExprBuilder::new(count_udaf(), vec![Arc::new(Literal::new(1i8.into()))])
            .schema(Arc::clone(&schema))
            .alias("count")
            .build()?,
    );

    let plan = Arc::new(AggregateExec::try_new(
        AggregateMode::Partial,
        group_by.clone(),
        vec![count_expr.clone()],
        vec![None],
        plan,
        schema,
    )?);
    let partial_schema = plan.schema();

    let plan = Arc::new(RepartitionExec::try_new(
        plan,
        Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], HASH_PARTITIONS),
    )?);

    let plan = Arc::new(AggregateExec::try_new(
        AggregateMode::FinalPartitioned,
        group_by,
        vec![count_expr],
        vec![None],
        plan,
        partial_schema,
    )?);

    let plan: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(plan));
    Ok(plan)
}

fn make_session_ctx() -> Result<(SessionContext, Arc<dyn MemoryPool>), Box<dyn std::error::Error>> {
    let runtime = RuntimeEnvBuilder::new()
        .with_memory_limit(POOL_SIZE, 1.0)
        .with_disk_manager_builder(
            DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled),
        )
        .build_arc()?;
    let pool: Arc<dyn MemoryPool> = Arc::clone(&runtime.memory_pool);

    let config = SessionConfig::new();
    let ctx = SessionContext::new_with_config_rt(config, runtime);
    Ok((ctx, pool))
}