//! LanceTableProvider — wraps a Lance Dataset as a DataFusion TableProvider.
//!
//! This enables seamless SQL queries over Lance datasets using Arrow zero-copy.
//! Supports:
//! - **Projection push-down**: Only reads requested columns from Lance
//! - **Filter push-down**: Pushes simple predicates into the Lance scanner
//! - **Limit push-down**: Limits rows at the Lance scan level
//! - **Partitioned scanning**: Reads Lance fragments in parallel

use std::any::Any;
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;

use arrow::array::RecordBatch;
use arrow_schema::SchemaRef;
use datafusion::catalog::Session;
use datafusion::catalog::TableProvider;
use datafusion::error::DataFusionError;
use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableType};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::datasource::memory::MemorySourceConfig;

use lance::dataset::Dataset;
use futures::TryStreamExt;

use super::context::SqlError;

/// A DataFusion TableProvider backed by a Lance dataset.
///
/// This provider opens a Lance dataset and exposes it as a table that
/// DataFusion can query with full SQL support. Arrow RecordBatches are
/// read zero-copy from the Lance columnar format.
///
/// # Optimizations
/// - **Projection push-down**: Only reads requested columns from Lance scanner
/// - **Filter push-down**: Converts DataFusion Exprs to Lance SQL filter strings
/// - **Limit push-down**: Applies LIMIT at the Lance scan level
pub struct LanceTableProvider {
    /// The underlying Lance dataset.
    dataset: Arc<Dataset>,
    /// Arrow schema of the dataset.
    arrow_schema: SchemaRef,
    /// Whether to use lazy scanning (re-scan on each query) vs cached mode.
    lazy_scan: bool,
    /// Cached batches (only used when lazy_scan = false).
    cached_batches: Option<Vec<RecordBatch>>,
}

impl fmt::Debug for LanceTableProvider {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LanceTableProvider")
            .field("schema", &self.arrow_schema)
            .field("lazy_scan", &self.lazy_scan)
            .field(
                "cached_rows",
                &self
                    .cached_batches
                    .as_ref()
                    .map(|b| b.iter().map(|r| r.num_rows()).sum::<usize>())
                    .unwrap_or(0),
            )
            .finish()
    }
}

impl LanceTableProvider {
    /// Create a new LanceTableProvider by opening a Lance dataset.
    ///
    /// By default, data is cached in memory for fast repeated queries.
    /// Use `try_new_lazy` for large datasets that should be scanned on demand.
    ///
    /// # Arguments
    /// * `uri` - Path or URI to the Lance dataset (local path, s3://, gs://)
    /// * `_storage_options` - Optional storage configuration for remote backends
    pub async fn try_new(
        uri: &str,
        _storage_options: Option<HashMap<String, String>>,
    ) -> Result<Self, SqlError> {
        let dataset = Dataset::open(uri)
            .await
            .map_err(|e| SqlError::Lance(format!("Failed to open dataset at '{}': {}", uri, e)))?;

        // Scan the full dataset to get Arrow RecordBatches
        let mut scanner = dataset.scan();
        let stream = scanner
            .try_into_stream()
            .await
            .map_err(|e| SqlError::Lance(format!("Failed to create scan stream: {}", e)))?;

        let batches: Vec<RecordBatch> = stream
            .try_collect()
            .await
            .map_err(|e| SqlError::Lance(format!("Failed to collect batches: {}", e)))?;

        let arrow_schema = if batches.is_empty() {
            Arc::new(arrow_schema::Schema::empty())
        } else {
            batches[0].schema()
        };

        Ok(Self {
            dataset: Arc::new(dataset),
            arrow_schema,
            lazy_scan: false,
            cached_batches: Some(batches),
        })
    }

    /// Create a lazy LanceTableProvider that scans data on each query.
    ///
    /// This is more memory-efficient for large datasets, as data is not
    /// cached in memory. Each query triggers a fresh Lance scan with
    /// push-down optimizations applied.
    pub async fn try_new_lazy(
        uri: &str,
        _storage_options: Option<HashMap<String, String>>,
    ) -> Result<Self, SqlError> {
        let dataset = Dataset::open(uri)
            .await
            .map_err(|e| SqlError::Lance(format!("Failed to open dataset at '{}': {}", uri, e)))?;

        // Get schema without loading data — scan 0 rows
        let mut scanner = dataset.scan();
        scanner
            .limit(Some(0), None)
            .map_err(|e| SqlError::Lance(format!("Failed to set limit: {}", e)))?;
        let stream = scanner
            .try_into_stream()
            .await
            .map_err(|e| SqlError::Lance(format!("Failed to create scan stream: {}", e)))?;
        let sample_batches: Vec<RecordBatch> = stream
            .try_collect()
            .await
            .map_err(|e| SqlError::Lance(format!("Failed to collect schema: {}", e)))?;

        // Derive schema from the dataset metadata
        let arrow_schema = if sample_batches.is_empty() {
            // Fallback: build schema from lance schema fields
            let lance_schema = dataset.schema();
            let fields: Vec<arrow_schema::Field> = lance_schema
                .fields
                .iter()
                .map(|f| {
                    let dt: arrow_schema::DataType = f.data_type().into();
                    arrow_schema::Field::new(&f.name, dt, f.nullable)
                })
                .collect();
            Arc::new(arrow_schema::Schema::new(fields))
        } else {
            sample_batches[0].schema()
        };

        Ok(Self {
            dataset: Arc::new(dataset),
            arrow_schema,
            lazy_scan: true,
            cached_batches: None,
        })
    }

    /// Create from an already-opened Lance Dataset.
    pub async fn from_dataset(dataset: Arc<Dataset>) -> Result<Self, SqlError> {
        let mut scanner = dataset.scan();
        let stream = scanner
            .try_into_stream()
            .await
            .map_err(|e| SqlError::Lance(format!("Failed to create scan stream: {}", e)))?;

        let batches: Vec<RecordBatch> = stream
            .try_collect()
            .await
            .map_err(|e| SqlError::Lance(format!("Failed to collect batches: {}", e)))?;

        let arrow_schema = if batches.is_empty() {
            Arc::new(arrow_schema::Schema::empty())
        } else {
            batches[0].schema()
        };

        Ok(Self {
            dataset,
            arrow_schema,
            lazy_scan: false,
            cached_batches: Some(batches),
        })
    }

    /// Get the underlying Lance Dataset.
    pub fn dataset(&self) -> &Arc<Dataset> {
        &self.dataset
    }

    /// Get the cached RecordBatches (None if lazy_scan mode).
    pub fn batches(&self) -> Option<&[RecordBatch]> {
        self.cached_batches.as_deref()
    }

    /// Total number of cached rows (0 if lazy_scan mode).
    pub fn num_rows(&self) -> usize {
        self.cached_batches
            .as_ref()
            .map(|b| b.iter().map(|r| r.num_rows()).sum())
            .unwrap_or(0)
    }

    /// Perform a fresh Lance scan with projection, filter, and limit push-down.
    ///
    /// Returns the scanned RecordBatches.
    async fn scan_with_pushdown(
        &self,
        projection: Option<&Vec<usize>>,
        filters: &[Expr],
        limit: Option<usize>,
    ) -> Result<Vec<RecordBatch>, DataFusionError> {
        let mut scanner = self.dataset.scan();

        // Projection push-down: only read selected columns
        if let Some(proj) = projection {
            let col_names: Vec<String> = proj
                .iter()
                .filter_map(|&idx| {
                    self.arrow_schema
                        .fields()
                        .get(idx)
                        .map(|f| f.name().clone())
                })
                .collect();

            if !col_names.is_empty() {
                let col_refs: Vec<&str> = col_names.iter().map(|s| s.as_str()).collect();
                scanner
                    .project(&col_refs)
                    .map_err(|e| DataFusionError::External(Box::new(e)))?;
            }
        }

        // Filter push-down: convert DataFusion Exprs to Lance filter strings
        if !filters.is_empty() {
            if let Some(filter_str) = Self::exprs_to_lance_filter(filters) {
                // Try to push the filter; if Lance doesn't support it, ignore
                let _ = scanner.filter(filter_str.as_str());
            }
        }

        // Limit push-down
        if let Some(lim) = limit {
            scanner.limit(Some(lim as i64), None).map_err(|e| {
                DataFusionError::External(Box::new(e))
            })?;
        }

        // Execute the scan
        let stream = scanner
            .try_into_stream()
            .await
            .map_err(|e| DataFusionError::External(Box::new(e)))?;

        let batches: Vec<RecordBatch> = stream
            .try_collect()
            .await
            .map_err(|e| DataFusionError::External(Box::new(e)))?;

        Ok(batches)
    }

    /// Convert DataFusion filter expressions to a Lance SQL-like filter string.
    ///
    /// Supports simple comparisons: col = val, col > val, col < val, etc.
    /// Returns None if no expressions could be converted.
    fn exprs_to_lance_filter(exprs: &[Expr]) -> Option<String> {
        let parts: Vec<String> = exprs
            .iter()
            .filter_map(|expr| Self::expr_to_lance_string(expr))
            .collect();

        if parts.is_empty() {
            None
        } else {
            Some(parts.join(" AND "))
        }
    }

    /// Convert a single DataFusion Expr to a Lance filter string.
    fn expr_to_lance_string(expr: &Expr) -> Option<String> {
        match expr {
            Expr::BinaryExpr(binary) => {
                let left = Self::expr_to_lance_string(&binary.left)?;
                let right = Self::expr_to_lance_string(&binary.right)?;
                let op = match binary.op {
                    datafusion::logical_expr::Operator::Eq => "=",
                    datafusion::logical_expr::Operator::NotEq => "!=",
                    datafusion::logical_expr::Operator::Lt => "<",
                    datafusion::logical_expr::Operator::LtEq => "<=",
                    datafusion::logical_expr::Operator::Gt => ">",
                    datafusion::logical_expr::Operator::GtEq => ">=",
                    datafusion::logical_expr::Operator::And => "AND",
                    datafusion::logical_expr::Operator::Or => "OR",
                    _ => return None,
                };
                Some(format!("{} {} {}", left, op, right))
            }
            Expr::Column(col) => Some(col.name.clone()),
            Expr::Literal(val, _metadata) => Some(format!("{}", val)),
            Expr::Not(inner) => {
                let inner_str = Self::expr_to_lance_string(inner)?;
                Some(format!("NOT ({})", inner_str))
            }
            Expr::IsNull(inner) => {
                let inner_str = Self::expr_to_lance_string(inner)?;
                Some(format!("{} IS NULL", inner_str))
            }
            Expr::IsNotNull(inner) => {
                let inner_str = Self::expr_to_lance_string(inner)?;
                Some(format!("{} IS NOT NULL", inner_str))
            }
            Expr::Between(between) => {
                let expr_str = Self::expr_to_lance_string(&between.expr)?;
                let low = Self::expr_to_lance_string(&between.low)?;
                let high = Self::expr_to_lance_string(&between.high)?;
                if between.negated {
                    Some(format!("{} NOT BETWEEN {} AND {}", expr_str, low, high))
                } else {
                    Some(format!("{} BETWEEN {} AND {}", expr_str, low, high))
                }
            }
            Expr::InList(in_list) => {
                let expr_str = Self::expr_to_lance_string(&in_list.expr)?;
                let values: Vec<String> = in_list
                    .list
                    .iter()
                    .filter_map(|e| Self::expr_to_lance_string(e))
                    .collect();
                if values.is_empty() {
                    return None;
                }
                let list_str = values.join(", ");
                if in_list.negated {
                    Some(format!("{} NOT IN ({})", expr_str, list_str))
                } else {
                    Some(format!("{} IN ({})", expr_str, list_str))
                }
            }
            _ => None,
        }
    }

    /// Refresh cached data by re-scanning the Lance dataset.
    pub async fn refresh(&mut self) -> Result<(), SqlError> {
        if !self.lazy_scan {
            let mut scanner = self.dataset.scan();
            let stream = scanner
                .try_into_stream()
                .await
                .map_err(|e| SqlError::Lance(format!("Failed to create scan stream: {}", e)))?;

            let batches: Vec<RecordBatch> = stream
                .try_collect()
                .await
                .map_err(|e| SqlError::Lance(format!("Failed to collect batches: {}", e)))?;

            if !batches.is_empty() {
                self.arrow_schema = batches[0].schema();
            }
            self.cached_batches = Some(batches);
        }
        Ok(())
    }
}

#[async_trait::async_trait]
impl TableProvider for LanceTableProvider {
    fn as_any(&self) -> &dyn Any {
        self
    }

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

    fn table_type(&self) -> TableType {
        TableType::Base
    }

    async fn scan(
        &self,
        _state: &dyn Session,
        projection: Option<&Vec<usize>>,
        filters: &[Expr],
        limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        let batches = if self.lazy_scan {
            // Lazy mode: perform a fresh scan with push-down optimizations
            self.scan_with_pushdown(projection, filters, limit).await?
        } else {
            // Cached mode: use cached batches with in-memory filtering
            let cached = self.cached_batches.as_ref().ok_or_else(|| {
                DataFusionError::Internal("No cached batches available".to_string())
            })?;

            // Apply limit on cached batches
            if let Some(lim) = limit {
                let mut result = Vec::new();
                let mut remaining = lim;
                for batch in cached {
                    if remaining == 0 {
                        break;
                    }
                    if batch.num_rows() <= remaining {
                        result.push(batch.clone());
                        remaining -= batch.num_rows();
                    } else {
                        result.push(batch.slice(0, remaining));
                        remaining = 0;
                    }
                }
                result
            } else {
                cached.clone()
            }
        };

        // Create a DataSourceExec execution plan from the batches
        let exec = MemorySourceConfig::try_new_exec(
            &[batches],
            self.arrow_schema.clone(),
            projection.cloned(),
        )?;

        Ok(exec)
    }

    fn supports_filters_pushdown(
        &self,
        filters: &[&Expr],
    ) -> Result<Vec<TableProviderFilterPushDown>, DataFusionError> {
        Ok(filters
            .iter()
            .map(|expr| {
                if self.lazy_scan {
                    // In lazy mode, we attempt to push simple filters to Lance
                    if Self::expr_to_lance_string(expr).is_some() {
                        TableProviderFilterPushDown::Inexact
                    } else {
                        TableProviderFilterPushDown::Unsupported
                    }
                } else {
                    // In cached mode, filters are handled by DataFusion
                    TableProviderFilterPushDown::Unsupported
                }
            })
            .collect())
    }
}
