// 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 crate::execution::operators::{copy_or_unpack_array, AlignedArrowStreamReader, CopyMode};
use crate::{errors::CometError, execution::planner::TEST_EXEC_CONTEXT_ID};
use arrow::array::{ArrayRef, RecordBatch, RecordBatchOptions};
use arrow::compute::{cast_with_options, CastOptions};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::common::{arrow_datafusion_err, DataFusionError, Result as DataFusionResult};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::metrics::{
    BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time,
};
use datafusion::{
    execution::TaskContext,
    physical_expr::*,
    physical_plan::{ExecutionPlan, *},
};
use futures::Stream;
use itertools::Itertools;
use std::{
    any::Any,
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll},
};

/// `ScanExec` reads batches of data from Spark over the Arrow C Stream Interface. The
/// `input_source` is moved out of the JVM-exported `ArrowArrayStream` at plan-construction time;
/// dropping the reader (when this exec drops) fires the stream's release callback, which closes
/// the JVM-side `ArrowReader` and its `VectorSchemaRoot`.
#[derive(Debug, Clone)]
pub struct ScanExec {
    /// JVM execution-context id used to look up the `JNIEnv` for callbacks.
    pub exec_context_id: i64,
    /// The C Stream Interface reader. `None` only in unit tests that seed input via
    /// `set_input_batch`.
    pub input_source: Option<Arc<Mutex<AlignedArrowStreamReader>>>,
    pub input_source_description: String,
    pub data_types: Vec<DataType>,
    pub schema: SchemaRef,
    /// Used in unit tests to mock the input batch; otherwise written by `pull_next` on each
    /// poll.
    pub batch: Arc<Mutex<Option<InputBatch>>>,
    cache: Arc<PlanProperties>,
    metrics: ExecutionPlanMetricsSet,
    baseline_metrics: BaselineMetrics,
}

impl ScanExec {
    pub fn new(
        exec_context_id: i64,
        input_source: Option<Arc<Mutex<AlignedArrowStreamReader>>>,
        input_source_description: &str,
        data_types: Vec<DataType>,
    ) -> Result<Self, CometError> {
        let metrics_set = ExecutionPlanMetricsSet::default();
        let baseline_metrics = BaselineMetrics::new(&metrics_set, 0);

        // Build schema directly from data types since get_next now always unpacks dictionaries
        let schema = schema_from_data_types(&data_types);

        let cache = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(Arc::clone(&schema)),
            // The partitioning is not important because we are not using DataFusion's
            // query planner or optimizer
            Partitioning::UnknownPartitioning(1),
            EmissionType::Final,
            Boundedness::Bounded,
        ));

        Ok(Self {
            exec_context_id,
            input_source,
            input_source_description: input_source_description.to_string(),
            data_types,
            batch: Arc::new(Mutex::new(None)),
            cache,
            metrics: metrics_set,
            baseline_metrics,
            schema,
        })
    }

    /// Unpack all dictionary types because some DataFusion operators
    /// and expressions do not support dictionary types
    fn unpack_dictionary_type(dt: &DataType) -> DataType {
        if let DataType::Dictionary(_, vt) = dt {
            // return the underlying data type
            return vt.as_ref().clone();
        }
        dt.clone()
    }

    /// Feeds input batch into this `Scan`. Only used in unit test.
    pub fn set_input_batch(&mut self, input: InputBatch) {
        *self.batch.try_lock().unwrap() = Some(input);
    }

    /// Pull next input batch from the upstream `ArrowArrayStreamReader`.
    pub fn get_next_batch(&mut self) -> Result<(), CometError> {
        if self.input_source.is_none() {
            // This is a unit test. Input batches are seeded via `set_input_batch`.
            return Ok(());
        }
        let mut timer = self.baseline_metrics.elapsed_compute().timer();

        let mut current_batch = self.batch.try_lock().unwrap();
        if current_batch.is_none() {
            let next_batch =
                ScanExec::pull_next(self.exec_context_id, self.input_source.as_ref().unwrap())?;
            *current_batch = Some(next_batch);
        }

        timer.stop();

        Ok(())
    }

    /// Pull the next `RecordBatch` from the stream and convert it to an `InputBatch`. Dictionary
    /// columns are unpacked because Comet's downstream operators do not handle them.
    fn pull_next(
        exec_context_id: i64,
        reader: &Arc<Mutex<AlignedArrowStreamReader>>,
    ) -> Result<InputBatch, CometError> {
        if exec_context_id == TEST_EXEC_CONTEXT_ID {
            // Unit test path; input batches are seeded directly.
            return Ok(InputBatch::EOF);
        }

        // The `Mutex` is for interior mutability (`next` needs `&mut`, but the exec holds the
        // reader behind an `Arc`); access is already serialized by the `self.batch` lock held in
        // `get_next_batch`, so a contended `try_lock` here would signal a caller bug, not races.
        let mut reader = reader
            .try_lock()
            .map_err(|_| CometError::Internal("AlignedArrowStreamReader contended".to_string()))?;

        let next = reader.next();
        match next {
            None => Ok(InputBatch::EOF),
            Some(Err(e)) => Err(CometError::from(e)),
            Some(Ok(record_batch)) => {
                let num_rows = record_batch.num_rows();
                let columns = record_batch.columns();
                let mut inputs: Vec<ArrayRef> = Vec::with_capacity(columns.len());
                for col in columns {
                    inputs.push(copy_or_unpack_array(col, &CopyMode::UnpackOrClone)?);
                }
                Ok(InputBatch::new(inputs, Some(num_rows)))
            }
        }
    }
}

fn schema_from_data_types(data_types: &[DataType]) -> SchemaRef {
    let fields = data_types
        .iter()
        .enumerate()
        .map(|(idx, dt)| {
            let datatype = ScanExec::unpack_dictionary_type(dt);
            // We don't use the field name. Put a placeholder.
            Field::new(format!("col_{idx}"), datatype, true)
        })
        .collect::<Vec<Field>>();

    Arc::new(Schema::new(fields))
}

impl ExecutionPlan for ScanExec {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![]
    }

    fn with_new_children(
        self: Arc<Self>,
        _: Vec<Arc<dyn ExecutionPlan>>,
    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
        Ok(self)
    }

    fn execute(
        &self,
        partition: usize,
        _: Arc<TaskContext>,
    ) -> datafusion::common::Result<SendableRecordBatchStream> {
        Ok(Box::pin(ScanStream::new(
            self.clone(),
            self.schema(),
            partition,
            self.baseline_metrics.clone(),
        )))
    }

    fn properties(&self) -> &Arc<PlanProperties> {
        &self.cache
    }

    fn name(&self) -> &str {
        "ScanExec"
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metrics.clone_inner())
    }
}

impl DisplayAs for ScanExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(f, "ScanExec: source=[{}], ", self.input_source_description)?;
                let fields: Vec<String> = self
                    .data_types
                    .iter()
                    .enumerate()
                    .map(|(idx, dt)| format!("col_{idx:}: {dt:}"))
                    .collect();
                write!(f, "schema=[{}]", fields.join(", "))?;
            }
            DisplayFormatType::TreeRender => unimplemented!(),
        }
        Ok(())
    }
}

/// A async-stream feeds input batch from `Scan` into DataFusion physical plan.
struct ScanStream<'a> {
    /// The `Scan` node producing input batches
    scan: ScanExec,
    /// Schema representing the data
    schema: SchemaRef,
    /// Metrics
    baseline_metrics: BaselineMetrics,
    /// Cast options
    cast_options: CastOptions<'a>,
    /// elapsed time for casting columns to different data types during scan
    cast_time: Time,
}

impl ScanStream<'_> {
    pub fn new(
        scan: ScanExec,
        schema: SchemaRef,
        partition: usize,
        baseline_metrics: BaselineMetrics,
    ) -> Self {
        let cast_time = MetricBuilder::new(&scan.metrics).subset_time("cast_time", partition);
        Self {
            scan,
            schema,
            baseline_metrics,
            cast_options: CastOptions::default(),
            cast_time,
        }
    }

    fn build_record_batch(
        &self,
        columns: &[ArrayRef],
        num_rows: usize,
    ) -> DataFusionResult<RecordBatch, DataFusionError> {
        let schema_fields = self.schema.fields();
        assert_eq!(columns.len(), schema_fields.len());

        // Cast dictionary-encoded primitive arrays to regular arrays and cast
        // Utf8/LargeUtf8/Binary arrays to dictionary-encoded if the schema is
        // defined as dictionary-encoded and the data in this batch is not
        // dictionary-encoded (could also be the other way around)
        let new_columns: Vec<ArrayRef> = columns
            .iter()
            .zip(schema_fields.iter())
            .map(|(column, f)| {
                if column.data_type() != f.data_type() {
                    let mut timer = self.cast_time.timer();
                    let cast_array = cast_with_options(column, f.data_type(), &self.cast_options);
                    timer.stop();
                    cast_array
                } else {
                    Ok(Arc::clone(column))
                }
            })
            .collect::<Result<Vec<_>, _>>()?;
        let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
        RecordBatch::try_new_with_options(Arc::clone(&self.schema), new_columns, &options)
            .map_err(|e| arrow_datafusion_err!(e))
    }
}

impl Stream for ScanStream<'_> {
    type Item = DataFusionResult<RecordBatch>;

    fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut timer = self.baseline_metrics.elapsed_compute().timer();
        let mut scan_batch = self.scan.batch.try_lock().unwrap();

        let input_batch = &*scan_batch;
        let input_batch = if let Some(batch) = input_batch {
            batch
        } else {
            timer.stop();
            return Poll::Pending;
        };

        let result = match input_batch {
            InputBatch::EOF => Poll::Ready(None),
            InputBatch::Batch(columns, num_rows) => {
                self.baseline_metrics.record_output(*num_rows);
                let maybe_batch = self.build_record_batch(columns, *num_rows);
                Poll::Ready(Some(maybe_batch))
            }
        };

        *scan_batch = None;

        timer.stop();

        result
    }
}

impl RecordBatchStream for ScanStream<'_> {
    /// Get the schema
    fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }
}

#[derive(Clone, Debug)]
pub enum InputBatch {
    /// The end of input batches.
    EOF,

    /// A normal batch with columns and a number of rows.
    /// It is possible to have a zero-column batch with a non-zero number of rows,
    /// i.e. reading empty schema from scan.
    Batch(Vec<ArrayRef>, usize),
}

impl InputBatch {
    /// Constructs an ` InputBatch ` from columns and an optional number of rows.
    /// If `num_rows` is none, this function will calculate it from given
    /// columns.
    pub fn new(columns: Vec<ArrayRef>, num_rows: Option<usize>) -> Self {
        let num_rows = num_rows.unwrap_or_else(|| {
            let lengths = columns.iter().map(|a| a.len()).unique().collect::<Vec<_>>();
            assert!(lengths.len() <= 1, "Columns have different lengths.");

            if lengths.is_empty() {
                // All are scalar values
                1
            } else {
                lengths[0]
            }
        });

        InputBatch::Batch(columns, num_rows)
    }
}
