use datafusion::arrow::array::{ArrayRef, UInt64Array};
use datafusion::arrow::datatypes::{DataType, Field, SchemaBuilder};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::arrow::util::pretty::print_batches;
use datafusion::error::DataFusionError;
use datafusion::physical_plan::stream::RecordBatchReceiverStreamBuilder;
use futures::StreamExt;
use std::sync::Arc;
use std::time::Duration;

fn make_dummy_record_batch() -> Result<RecordBatch, DataFusionError> {
    let a: ArrayRef = Arc::new(UInt64Array::from(vec![1, 2]));
    let record_batch = RecordBatch::try_from_iter(vec![("a", a)])
        .map_err(|error| DataFusionError::External(Box::new(error)))?;
    Ok(record_batch)
}

#[tokio::main]
async fn main() {
    let schema = {
        let mut builder = SchemaBuilder::new();
        builder.push(Field::new("a", DataType::UInt64, false));
        builder.finish()
    };
    let mut stream = {
        let mut builder = RecordBatchReceiverStreamBuilder::new(Arc::new(schema), 1);
        {
            let tx = builder.tx();
            builder.spawn_blocking(move || loop {
                let record_batch = make_dummy_record_batch()?;
                match tx.blocking_send(Ok(record_batch)) {
                    Ok(()) => (),
                    Err(_sent) => {
                        // If we ignore this error, nothing "aborts" us.
                        eprintln!("ignoring send error for demonstration purposes");
                    }
                };
                println!("tick");
                std::thread::sleep(Duration::from_millis(300));
            });
        }
        builder.build()
    };

    for i in 0..3 {
        let record_batch = stream.next().await.unwrap().unwrap();
        println!("Record Batch #{i}");
        print_batches(&[record_batch]).unwrap();
    }

    // Docs claim this will abort the producer.
    drop(stream);

    // If it were aborted, we should stop seeing "tick" messages any moment now!
    tokio::time::sleep(Duration::new(5, 0)).await;
    println!("is it still going? bailing out.");
    std::process::exit(0);
}