//! Shared helpers for the `lines` and `json` file scans.
//!
//! Both formats are file-granularity-only — they decode each file as
//! one logical stream and never split byte ranges within a file. The
//! default DataFusion partitioner (`FileGroupPartitioner`) assigns
//! mid-file byte ranges, which would corrupt our streams. We override
//! `FileSource::repartitioned` with a file-level LPT
//! (longest-processing-time) bucket-packing scheme — distribute whole
//! files across partitions, biggest first, each going to the
//! currently-smallest bucket. Keeps partitions close to equal total
//! bytes without splitting files.
//!
//! Unlike the default impl we also do NOT bail on `is_compressed()`:
//! gzip files can still be distributed across partitions, we just
//! can't split them mid-stream.

use datafusion::common::Result;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::physical_plan::{FileGroup, FileScanConfig};
use datafusion::error::DataFusionError;

/// File-granularity-only repartitioner. `caller` is the format name
/// used in the error message of the (logically unreachable) "no
/// buckets" branch.
pub fn lpt_repartition(
    caller: &'static str,
    target_partitions: usize,
    repartition_file_min_size: usize,
    config: &FileScanConfig,
) -> Result<Option<FileScanConfig>> {
    let flat: Vec<PartitionedFile> = config
        .file_groups
        .iter()
        .flat_map(|g| g.iter().cloned())
        .collect();

    if flat.is_empty() || target_partitions <= 1 {
        return Ok(None);
    }

    let total_size: u64 = flat.iter().map(|f| f.object_meta.size).sum();
    if total_size < repartition_file_min_size as u64 {
        return Ok(None);
    }

    let mut files = flat;
    files.sort_by(|a, b| b.object_meta.size.cmp(&a.object_meta.size));

    let mut buckets: Vec<(u64, Vec<PartitionedFile>)> =
        (0..target_partitions).map(|_| (0u64, Vec::new())).collect();

    for file in files {
        // `buckets` was built above with `target_partitions` entries,
        // and we returned early when `target_partitions <= 1`, so
        // `min_by_key` always sees at least one element. State the
        // invariant as an Internal error rather than .expect() per
        // CLAUDE.md.
        let (idx, _) = buckets
            .iter()
            .enumerate()
            .min_by_key(|(_, (size, _))| *size)
            .ok_or_else(|| {
                DataFusionError::Internal(format!("{caller} LPT packing produced no buckets"))
            })?;
        let size = file.object_meta.size;
        buckets[idx].0 += size;
        buckets[idx].1.push(file);
    }

    // Drop empty buckets (target_partitions > file count). An empty
    // FileGroup confuses downstream planning.
    let groups: Vec<FileGroup> = buckets
        .into_iter()
        .filter_map(|(_, files)| {
            if files.is_empty() {
                None
            } else {
                Some(FileGroup::new(files))
            }
        })
        .collect();

    if groups.len() <= 1 {
        return Ok(None);
    }

    let mut new_config = config.clone();
    new_config.file_groups = groups;
    Ok(Some(new_config))
}
