// Copyright (c) 2023-2026 ParadeDB, Inc.
//
// This file is part of ParadeDB - Postgres for Search and Analytics
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use std::sync::Arc;

use arrow_array::{Array, RecordBatch};
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::physical_plan::ExecutionPlan;
use futures::StreamExt;
use tokio::runtime::Runtime;

use crate::api::HashMap;
use crate::index::fast_fields_helper::{build_arrow_schema, FFHelper, WhichFastField};
use crate::nodecast;
use crate::postgres::customscan::basescan::exec_methods::{ExecMethod, ExecState};
use crate::postgres::customscan::basescan::scan_state::BaseScanState;
use crate::postgres::customscan::builders::custom_path::ExecMethodType;
use crate::postgres::heap::VisibilityChecker;
use crate::postgres::rel::PgSearchRelation;
use crate::postgres::types_arrow::arrow_array_to_datum;
use crate::scan::execution_plan::{PgSearchScanPlan, ScanState};

use pgrx::{pg_sys, IntoDatum, PgOid, PgTupleDesc};

// ============================================================================
// Synchronous stream polling utilities
// ============================================================================

/// Polls a stream for the next item synchronously using a tokio runtime.
/// This properly handles `Poll::Pending` by driving the stream to completion,
/// which is necessary for DataFusion operators that may buffer data across polls.
///
/// Uses the same tokio runtime pattern as JoinScan for consistency.
fn poll_next_sync<S: futures::Stream + Unpin>(
    runtime: &Runtime,
    stream: &mut S,
) -> Option<S::Item> {
    // Check for query cancellation before blocking
    pgrx::check_for_interrupts!();

    // Use tokio runtime to drive the stream (same pattern as JoinScan)
    runtime.block_on(async { stream.next().await })
}

struct Inner {
    heaprel: Option<PgSearchRelation>,
    tupdesc: Option<PgTupleDesc<'static>>,

    /// Execution time WhichFastFields.
    pub which_fast_fields: Vec<WhichFastField>,

    /// Fast field helper wrapped in Arc for sharing with DataFusion plans.
    pub ffhelper: Option<Arc<FFHelper>>,

    pub slot: *mut pg_sys::TupleTableSlot,

    did_query: bool,
}

impl Inner {
    pub fn new(which_fast_fields: Vec<WhichFastField>) -> Self {
        Self {
            heaprel: None,
            tupdesc: None,
            which_fast_fields,
            ffhelper: None,
            slot: std::ptr::null_mut(),
            did_query: false,
        }
    }

    pub fn init(&mut self, state: &mut BaseScanState, cstate: *mut pg_sys::CustomScanState) {
        unsafe {
            self.heaprel = Some(Clone::clone(state.heaprel()));
            self.tupdesc = Some(PgTupleDesc::from_pg_unchecked(
                (*cstate).ss.ps.ps_ResultTupleDesc,
            ));
            self.slot = pg_sys::MakeTupleTableSlot(
                (*cstate).ss.ps.ps_ResultTupleDesc,
                &pg_sys::TTSOpsVirtual,
            );
            // Initialize the fast field helper wrapped in Arc for sharing
            self.ffhelper = Some(Arc::new(FFHelper::with_fields(
                state.search_reader.as_ref().unwrap(),
                &self.which_fast_fields,
            )));
        }
    }

    pub fn reset(&mut self, _state: &mut BaseScanState) {
        self.did_query = false;
    }
}

/// Execution state for columnar retrieval using DataFusion execution.
///
/// This execution state is designed to handle two scenarios:
/// 1. Multiple string fast fields in a single query
/// 2. A mix of string and numeric fast fields in a single query
///
/// The execution method produces data through DataFusion's execution engine,
/// consuming results as Arrow RecordBatches from a DataFusion stream. Segments are processed
/// lazily via PostgreSQL's parallel query infrastructure, with DataFusion producing batches
/// for each segment.
///
/// # Usage Context
/// This execution method is selected when a query uses multiple fast fields with at least one
/// string fast field. It processes both string and numeric fields directly from the index's
/// fast field data structures, avoiding the need to fetch full documents.
///
/// # Feature Flag
/// This execution method is controlled by the `paradedb.enable_columnar_exec` GUC setting.
/// It is disabled by default and can be enabled with:
/// ```sql
/// SET paradedb.enable_columnar_exec = true;
/// ```
pub struct ColumnarExecState {
    /// Core functionality shared with other fast field execution methods
    inner: Inner,

    /// Fast fields to fetch from the scanner (includes projected fields + Ctid + Sort column)
    scanner_fast_fields: Vec<WhichFastField>,

    /// The batch size hint to use for this execution.
    batch_size_hint: Option<usize>,

    /// Tokio runtime for driving async DataFusion streams synchronously.
    /// Created once and reused (same pattern as JoinScan).
    runtime: Option<Runtime>,

    /// The DataFusion stream producing RecordBatches.
    stream: Option<SendableRecordBatchStream>,

    /// The current RecordBatch of fast field values
    current_record_batch: Option<RecordBatch>,
    current_batch_row_idx: usize,

    /// Column index for ctid in the RecordBatch
    ctid_column_idx: Option<usize>,

    /// Statistics tracking the number of visible rows
    num_visible: usize,

    /// Const values extracted from the target list to be projected into the slot
    const_values: HashMap<usize, (pg_sys::Datum, bool)>,
}

/// Populates the target slot with values from a RecordBatch.
///
/// Extracts values from Arrow columns and converts them to PostgreSQL datums.
/// Special handling for ctid and tableoid which are set on the slot directly.
#[allow(clippy::too_many_arguments)]
fn populate_slot_from_record_batch(
    const_values: &HashMap<usize, (pg_sys::Datum, bool)>,
    record_batch: &RecordBatch,
    row_idx: usize,
    which_fast_fields: &[WhichFastField],
    tupdesc: &pgrx::PgTupleDesc,
    slot: &mut pg_sys::TupleTableSlot,
    datums: &mut [pg_sys::Datum],
    isnull: &mut [bool],
) {
    for (i, (att, which_fast_field)) in tupdesc.iter().zip(which_fast_fields).enumerate() {
        let column = record_batch.column(i);

        // Handle Junk columns first (before null check) - they use const_values, not Arrow data
        if matches!(which_fast_field, WhichFastField::Junk(_)) {
            if let Some((val, is_null)) = const_values.get(&i) {
                datums[i] = *val;
                isnull[i] = *is_null;
            } else {
                pgrx::error!(
                    "Expression in target list is not yet supported. \
                        Please file an issue at https://github.com/paradedb/paradedb/issues."
                );
            }
            continue;
        }

        // Check if this column has a null at this row
        if column.is_null(row_idx) {
            // Check for constant values
            if let Some((val, is_null)) = const_values.get(&i) {
                datums[i] = *val;
                isnull[i] = *is_null;
            }
            // Otherwise leave as null (already initialized)
            continue;
        }

        // Handle special fields that don't need datum conversion
        match which_fast_field {
            WhichFastField::Ctid => {
                // ctid is already set on slot.tts_tid before calling this function
                datums[i] = slot.tts_tid.into_datum().unwrap_or(pg_sys::Datum::null());
                isnull[i] = false;
                continue;
            }
            WhichFastField::TableOid => {
                // tableoid is already set on slot.tts_tableOid before calling this function
                datums[i] = slot
                    .tts_tableOid
                    .into_datum()
                    .unwrap_or(pg_sys::Datum::null());
                isnull[i] = false;
                continue;
            }
            _ => {}
        }

        // Extract numeric scale if this is a Numeric64 or NumericBytes field
        let numeric_scale = which_fast_field
            .field_type()
            .and_then(|ft| ft.numeric_scale());

        // Convert Arrow array value to datum
        match arrow_array_to_datum(
            column.as_ref(),
            row_idx,
            PgOid::from(att.atttypid),
            numeric_scale,
        ) {
            Ok(Some(datum)) => {
                datums[i] = datum;
                isnull[i] = false;
            }
            Ok(None) => {
                // Null datum - check for const value
                if let Some((val, is_null)) = const_values.get(&i) {
                    datums[i] = *val;
                    isnull[i] = *is_null;
                }
            }
            Err(e) => {
                // This panic indicates a bug in type mapping between Arrow and PostgreSQL.
                // The schema was computed at planning time, so a mismatch here means
                // either the schema computation is wrong or Arrow returned unexpected data.
                panic!(
                    "BUG: Failed to convert Arrow value to PostgreSQL datum. \
                    Attribute OID: {:?}, Fast field: {which_fast_field:?}, Error: {e}. \
                    This indicates a type mapping bug in the Arrow-to-Postgres conversion.",
                    att.atttypid
                );
            }
        }
    }
}

impl ColumnarExecState {
    /// Creates a new ColumnarExecState from a list of fast fields.
    ///
    /// This constructor analyzes the provided fast fields and categorizes them
    /// into string and numeric types for optimized processing.
    ///
    /// # Arguments
    ///
    /// * `which_fast_fields` - Vector of fast fields that will be processed (projected fields)
    ///
    /// # Returns
    ///
    /// A new ColumnarExecState instance
    pub fn new(which_fast_fields: Vec<WhichFastField>) -> Self {
        // Build scanner fields: projected + [Ctid]
        let mut scanner_fast_fields = which_fast_fields.clone();

        // Ensure Ctid is present (always needed for execution)
        let ctid_column_idx = if let Some(idx) = scanner_fast_fields
            .iter()
            .position(|f| matches!(f, WhichFastField::Ctid))
        {
            Some(idx)
        } else {
            scanner_fast_fields.push(WhichFastField::Ctid);
            Some(scanner_fast_fields.len() - 1)
        };

        // batch_size_hint is computed in `init` from the resolved LIMIT/OFFSET,
        // when EState is available.
        Self {
            inner: Inner::new(which_fast_fields),
            scanner_fast_fields,
            batch_size_hint: None,
            runtime: None,
            stream: None,
            current_record_batch: None,
            current_batch_row_idx: 0,
            ctid_column_idx,
            num_visible: 0,
            const_values: HashMap::default(),
        }
    }

    /// Creates a DataFusion stream for the unsorted path.
    ///
    /// Uses PostgreSQL's lazy segment checkout - one segment at a time.
    /// Each segment is processed through a single-partition [`PgSearchScanPlan`].
    fn create_unsorted_stream(
        &mut self,
        state: &mut BaseScanState,
    ) -> Option<SendableRecordBatchStream> {
        if self.inner.did_query {
            return None;
        }
        self.inner.did_query = true;

        let search_reader = state.search_reader.as_ref().unwrap();
        let index_rel = state.indexrel.as_ref().unwrap();
        let heap_rel = state.heaprel.as_ref().unwrap();
        let ffhelper = self
            .inner
            .ffhelper
            .as_ref()
            .expect("ColumnarExecState: ffhelper should be initialized");

        // Clone visibility checker for the plan
        // TODO: This will cause metrics to be lost for fast field scans: see `impl Clone for VisibilityChecker`.
        let visibility = state
            .visibility_checker
            .as_ref()
            .expect("ColumnarExecState: visibility_checker should be initialized")
            .clone();

        let scanner_config = crate::scan::execution_plan::ScannerConfig {
            which_fast_fields: self.scanner_fast_fields.clone(),
            heap_relid: heap_rel.oid().to_u32(),
            batch_size_hint: self.batch_size_hint,
            // Basescan is never leader-dispatched; mirror the reader's scoring from the fields.
            score_needed: self
                .scanner_fast_fields
                .iter()
                .any(|f| matches!(f, crate::index::fast_fields_helper::WhichFastField::Score)),
        };

        // Create PgSearchScanPlan and execute via DataFusion
        let state_partition = ScanState {
            source_idx: None,
            // Basescan is never an MPP source.
            planner_estimated_rows: 0,
            scanner_config,
            ffhelper: Arc::clone(ffhelper),
            visibility: Box::new(visibility) as Box<VisibilityChecker>,
            reader: search_reader.clone(),
        };

        let plan = PgSearchScanPlan::new(
            Some(state_partition),
            build_arrow_schema(&self.scanner_fast_fields),
            // TODO: Switch to an Arc in the scan state.
            state.search_query_input().clone(),
            None,
            Vec::new(),
            None,
            index_rel.oid().to_u32(),
            None,
            1,
            state.parallel_state(),
            None,
        );

        let task_ctx = Arc::new(TaskContext::default());
        match plan.execute(0, task_ctx) {
            Ok(stream) => Some(stream),
            Err(e) => {
                pgrx::error!("Failed to execute plan: {e}");
            }
        }
    }
}

impl ExecMethod for ColumnarExecState {
    /// Initializes the execution state with the necessary context.
    ///
    /// # Arguments
    ///
    /// * `state` - The current scan state containing query information
    /// * `cstate` - PostgreSQL's custom scan state pointer
    fn init(&mut self, state: &mut BaseScanState, cstate: *mut pg_sys::CustomScanState) {
        // Resolve LIMIT/OFFSET for the batch size hint. `resolve_mut`
        // converts each Param → Static in place, so the LimitOffset stored
        // on ExecMethodType ends up holding resolved values for EXPLAIN.
        if let ExecMethodType::Columnar {
            limit_offset: Some(ref mut lo),
            ..
        } = state.exec_method_type
        {
            unsafe {
                let estate = (*cstate).ss.ps.state;
                if let Some(fetch) = lo.resolve_mut(estate).and_then(|lo| lo.static_fetch()) {
                    self.batch_size_hint = Some(fetch * 2);
                }
            }
        }

        // Initialize the inner FastFieldExecState
        self.inner.init(state, cstate);

        unsafe {
            let targetlist = (*(*cstate).ss.ps.plan).targetlist;
            let len = pg_sys::list_length(targetlist);
            self.const_values.clear();
            self.const_values.reserve(len as usize);

            for i in 0..len {
                let tle = pg_sys::list_nth(targetlist, i) as *mut pg_sys::TargetEntry;
                if !tle.is_null() && !(*tle).expr.is_null() {
                    if let Some(expr) = nodecast!(Const, T_Const, (*tle).expr) {
                        self.const_values
                            .insert(i as usize, ((*expr).constvalue, (*expr).constisnull));
                    }
                }
            }
        }

        // Reset mixed field specific state
        self.stream = None;
        self.current_record_batch = None;
        self.current_batch_row_idx = 0;
        self.num_visible = 0;
    }

    /// Executes the search query and prepares result processing.
    ///
    /// Segments are processed lazily via PostgreSQL's parallel query infrastructure, with each
    /// segment producing its own DataFusion stream.
    ///
    /// # Arguments
    ///
    /// * `state` - The current scan state containing query information
    ///
    /// # Returns
    ///
    /// `true` if there are results to process, `false` otherwise
    fn query(&mut self, state: &mut BaseScanState) -> bool {
        // Create tokio runtime on first use (same pattern as JoinScan).
        // This is a single-threaded runtime used to drive DataFusion's async streams
        // synchronously within PostgreSQL's execution model.
        if self.runtime.is_none() {
            self.runtime = Some(
                tokio::runtime::Builder::new_current_thread()
                    .build()
                    .expect("Failed to create tokio runtime for DataFusion stream execution"),
            );
        }

        loop {
            // Try to get next batch from existing stream
            // Scope the runtime borrow to avoid conflicts with mutable borrows below
            if let Some(stream) = &mut self.stream {
                let runtime = self.runtime.as_ref().unwrap();
                match poll_next_sync(runtime, stream) {
                    Some(Ok(batch)) => {
                        self.current_record_batch = Some(batch);
                        self.current_batch_row_idx = 0;
                        return true;
                    }
                    Some(Err(e)) => {
                        pgrx::error!("Error polling DataFusion stream: {e}");
                    }
                    None => {
                        // Stream exhausted; try to get another stream (next segment).
                        self.stream = None;
                        continue;
                    }
                }
            }

            // Create a new DataFusion stream
            let new_stream = self.create_unsorted_stream(state);

            match new_stream {
                Some(stream) => {
                    self.stream = Some(stream);
                    // Continue loop to poll the new stream
                }
                None => {
                    return false;
                }
            }
        }
    }

    /// Fetches the next result and prepares it for returning to PostgreSQL.
    ///
    /// This method converts DataFusion RecordBatch results into PostgreSQL tuple format,
    /// handling value retrieval for all field types from Arrow columns.
    ///
    /// # Arguments
    ///
    /// * `state` - The current scan state
    ///
    /// # Returns
    ///
    /// The next execution state containing the result or EOF
    fn internal_next(&mut self, _state: &mut BaseScanState) -> ExecState {
        unsafe {
            let record_batch = match self.current_record_batch.as_ref() {
                Some(batch) => batch,
                None => return ExecState::Eof,
            };

            let row_idx = self.current_batch_row_idx;
            if row_idx >= record_batch.num_rows() {
                // This batch is exhausted.
                self.current_record_batch = None;
                return ExecState::Eof;
            }

            self.current_batch_row_idx += 1;

            let heaprel = self
                .inner
                .heaprel
                .as_ref()
                .expect("ColumnarExecState: heaprel should be initialized");
            let slot = self.inner.slot;
            let natts = (*(*slot).tts_tupleDescriptor).natts as usize;

            // Extract ctid from the RecordBatch
            let ctid = if let Some(ctid_idx) = self.ctid_column_idx {
                let ctid_array = record_batch
                    .column(ctid_idx)
                    .as_any()
                    .downcast_ref::<arrow_array::UInt64Array>()
                    .expect("ctid column should be UInt64Array");
                ctid_array.value(row_idx)
            } else {
                panic!("ctid column not found in fast field execution");
            };

            // Set ctid and table OID on the slot
            crate::postgres::utils::u64_to_item_pointer(ctid, &mut (*slot).tts_tid);
            (*slot).tts_tableOid = heaprel.oid();

            // Setup slot for returning data
            (*slot).tts_flags &= !pg_sys::TTS_FLAG_EMPTY as u16;
            (*slot).tts_flags |= pg_sys::TTS_FLAG_SHOULDFREE as u16;
            (*slot).tts_nvalid = natts as _;

            let datums = std::slice::from_raw_parts_mut((*slot).tts_values, natts);
            let isnull = std::slice::from_raw_parts_mut((*slot).tts_isnull, natts);

            // Initialize all values to NULL
            for i in 0..natts {
                datums[i] = pg_sys::Datum::null();
                isnull[i] = true;
            }

            let which_fast_fields = &self.inner.which_fast_fields;
            let tupdesc = self.inner.tupdesc.as_ref().unwrap();
            debug_assert!(natts == which_fast_fields.len());

            populate_slot_from_record_batch(
                &self.const_values,
                record_batch,
                row_idx,
                which_fast_fields,
                tupdesc,
                &mut *slot,
                datums,
                isnull,
            );

            ExecState::Virtual { slot }
        }
    }

    /// Resets the execution state to its initial state.
    ///
    /// # Arguments
    ///
    /// * `state` - The current scan state
    fn reset(&mut self, state: &mut BaseScanState) {
        // Reset inner FastFieldExecState
        self.inner.reset(state);

        // Reset DataFusion stream state
        self.stream = None;
        self.current_record_batch = None;
        self.current_batch_row_idx = 0;

        // Reset statistics
        self.num_visible = 0;
    }

    /// Increments the count of visible rows.
    ///
    /// Called when a row passes visibility checks.
    fn increment_visible(&mut self) {
        self.num_visible += 1;
    }
}
