//! The read primitive: full table → Arrow in one call.

use std::sync::Arc;

use arrow_array::RecordBatch;
use arrow_schema::SchemaRef as ArrowSchemaRef;
use futures::{StreamExt, TryStreamExt};
use iceberg::expr::{Predicate, Reference};
use iceberg::spec::Datum;
use iceberg::table::Table as IceTable;

use crate::error::{Result, SkadeError};

/// The Arrow schema (with Iceberg field-id metadata) of `table`'s current
/// schema — the schema batches must carry to be written, and the schema scans
/// come back with.
pub fn arrow_schema_of(table: &IceTable) -> Result<ArrowSchemaRef> {
    let s = iceberg::arrow::schema_to_arrow_schema(table.metadata().current_schema())?;
    Ok(Arc::new(s))
}

/// Full-scan `table` (current snapshot, all columns) into Arrow record batches:
/// `scan().select_all().build()?.to_arrow()` as one ergonomic call.
pub async fn read_all(table: &IceTable) -> Result<Vec<RecordBatch>> {
    let stream = table.scan().select_all().build()?.to_arrow().await?;
    Ok(stream.try_collect().await?)
}

/// Full-scan but **project only `columns`** — column pruning pushed into the
/// scan, so the engine reads only those columns from Parquet (less I/O + decode
/// than [`read_all`]). Empty `columns` is equivalent to [`read_all`].
pub async fn read_columns(table: &IceTable, columns: &[&str]) -> Result<Vec<RecordBatch>> {
    if columns.is_empty() {
        return read_all(table).await;
    }
    let stream = table
        .scan()
        .select(columns.iter())
        .build()?
        .to_arrow()
        .await?;
    Ok(stream.try_collect().await?)
}

/// A backend-neutral scalar value used in a [`ScanFilter`] predicate.
///
/// This is deliberately a tiny, engine-agnostic set (string / signed-int /
/// float / bool) — *not* an iceberg `Datum` — so the predicate surface does not
/// leak the Iceberg engine into skade's public API. A future non-Iceberg
/// backend (e.g. DuckDB) can implement the same [`ScanFilter`] contract by
/// mapping these scalars onto its own predicate type. The Iceberg backend maps
/// each variant onto an `iceberg::spec::Datum` in [`Scalar::to_datum`].
#[derive(Debug, Clone, PartialEq)]
pub enum Scalar {
    /// A UTF-8 string value (the common case — partition columns like `repo`).
    Str(String),
    /// A signed 64-bit integer (`Datum::long`).
    I64(i64),
    /// A signed 32-bit integer (`Datum::int`).
    I32(i32),
    /// A 64-bit float (`Datum::double`).
    F64(f64),
    /// A boolean (`Datum::bool`).
    Bool(bool),
}

impl Scalar {
    /// Map this backend-neutral scalar onto the Iceberg engine's `Datum`.
    fn to_datum(&self) -> Datum {
        match self {
            Scalar::Str(s) => Datum::string(s),
            Scalar::I64(v) => Datum::long(*v),
            Scalar::I32(v) => Datum::int(*v),
            Scalar::F64(v) => Datum::double(*v),
            Scalar::Bool(v) => Datum::bool(*v),
        }
    }
}

impl From<&str> for Scalar {
    fn from(s: &str) -> Self {
        Scalar::Str(s.to_string())
    }
}
impl From<String> for Scalar {
    fn from(s: String) -> Self {
        Scalar::Str(s)
    }
}
impl From<i64> for Scalar {
    fn from(v: i64) -> Self {
        Scalar::I64(v)
    }
}
impl From<i32> for Scalar {
    fn from(v: i32) -> Self {
        Scalar::I32(v)
    }
}
impl From<f64> for Scalar {
    fn from(v: f64) -> Self {
        Scalar::F64(v)
    }
}
impl From<bool> for Scalar {
    fn from(v: bool) -> Self {
        Scalar::Bool(v)
    }
}

/// A backend-neutral scan predicate pushed **into** the read path so the engine
/// prunes data files / row-groups instead of reading them.
///
/// Kept intentionally small — `column <op> value` on a single column, which
/// covers every pushdown nornir's warehouse does today (filter by the `repo`
/// partition column). It is **not** the Iceberg `Predicate` type: that would
/// leak the engine into the public signature. A future backend (DuckDB, …)
/// implements the same enum by mapping it onto its own predicate. The Iceberg
/// backend lowers each variant to an `iceberg::expr::Predicate` in
/// [`ScanFilter::to_iceberg`].
///
/// Use the constructors ([`ScanFilter::eq`], [`ScanFilter::is_in`],
/// [`ScanFilter::range`]) for readable call sites:
/// `ScanFilter::eq("repo", "znippy")`.
#[derive(Debug, Clone, PartialEq)]
pub enum ScanFilter {
    /// `column == value` (today: `repo == "<repo>"`, partition pruning).
    Eq { column: String, value: Scalar },
    /// `column IN (values…)` — matches any of the values.
    In {
        column: String,
        values: Vec<Scalar>,
    },
    /// A range on a column: `lo <= column` and/or `column <= hi` (inclusive
    /// bounds). Either bound may be absent (`None`) for a one-sided range.
    Range {
        column: String,
        lo: Option<Scalar>,
        hi: Option<Scalar>,
    },
}

impl ScanFilter {
    /// `column == value`. The common partition-pruning case.
    pub fn eq(column: impl Into<String>, value: impl Into<Scalar>) -> Self {
        ScanFilter::Eq {
            column: column.into(),
            value: value.into(),
        }
    }

    /// `column IN (values…)`.
    pub fn is_in<S: Into<Scalar>>(
        column: impl Into<String>,
        values: impl IntoIterator<Item = S>,
    ) -> Self {
        ScanFilter::In {
            column: column.into(),
            values: values.into_iter().map(Into::into).collect(),
        }
    }

    /// An inclusive range `lo <= column <= hi`; pass `None` for an open bound.
    pub fn range(
        column: impl Into<String>,
        lo: Option<impl Into<Scalar>>,
        hi: Option<impl Into<Scalar>>,
    ) -> Self {
        ScanFilter::Range {
            column: column.into(),
            lo: lo.map(Into::into),
            hi: hi.map(Into::into),
        }
    }

    /// Lower this backend-neutral filter onto the Iceberg engine's `Predicate`,
    /// which the scan builder pushes down (file/row-group pruning). Returns
    /// `None` for an empty range (no bounds) — caller skips the filter.
    fn to_iceberg(&self) -> Option<Predicate> {
        match self {
            ScanFilter::Eq { column, value } => {
                Some(Reference::new(column).equal_to(value.to_datum()))
            }
            ScanFilter::In { column, values } => {
                if values.is_empty() {
                    // `IN ()` matches nothing. iceberg's set predicate needs ≥1
                    // literal, and leaving the scan *unfiltered* would wrongly
                    // match everything — so emit an always-false predicate the
                    // engine can evaluate: `col IS NULL AND col IS NOT NULL`.
                    return Some(
                        Reference::new(column)
                            .is_null()
                            .and(Reference::new(column).is_not_null()),
                    );
                }
                Some(Reference::new(column).is_in(values.iter().map(|v| v.to_datum())))
            }
            ScanFilter::Range { column, lo, hi } => {
                let lo_pred = lo
                    .as_ref()
                    .map(|v| Reference::new(column).greater_than_or_equal_to(v.to_datum()));
                let hi_pred = hi
                    .as_ref()
                    .map(|v| Reference::new(column).less_than_or_equal_to(v.to_datum()));
                match (lo_pred, hi_pred) {
                    (Some(l), Some(h)) => Some(l.and(h)),
                    (Some(l), None) => Some(l),
                    (None, Some(h)) => Some(h),
                    (None, None) => None,
                }
            }
        }
    }
}

/// **Filtered / pushdown read.** Scan `table` but push `filter` into the scan
/// planner so the engine prunes other partitions' data files / row-groups
/// instead of reading them, and project only `columns` (empty = all columns,
/// preserving column order so positional downcasts stay valid).
///
/// This is the public, backend-neutral equivalent of reaching through to the
/// raw Iceberg scan builder (`table.scan().with_filter(Reference::new("repo")
/// .equal_to(Datum::string(r)))`) — the predicate is expressed as the engine-
/// agnostic [`ScanFilter`], so the public signature never names an Iceberg type.
///
/// Pushdown prunes at file/row-group granularity, **not** per row: the result
/// may still contain rows the predicate would reject, so a caller that needs
/// exact filtering must keep a residual per-row guard. (This matches Iceberg's
/// partition-pruning semantics — pruning is a read-amplification optimization,
/// not a row filter.)
pub async fn read_filtered(
    table: &IceTable,
    filter: &ScanFilter,
    columns: &[&str],
) -> Result<Vec<RecordBatch>> {
    let mut builder = table.scan();
    if let Some(pred) = filter.to_iceberg() {
        builder = builder.with_filter(pred);
    }
    if columns.is_empty() {
        builder = builder.select_all();
    } else {
        builder = builder.select(columns.iter());
    }
    let stream = builder.build()?.to_arrow().await?;
    Ok(stream.try_collect().await?)
}

/// Plan-time data-skipping stats: how many data files / file-splits / rows
/// SURVIVE pruning (Iceberg manifest min/max + partition pruning) for a scan,
/// computed from the scan plan WITHOUT reading any data.
///
/// This is the observable headline metric for the data-skipping win — it counts
/// what the engine will actually OPEN, not a guess. Diff a filtered plan against
/// the unfiltered baseline to get files SKIPPED:
/// `files_skipped = plan_stats(t, None).data_files − plan_stats(t, Some(f)).data_files`.
///
/// Row-group / bloom skipping happens later, inside the Parquet reader, and is
/// NOT visible at plan time — measure that via wall time + the file's row-group
/// count (Parquet metadata), not here.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ScanPlanStats {
    /// Distinct data files the plan will open after pruning.
    pub data_files: u64,
    /// File-scan tasks (one file may split into several) — the unit of scan work.
    pub splits: u64,
    /// Rows the plan will read (sum of the planned files' record counts).
    pub rows_planned: u64,
}

/// Plan (do NOT execute) a scan and report its [`ScanPlanStats`]. Pushes
/// `filter` into the planner exactly like [`read_filtered`]; `filter = None`
/// plans the whole table (the baseline to diff a filtered plan against).
pub async fn plan_stats(table: &IceTable, filter: Option<&ScanFilter>) -> Result<ScanPlanStats> {
    let mut builder = table.scan().select_all();
    if let Some(pred) = filter.and_then(|f| f.to_iceberg()) {
        builder = builder.with_filter(pred);
    }
    let tasks: Vec<_> = builder.build()?.plan_files().await?.try_collect().await?;
    let mut files = std::collections::HashSet::new();
    let mut rows_planned = 0u64;
    for t in &tasks {
        files.insert(t.data_file_path.clone());
        rows_planned += t.record_count.unwrap_or(0);
    }
    Ok(ScanPlanStats {
        data_files: files.len() as u64,
        splits: tasks.len() as u64,
        rows_planned,
    })
}

/// **Limit / early-break streaming read.** Scan `table` but **stop once
/// `max_rows` rows are in hand** instead of materializing the whole table the
/// way [`read_all`] does (it `try_collect`s every data file before the caller
/// can truncate).
///
/// It drives the Iceberg scan's Arrow **stream** and breaks as soon as it has
/// accumulated `max_rows` rows; dropping the stream cancels the rest of the
/// scan, so the remaining data files are never opened or decoded. On a big
/// table (e.g. a multi-GB preview) this turns a multi-second full scan into a
/// few-millisecond read of the first data file(s). A small `with_batch_size`
/// keeps the first batch from overshooting `max_rows` by much.
///
/// Returns *at most* `max_rows` worth of rows — actually the smallest whole
/// number of batches whose row total reaches `max_rows` (the last batch may
/// carry the count slightly over `max_rows`; the caller truncates if it needs
/// an exact count). `max_rows == 0` falls back to a full [`read_all`].
///
/// Caveat: Iceberg gives no ordering guarantee, so this returns *some*
/// `max_rows` rows, not a deterministic top-N — fine for previews (which never
/// promised an order) but not for ranked reads.
pub async fn read_limited(table: &IceTable, max_rows: usize) -> Result<Vec<RecordBatch>> {
    if max_rows == 0 {
        return read_all(table).await;
    }
    // Cap per-batch rows so the stream can stop close to `max_rows` rather than
    // pulling one giant batch. Clamp to a sane floor for tiny limits.
    let batch_size = max_rows.clamp(256, 8192);
    let mut stream = table
        .scan()
        .select_all()
        .with_batch_size(Some(batch_size))
        .build()?
        .to_arrow()
        .await?;
    let mut batches = Vec::new();
    let mut have = 0usize;
    while have < max_rows {
        match stream.next().await {
            Some(b) => {
                let b = b?;
                have += b.num_rows();
                batches.push(b);
            }
            None => break, // table smaller than the limit
        }
    }
    // Dropping `stream` here cancels the rest of the scan: no further parquet
    // files are opened or decoded.
    Ok(batches)
}

/// Full-scan `table` and count rows without materializing the batches.
pub async fn scan_count(table: &IceTable) -> Result<u64> {
    let stream = table.scan().select_all().build()?.to_arrow().await?;
    Ok(stream
        .try_fold(0u64, |acc, batch| async move {
            Ok(acc + batch.num_rows() as u64)
        })
        .await?)
}

/// One equality-delete file added in a delta window: the file path plus the
/// Iceberg field-ids that define row equality (`equality_ids`). The delete
/// applies to any row whose values in those fields match a row in the file —
/// for a graph that means "remove the node(s) with this identity".
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EqualityDeleteFile {
    /// Absolute path of the equality-delete parquet file.
    pub path: String,
    /// Iceberg field-ids that define equality (the identity columns).
    pub equality_ids: Vec<i32>,
}

/// What a [`read_delta`] resolved to: the new snapshots and the data files they
/// appended. Returned so a caller can see/assert *what* the delta was (which
/// snapshots advanced, which files were read) without re-deriving it.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct DeltaPlan {
    /// Snapshot ids strictly after `from`, up to and including `to`, oldest-first.
    pub snapshots: Vec<i64>,
    /// Absolute data-file paths added by those snapshots (the rows to read).
    pub added_files: Vec<String>,
    /// Equality-delete files added in the window — the rows they name are
    /// **removed** (CDC deletes). Read with [`read_equality_deletes`].
    pub delete_files: Vec<EqualityDeleteFile>,
    /// `true` when a commit in the window removes rows in a way the append +
    /// equality-delete delta can't express — a **position**-delete file
    /// (deletes by file+row-offset, not by identity) or a non-`Append`/non-
    /// `Delete` operation (overwrite/replace). The caller should re-read the
    /// whole `to` snapshot instead. An `Append` (inserts) or a `Delete`
    /// operation that only adds *equality*-delete files does **not** set this.
    pub needs_full_reload: bool,
}

/// **Incremental read.** The data files appended to `table` by the snapshots
/// strictly after `from` up to (and including) `to`, read into Arrow batches —
/// *only* those files, not the whole table.
///
/// iceberg-rust 0.9.1 exposes no incremental/changelog scan (only time-travel
/// `snapshot_id()`), so this is built from the manifest layer: walk the snapshot
/// lineage (`parent_snapshot_id`) from `to` back to `from`, then for each new
/// snapshot take the **data** manifests it added (`ManifestFile.added_snapshot_id`)
/// and their `Added` entries — those `DataFile`s are exactly the appended rows.
///
/// `from = None` means "from the beginning" (every data file up to `to`).
///
/// **CDC-aware.** It sees inserts (appended data files) **and** identity-based
/// deletes (equality-delete files added by `Delete` snapshots): the returned
/// [`DeltaPlan::delete_files`] names them, and [`read_equality_deletes`] reads
/// the deleted-row identities. It does **not** handle *position* deletes (delete
/// by file+row-offset) or overwrite/replace commits — those set
/// `DeltaPlan.needs_full_reload` and return **no** batches (the caller does a
/// full [`read_all`] of `to` instead).
pub async fn read_delta(
    table: &IceTable,
    from: Option<i64>,
    to: i64,
) -> Result<(Vec<RecordBatch>, DeltaPlan)> {
    use iceberg::spec::{
        DataContentType, Manifest, ManifestContentType, ManifestStatus, Operation,
    };
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

    let metadata = table.metadata();
    let file_io = table.file_io();
    let mut plan = DeltaPlan::default();
    if from == Some(to) {
        return Ok((Vec::new(), plan)); // already caught up
    }

    // Lineage: to → parent → … until we hit `from` (exclusive) or the root.
    let mut chain: Vec<i64> = Vec::new();
    let mut cur = Some(to);
    while let Some(id) = cur {
        if Some(id) == from {
            break;
        }
        let snap = metadata
            .snapshot_by_id(id)
            .ok_or_else(|| SkadeError::Other(format!("snapshot {id} not in lineage")))?;
        // Append (inserts) and Delete (equality-delete files) are both
        // expressible as a CDC delta. Overwrite/Replace rewrite rows in place,
        // which an additive delta can't follow — fall back to a full re-read.
        if !matches!(
            snap.summary().operation,
            Operation::Append | Operation::Delete
        ) {
            plan.needs_full_reload = true;
        }
        chain.push(id);
        cur = snap.parent_snapshot_id();
    }
    chain.reverse();
    plan.snapshots = chain.clone();
    if plan.needs_full_reload {
        return Ok((Vec::new(), plan));
    }
    let want: std::collections::HashSet<i64> = chain.into_iter().collect();
    if want.is_empty() {
        return Ok((Vec::new(), plan));
    }

    // The target's manifest list names every live manifest; each carries
    // `added_snapshot_id`. Take the manifests added in our window — both Data
    // (inserts) and Deletes (equality/position deletes).
    let target = metadata
        .snapshot_by_id(to)
        .ok_or_else(|| SkadeError::Other(format!("target snapshot {to} not found")))?;
    let manifest_list = target.load_manifest_list(file_io, metadata).await?;
    for mf in manifest_list.entries() {
        if !want.contains(&mf.added_snapshot_id) {
            continue;
        }
        let bytes = file_io.new_input(&mf.manifest_path)?.read().await?;
        let manifest = Manifest::parse_avro(&bytes)?;
        for entry in manifest.entries() {
            if entry.status() != ManifestStatus::Added {
                continue;
            }
            let df = entry.data_file();
            match df.content_type() {
                DataContentType::Data if mf.content == ManifestContentType::Data => {
                    plan.added_files.push(df.file_path().to_string());
                }
                DataContentType::EqualityDeletes => {
                    plan.delete_files.push(EqualityDeleteFile {
                        path: df.file_path().to_string(),
                        equality_ids: df.equality_ids().unwrap_or_default(),
                    });
                }
                // Position deletes name rows by (file, offset), not identity —
                // a graph upsert keyed on a business id can't resolve them, so
                // re-read the whole snapshot to stay correct.
                DataContentType::PositionDeletes => {
                    plan.needs_full_reload = true;
                }
                _ => {}
            }
        }
    }
    if plan.needs_full_reload {
        plan.added_files.clear();
        plan.delete_files.clear();
        return Ok((Vec::new(), plan));
    }
    plan.added_files.sort();
    plan.added_files.dedup();
    plan.delete_files.sort_by(|a, b| a.path.cmp(&b.path));
    plan.delete_files.dedup();

    // Read only those data files, through Iceberg's FileIO.
    let mut batches = Vec::new();
    for path in &plan.added_files {
        let bytes = file_io.new_input(path)?.read().await?;
        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes)?.build()?;
        for b in reader {
            batches.push(b?);
        }
    }
    Ok((batches, plan))
}

/// Read the deleted-row identities out of the equality-delete files in a
/// [`DeltaPlan`]. An equality-delete file is a parquet file whose columns are
/// the equality fields; each row names rows to delete (any data row whose
/// equality-field values match). Returns one `(equality_ids, batches)` pair per
/// file so the caller knows which columns define the deleted identity.
pub async fn read_equality_deletes(
    table: &IceTable,
    plan: &DeltaPlan,
) -> Result<Vec<(Vec<i32>, Vec<RecordBatch>)>> {
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    let file_io = table.file_io();
    let mut out = Vec::with_capacity(plan.delete_files.len());
    for d in &plan.delete_files {
        let bytes = file_io.new_input(&d.path)?.read().await?;
        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes)?.build()?;
        let mut batches = Vec::new();
        for b in reader {
            batches.push(b?);
        }
        out.push((d.equality_ids.clone(), batches));
    }
    Ok(out)
}
