// 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.

use std::io::Cursor;
use std::ops::Range;
use std::sync::Arc;
use std::time::SystemTime;

use arrow::array::{ArrayRef, Int8Array, Int64Array, StringArray};
use arrow::datatypes::{Field, Schema, SchemaBuilder};
use arrow::record_batch::RecordBatch;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{
    ParquetFileMetrics, ParquetFileReaderFactory, ParquetSource,
};
use datafusion::physical_plan::collect;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::prelude::SessionContext;
use datafusion_common::Result;
use datafusion_common::test_util::batches_to_sort_string;

use bytes::Bytes;
use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
use datafusion_datasource::source::DataSourceExec;
use datafusion_datasource_parquet::metadata::DFParquetMetadata;
use datafusion_datasource_parquet::{ParquetAccessPlan, RowGroupAccess};
use futures::future::BoxFuture;
use futures::{FutureExt, TryFutureExt};
use insta::assert_snapshot;
use object_store::memory::InMemory;
use object_store::path::Path;
use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
use parquet::arrow::ArrowWriter;
use parquet::arrow::arrow_reader::ArrowReaderOptions;
use parquet::arrow::async_reader::AsyncFileReader;
use parquet::errors::ParquetError;
use parquet::file::metadata::ParquetMetaData;
use parquet::file::properties::WriterProperties;

const EXPECTED_USER_DEFINED_METADATA: &str = "some-user-defined-metadata";

#[tokio::test]
async fn route_data_access_ops_to_parquet_file_reader_factory() {
    let c1: ArrayRef = Arc::new(StringArray::from(vec![Some("Foo"), None, Some("bar")]));
    let c2: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), Some(2), None]));
    let c3: ArrayRef = Arc::new(Int8Array::from(vec![Some(10), Some(20), None]));

    let batch = create_batch(vec![
        ("c1", c1.clone()),
        ("c2", c2.clone()),
        ("c3", c3.clone()),
    ]);

    let file_schema = batch.schema().clone();
    let (in_memory_object_store, parquet_files_meta) =
        store_parquet_in_memory(vec![batch]).await;
    let file_group = parquet_files_meta
        .into_iter()
        .map(|meta| {
            PartitionedFile::new_from_meta(meta)
                .with_extension(String::from(EXPECTED_USER_DEFINED_METADATA))
        })
        .collect();

    let source = Arc::new(
        ParquetSource::new(file_schema.clone())
            // prepare the scan
            .with_parquet_file_reader_factory(Arc::new(
                InMemoryParquetFileReaderFactory(Arc::clone(&in_memory_object_store)),
            )),
    );
    let base_config = FileScanConfigBuilder::new(
        // just any url that doesn't point to in memory object store
        ObjectStoreUrl::local_filesystem(),
        source,
    )
    .with_file_group(file_group)
    .build();

    let parquet_exec = DataSourceExec::from_data_source(base_config);

    let session_ctx = SessionContext::new();
    let task_ctx = session_ctx.task_ctx();
    let read = collect(parquet_exec, task_ctx).await.unwrap();

    assert_snapshot!(batches_to_sort_string(&read), @r"
    +-----+----+----+
    | c1  | c2 | c3 |
    +-----+----+----+
    |     | 2  | 20 |
    | Foo | 1  | 10 |
    | bar |    |    |
    +-----+----+----+
    ");
}

/// Regression test for the type-keyed extensions map: independent components
/// must be able to attach their own per-file payloads on the same
/// [`PartitionedFile`] without colliding. Here we attach a custom reader
/// payload (the `String` checked by [`InMemoryParquetFileReaderFactory`])
/// *and* a [`ParquetAccessPlan`] that skips the first row group, then verify
/// (a) the factory still sees its payload (its internal `assert_eq!` would
/// fire if the slot got overwritten) and (b) the access plan is honored — so
/// only the second row group's 5 rows come out, not all 10.
#[tokio::test]
async fn custom_payload_and_access_plan_coexist() {
    // Two row groups of 5 rows each: values 0..=4 in row group 0, 5..=9 in
    // row group 1.
    let c1: ArrayRef = Arc::new(Int64Array::from((0..10).collect::<Vec<i64>>()));
    let batch = create_batch(vec![("c1", c1)]);
    let file_schema = batch.schema().clone();

    let in_memory = InMemory::new();
    let mut buf = Vec::<u8>::with_capacity(32 * 1024);
    let props = WriterProperties::builder()
        .set_max_row_group_row_count(Some(5))
        .build();
    let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
    writer.write(&batch).unwrap();
    writer.close().unwrap();

    let location = Path::parse("two-row-groups.parquet").unwrap();
    let size = buf.len() as u64;
    in_memory
        .put(&location, Bytes::from(buf).into())
        .await
        .unwrap();
    let meta = ObjectMeta {
        location,
        last_modified: chrono::DateTime::from(SystemTime::now()),
        size,
        e_tag: None,
        version: None,
    };

    let access_plan =
        ParquetAccessPlan::new(vec![RowGroupAccess::Skip, RowGroupAccess::Scan]);
    let pf = PartitionedFile::new_from_meta(meta)
        .with_extension(String::from(EXPECTED_USER_DEFINED_METADATA))
        .with_extension(access_plan);

    let store: Arc<dyn ObjectStore> = Arc::new(in_memory);
    let source = Arc::new(
        ParquetSource::new(file_schema.clone()).with_parquet_file_reader_factory(
            Arc::new(InMemoryParquetFileReaderFactory(Arc::clone(&store))),
        ),
    );
    let base_config =
        FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
            .with_file_group(vec![pf].into())
            .build();
    let parquet_exec = DataSourceExec::from_data_source(base_config);

    let session_ctx = SessionContext::new();
    let read = collect(parquet_exec, session_ctx.task_ctx()).await.unwrap();

    let total: usize = read.iter().map(|b| b.num_rows()).sum();
    assert_eq!(
        total, 5,
        "access plan should have skipped the first row group"
    );

    let values: Vec<i64> = read
        .iter()
        .flat_map(|b| {
            let arr = b.column(0).as_any().downcast_ref::<Int64Array>().unwrap();
            (0..arr.len()).map(|i| arr.value(i)).collect::<Vec<_>>()
        })
        .collect();
    assert_eq!(values, vec![5, 6, 7, 8, 9]);
}

#[derive(Debug)]
struct InMemoryParquetFileReaderFactory(Arc<dyn ObjectStore>);

impl ParquetFileReaderFactory for InMemoryParquetFileReaderFactory {
    fn create_reader(
        &self,
        partition_index: usize,
        partitioned_file: PartitionedFile,
        metadata_size_hint: Option<usize>,
        metrics: &ExecutionPlanMetricsSet,
    ) -> Result<Box<dyn AsyncFileReader + Send>> {
        let metadata = partitioned_file
            .extension::<String>()
            .expect("has user defined metadata");

        assert_eq!(EXPECTED_USER_DEFINED_METADATA, &metadata[..]);

        let parquet_file_metrics = ParquetFileMetrics::new(
            partition_index,
            partitioned_file.object_meta.location.as_ref(),
            metrics,
        );

        Ok(Box::new(ParquetFileReader {
            store: Arc::clone(&self.0),
            meta: partitioned_file.object_meta,
            metrics: parquet_file_metrics,
            metadata_size_hint,
        }))
    }
}

fn create_batch(columns: Vec<(&str, ArrayRef)>) -> RecordBatch {
    columns.into_iter().fold(
        RecordBatch::new_empty(Arc::new(Schema::empty())),
        |batch, (field_name, arr)| add_to_batch(&batch, field_name, arr.clone()),
    )
}

fn add_to_batch(batch: &RecordBatch, field_name: &str, array: ArrayRef) -> RecordBatch {
    let mut fields = SchemaBuilder::from(batch.schema().fields());
    fields.push(Field::new(field_name, array.data_type().clone(), true));
    let schema = Arc::new(fields.finish());

    let mut columns = batch.columns().to_vec();
    columns.push(array);
    RecordBatch::try_new(schema, columns).expect("error; creating record batch")
}

async fn store_parquet_in_memory(
    batches: Vec<RecordBatch>,
) -> (Arc<dyn ObjectStore>, Vec<ObjectMeta>) {
    let in_memory = InMemory::new();

    let parquet_batches: Vec<(ObjectMeta, Bytes)> = batches
        .into_iter()
        .enumerate()
        .map(|(offset, batch)| {
            let mut buf = Vec::<u8>::with_capacity(32 * 1024);
            let mut output = Cursor::new(&mut buf);

            let mut writer = ArrowWriter::try_new(&mut output, batch.schema(), None)
                .expect("creating writer");

            writer.write(&batch).expect("Writing batch");
            writer.close().unwrap();

            let meta = ObjectMeta {
                location: Path::parse(format!("file-{offset}.parquet"))
                    .expect("creating path"),
                last_modified: chrono::DateTime::from(SystemTime::now()),
                size: buf.len() as u64,
                e_tag: None,
                version: None,
            };

            (meta, Bytes::from(buf))
        })
        .collect();

    let mut objects = Vec::with_capacity(parquet_batches.len());
    for (meta, bytes) in parquet_batches {
        in_memory
            .put(&meta.location, bytes.into())
            .await
            .expect("put parquet file into in memory object store");
        objects.push(meta);
    }

    (Arc::new(in_memory), objects)
}

/// Implements [`AsyncFileReader`] for a parquet file in object storage
struct ParquetFileReader {
    store: Arc<dyn ObjectStore>,
    meta: ObjectMeta,
    metrics: ParquetFileMetrics,
    metadata_size_hint: Option<usize>,
}

impl AsyncFileReader for ParquetFileReader {
    fn get_bytes(
        &mut self,
        range: Range<u64>,
    ) -> BoxFuture<'_, parquet::errors::Result<Bytes>> {
        let bytes_scanned = range.end - range.start;
        self.metrics.bytes_scanned.add(bytes_scanned as usize);

        self.store
            .get_range(&self.meta.location, range)
            .map_err(|e| {
                ParquetError::General(format!("AsyncChunkReader::get_bytes error: {e}"))
            })
            .boxed()
    }

    fn get_metadata(
        &mut self,
        _options: Option<&ArrowReaderOptions>,
    ) -> BoxFuture<'_, parquet::errors::Result<Arc<ParquetMetaData>>> {
        Box::pin(async move {
            let metadata = DFParquetMetadata::new(self.store.as_ref(), &self.meta)
                .with_metadata_size_hint(self.metadata_size_hint)
                .fetch_metadata()
                .await
                .map_err(|e| {
                    ParquetError::General(format!(
                        "AsyncChunkReader::get_metadata error: {e}"
                    ))
                })?;
            Ok(metadata)
        })
    }
}
