//! The writer stack as one call: ParquetWriter → RollingFileWriter →
//! DataFileWriter → `fast_append` commit.

use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use arrow_array::{Array, LargeStringArray, RecordBatch, StringArray};
use arrow_schema::SchemaRef as ArrowSchemaRef;
use bytes::Bytes;
use iceberg::spec::{
    DataContentType, DataFile, DataFileBuilder, DataFileFormat, Literal, PartitionKey, Struct,
    Transform,
};
use iceberg::table::Table as IceTable;
use iceberg::transaction::{ApplyTransactionAction, Transaction};
use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
use iceberg::writer::file_writer::location_generator::{
    DefaultFileNameGenerator, DefaultLocationGenerator,
};
use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
use iceberg::writer::file_writer::ParquetWriterBuilder;
use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
use iceberg::Catalog;
use parquet::arrow::ArrowWriter;
use parquet::file::properties::WriterProperties;
use rayon::prelude::*;

use crate::bridge::recast;
use crate::error::{Result, SkadeError};
use crate::read::arrow_schema_of;

/// Process-unique data-file name prefix so repeated commits (and concurrent
/// writers in one process) never collide on file names.
fn unique_prefix() -> String {
    static SEQ: AtomicU64 = AtomicU64::new(0);
    let nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_nanos() as u64)
        .unwrap_or(0);
    format!(
        "skade-{}-{nanos:x}-{}",
        std::process::id(),
        SEQ.fetch_add(1, Ordering::Relaxed)
    )
}

/// Throughput of an [`ingest`] run.
#[derive(Debug, Clone)]
pub struct IngestStats {
    pub rows: u64,
    pub commits: u64,
    pub elapsed: Duration,
}

impl IngestStats {
    pub fn rows_per_sec(&self) -> f64 {
        let s = self.elapsed.as_secs_f64();
        if s > 0.0 { self.rows as f64 / s } else { 0.0 }
    }
    pub fn commits_per_sec(&self) -> f64 {
        let s = self.elapsed.as_secs_f64();
        if s > 0.0 { self.commits as f64 / s } else { 0.0 }
    }
}

/// Derive the single [`PartitionKey`] for `batches` from `table`'s default
/// partition spec, or `None` for an unpartitioned table (the common case).
///
/// Supports **identity** transforms over **string** columns (`Utf8`/`LargeUtf8`).
/// Every row across every batch must share the same partition value — an
/// [`append`] commits one group of data files as one snapshot, so it writes a
/// single partition. A batch that spans partitions, or a non-identity/non-string
/// partition column, errors loudly rather than mis-tagging the files (which
/// would silently break partition pruning on read).
fn partition_key_for(table: &IceTable, batches: &[RecordBatch]) -> Result<Option<PartitionKey>> {
    let meta = table.metadata();
    let spec = meta.default_partition_spec();
    if spec.fields().is_empty() {
        return Ok(None);
    }
    let schema = meta.current_schema();
    let mut lits: Vec<Option<Literal>> = Vec::with_capacity(spec.fields().len());
    for f in spec.fields() {
        if f.transform != Transform::Identity {
            return Err(SkadeError::other(format!(
                "skade::append: only identity partition transforms are supported (got {:?} on `{}`)",
                f.transform, f.name
            )));
        }
        let src = schema.field_by_id(f.source_id).ok_or_else(|| {
            SkadeError::other(format!("partition source field id {} not in schema", f.source_id))
        })?;
        // Single partition value across every batch (None ⇒ all batches empty).
        let mut value: Option<String> = None;
        for b in batches {
            let col = b.column_by_name(&src.name).ok_or_else(|| {
                SkadeError::other(format!("partition column `{}` missing from batch", src.name))
            })?;
            let read = |i: usize| -> Result<&str> {
                if let Some(a) = col.as_any().downcast_ref::<StringArray>() {
                    Ok(a.value(i))
                } else if let Some(a) = col.as_any().downcast_ref::<LargeStringArray>() {
                    Ok(a.value(i))
                } else {
                    Err(SkadeError::other(format!(
                        "partition column `{}` is not a string (only Utf8/LargeUtf8 identity partitions)",
                        src.name
                    )))
                }
            };
            for i in 0..col.len() {
                let v = read(i)?;
                match &value {
                    None => value = Some(v.to_string()),
                    Some(v0) if v0 != v => {
                        return Err(SkadeError::other(format!(
                            "skade::append: batch spans multiple `{}` partitions (`{v0}` vs `{v}`); \
                             single-partition appends only",
                            src.name
                        )));
                    }
                    _ => {}
                }
            }
        }
        match value {
            Some(v) => lits.push(Some(Literal::string(v))),
            None => return Ok(None),
        }
    }
    Ok(Some(PartitionKey::new(
        spec.as_ref().clone(),
        schema.clone(),
        Struct::from_iter(lits),
    )))
}

/// Write `batches` into one group of Parquet data files and commit them as a
/// single `fast_append` snapshot. Batches are [`recast`] to the table's
/// field-id Arrow schema first (covers `Utf8 → LargeUtf8` etc.), so any
/// column-compatible batches work. For a partitioned table the (single)
/// partition key is derived from the batch via [`partition_key_for`]. Returns
/// the updated table.
pub async fn append(
    catalog: &dyn Catalog,
    table: &IceTable,
    batches: &[RecordBatch],
) -> Result<IceTable> {
    let target = arrow_schema_of(table)?;
    let batches = recast(batches, target)?;
    let partition = partition_key_for(table, &batches)?;

    let schema = table.metadata().current_schema().clone();
    let data_location = format!("{}/data", table.metadata().location());
    let location_gen = DefaultLocationGenerator::with_data_location(data_location);
    let file_name_gen =
        DefaultFileNameGenerator::new(unique_prefix(), None, DataFileFormat::Parquet);

    let pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema);
    let rolling = RollingFileWriterBuilder::new_with_default_file_size(
        pw,
        table.file_io().clone(),
        location_gen,
        file_name_gen,
    );
    let mut writer = DataFileWriterBuilder::new(rolling).build(partition).await?;
    for b in &batches {
        writer.write(b.clone()).await?;
    }
    let data_files = writer.close().await?;

    let tx = Transaction::new(table);
    let action = tx.fast_append().add_data_files(data_files);
    let tx = action.apply(tx)?;
    Ok(tx.commit(catalog).await?)
}

/// Ingest all batches, committing every `batches_per_commit` batches via
/// [`append`]. Returns the final table + throughput.
pub async fn ingest(
    catalog: &dyn Catalog,
    table: IceTable,
    batches: impl IntoIterator<Item = RecordBatch>,
    batches_per_commit: usize,
) -> Result<(IceTable, IngestStats)> {
    let bpc = batches_per_commit.max(1);
    let mut table = table;
    let (mut rows, mut commits) = (0u64, 0u64);
    let mut group: Vec<RecordBatch> = Vec::with_capacity(bpc);
    let start = Instant::now();
    for batch in batches {
        rows += batch.num_rows() as u64;
        group.push(batch);
        if group.len() >= bpc {
            table = append(catalog, &table, &group).await?;
            commits += 1;
            group.clear();
        }
    }
    if !group.is_empty() {
        table = append(catalog, &table, &group).await?;
        commits += 1;
    }
    Ok((table, IngestStats { rows, commits, elapsed: start.elapsed() }))
}

/// Parallel bulk ingest — the gatling topology: **all cores encode** (rayon),
/// **one** sequential writer commits.
///
/// Each inner `Vec<RecordBatch>` is encoded into one in-memory Parquet file
/// across the rayon pool (pure CPU, no I/O), then a single writer streams the
/// files to the table's `FileIO` and commits every `files_per_commit` files via
/// `fast_append`. Single-partition per call (same rule as [`append`]); for a
/// partitioned table every file is tagged with the derived partition key.
///
/// The encode phase holds all files in memory before the write phase, so size
/// `groups` to your memory budget for very large loads (chunk across calls).
pub async fn ingest_parallel(
    catalog: &dyn Catalog,
    table: IceTable,
    groups: Vec<Vec<RecordBatch>>,
    files_per_commit: usize,
) -> Result<(IceTable, IngestStats)> {
    let fpc = files_per_commit.max(1);
    let target = arrow_schema_of(&table)?;
    let start = Instant::now();

    // Recast each group to the table's field-id Arrow schema (so the Parquet
    // columns carry the field ids the scan maps back through).
    let groups: Vec<Vec<RecordBatch>> =
        groups.into_iter().map(|g| recast(&g, target.clone())).collect::<Result<_>>()?;

    // One partition key for the whole call (single-partition rule), derived from
    // every batch — `None` for an unpartitioned table.
    let all: Vec<RecordBatch> = groups.iter().flatten().cloned().collect();
    let partition = partition_key_for(&table, &all)?;
    drop(all);

    // ── Encode stage: all cores, pure CPU, no I/O. ──
    let mut encoded: Vec<(usize, Vec<u8>, u64)> = groups
        .into_par_iter()
        .enumerate()
        .map(|(seq, g)| {
            let (bytes, rows) = encode_parquet(target.clone(), &g)?;
            Ok::<_, SkadeError>((seq, bytes, rows))
        })
        .collect::<Result<Vec<_>>>()?;
    encoded.sort_by_key(|(seq, _, _)| *seq);

    // ── Writer stage: exactly ONE sequential writer + committer. ──
    let prefix = unique_prefix();
    let data_dir = format!("{}/data", table.metadata().location());
    let mut table = table;
    let (mut rows, mut commits) = (0u64, 0u64);
    let mut pending: Vec<DataFile> = Vec::with_capacity(fpc);
    for (seq, bytes, record_count) in encoded {
        let path = format!("{data_dir}/{prefix}-{seq:08}.parquet");
        let len = bytes.len() as u64;
        table.file_io().new_output(&path)?.write(Bytes::from(bytes)).await?;
        rows += record_count;
        let mut b = DataFileBuilder::default();
        b.content(DataContentType::Data)
            .file_path(path)
            .file_format(DataFileFormat::Parquet)
            .record_count(record_count)
            .file_size_in_bytes(len);
        if let Some(pk) = partition.as_ref() {
            b.partition(pk.data().clone());
            b.partition_spec_id(pk.spec().spec_id());
        }
        pending.push(b.build().map_err(SkadeError::other)?);
        if pending.len() >= fpc {
            table = commit_files(catalog, &table, std::mem::take(&mut pending)).await?;
            commits += 1;
        }
    }
    if !pending.is_empty() {
        table = commit_files(catalog, &table, pending).await?;
        commits += 1;
    }
    Ok((table, IngestStats { rows, commits, elapsed: start.elapsed() }))
}

/// Encode a group of batches into one in-memory Parquet file — CPU only, the
/// part that fans out across cores. `schema` must carry the iceberg field ids.
fn encode_parquet(schema: ArrowSchemaRef, batches: &[RecordBatch]) -> Result<(Vec<u8>, u64)> {
    let mut buf: Vec<u8> = Vec::new();
    let mut rows = 0u64;
    {
        let mut w = ArrowWriter::try_new(&mut buf, schema, Some(WriterProperties::builder().build()))?;
        for b in batches {
            rows += b.num_rows() as u64;
            w.write(b)?;
        }
        w.close()?;
    }
    Ok((buf, rows))
}

/// One `fast_append` commit of already-written data files.
async fn commit_files(
    catalog: &dyn Catalog,
    table: &IceTable,
    files: Vec<DataFile>,
) -> Result<IceTable> {
    let tx = Transaction::new(table);
    let action = tx.fast_append().add_data_files(files);
    let tx = action.apply(tx)?;
    Ok(tx.commit(catalog).await?)
}
