//! Windowed, thread-parallel Parquet reading.
//!
//! The building blocks for bounded-memory bulk ingest (billions of rows):
//! [`parquet_layout`] inspects a file without reading data, [`rowgroup_windows`]
//! groups its row groups into ~`window_rows`-sized windows, and
//! [`read_row_groups`] decodes one window across N OS threads. A streaming
//! ingest reads a window, appends it, frees it, then reads the next — peak
//! memory is one window, not the whole file.
//!
//! (Lifted from nornir-catalog's bench data-plane, where it feeds the OSM /
//! GeoParquet ingest benchmarks.)

use std::path::Path;

use arrow_array::RecordBatch;
use arrow_schema::SchemaRef as ArrowSchemaRef;
use iceberg::spec::Schema;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

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

/// Open a Parquet file and return (derived Iceberg schema, the file's Arrow
/// schema, per-row-group row counts) — without reading any data.
pub fn parquet_layout(path: impl AsRef<Path>) -> Result<(Schema, ArrowSchemaRef, Vec<usize>)> {
    let path = path.as_ref();
    let file = std::fs::File::open(path)
        .map_err(|e| SkadeError::Other(format!("open parquet {}: {e}", path.display())))?;
    let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
    let ice = arrow_to_iceberg(builder.schema())?;
    let file_schema: ArrowSchemaRef = builder.schema().clone();
    let meta = builder.metadata();
    let rg_rows: Vec<usize> = (0..meta.num_row_groups())
        .map(|i| meta.row_group(i).num_rows() as usize)
        .collect();
    Ok((ice, file_schema, rg_rows))
}

/// Group row-group indices into windows whose row counts each sum to about
/// `window_rows`, stopping once `max_rows` total is covered. Returns one `Vec`
/// of row-group indices per window — feed each to [`read_row_groups`].
pub fn rowgroup_windows(rg_rows: &[usize], window_rows: usize, max_rows: usize) -> Vec<Vec<usize>> {
    let win = window_rows.max(1);
    let mut windows: Vec<Vec<usize>> = Vec::new();
    let mut cur: Vec<usize> = Vec::new();
    let (mut cur_rows, mut total) = (0usize, 0usize);
    for (i, &rows) in rg_rows.iter().enumerate() {
        if total >= max_rows {
            break;
        }
        cur.push(i);
        cur_rows += rows;
        total += rows;
        if cur_rows >= win {
            windows.push(std::mem::take(&mut cur));
            cur_rows = 0;
        }
    }
    if !cur.is_empty() {
        windows.push(cur);
    }
    windows
}

/// Read exactly the row groups in `rgs`, decoded in `rows_per_batch` batches,
/// parallel across up to `threads` OS threads (Parquet page decompression is
/// single-threaded per reader, so a many-row-group file otherwise pins one
/// core). Batches keep the file's own Arrow schema ([`crate::recast`] them to a
/// table's field-id schema before ingest). Batch order across threads is not
/// preserved.
pub fn read_row_groups(
    path: impl AsRef<Path>,
    rgs: &[usize],
    rows_per_batch: usize,
    threads: usize,
) -> Result<Vec<RecordBatch>> {
    let path = path.as_ref();
    if rgs.is_empty() {
        return Ok(Vec::new());
    }
    let nthreads = threads.max(1).min(rgs.len());
    // Round-robin row groups across threads for even load.
    let mut buckets: Vec<Vec<usize>> = vec![Vec::new(); nthreads];
    for (k, rg) in rgs.iter().copied().enumerate() {
        buckets[k % nthreads].push(rg);
    }
    let rpb = rows_per_batch.max(1);
    let mut handles = Vec::with_capacity(nthreads);
    for bucket in buckets.into_iter().filter(|b| !b.is_empty()) {
        let path = path.to_path_buf();
        handles.push(std::thread::spawn(move || -> Result<Vec<RecordBatch>> {
            let f = std::fs::File::open(&path)?;
            let rdr = ParquetRecordBatchReaderBuilder::try_new(f)?
                .with_row_groups(bucket)
                .with_batch_size(rpb)
                .build()?;
            rdr.collect::<std::result::Result<Vec<_>, _>>()
                .map_err(SkadeError::from)
        }));
    }
    let mut out = Vec::new();
    for h in handles {
        out.extend(
            h.join()
                .map_err(|_| SkadeError::Other("parquet read thread panicked".into()))??,
        );
    }
    Ok(out)
}
