/// Parquet writer: adaptive flush threshold based on available RAM.
///
/// flush_threshold is computed by the pipeline from coren::MachCap:
///   budget_per_worker = mem_available * 0.4 / n_workers
///   threshold = (budget / EST_BYTES_PER_ROW).clamp(5_000, 500_000)
///
/// Pi 4GB / 4 workers:  ~100K rows between flushes
/// Laptop 32GB / 8:     ~400K rows
/// HPC 512GB:           clamped at 500K
///
/// File layout: _workers/worker_N/table_name.parquet
use crate::schema::{self, DT, ALL_TABLES, TableDef};
use arrow::array::{ArrayRef, Int32Builder, StringBuilder};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use parquet::arrow::ArrowWriter;
use parquet::basic::{Compression, ZstdLevel};
use parquet::file::properties::WriterProperties;
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::Arc;

enum Builder {
    Int(Int32Builder),
    Str(StringBuilder),
}

fn make_props(row_group_size: usize) -> anyhow::Result<WriterProperties> {
    Ok(WriterProperties::builder()
        .set_compression(Compression::ZSTD(ZstdLevel::try_new(3)?))
        .set_dictionary_enabled(true)
        .set_max_row_group_size(row_group_size)
        .build())
}

struct TableWriter {
    builders: Vec<Builder>,
    n_rows: usize,
    flush_at: usize,
    schema: SchemaRef,
    writer: ArrowWriter<File>,
}

impl TableWriter {
    fn new(def: &'static TableDef, path: &Path, flush_at: usize) -> anyhow::Result<Self> {
        let schema = schema::arrow_schema(def);
        let file = File::create(path)?;
        // Row group size = flush_at. Each flush produces one complete row group.
        // ArrowWriter writes encoded pages to disk immediately on write().
        let writer = ArrowWriter::try_new(file, schema.clone(), Some(make_props(flush_at)?))?;
        let builders = def.columns.iter().map(|c| match c.dtype {
            DT::I32 => Builder::Int(Int32Builder::new()),
            DT::Str => Builder::Str(StringBuilder::new()),
        }).collect();
        Ok(TableWriter { builders, n_rows: 0, flush_at, schema, writer })
    }

    fn flush(&mut self) -> anyhow::Result<()> {
        if self.n_rows == 0 { return Ok(()); }
        let arrays: Vec<ArrayRef> = self.builders.iter_mut().map(|b| match b {
            Builder::Int(b) => Arc::new(b.finish()) as ArrayRef,
            Builder::Str(b) => Arc::new(b.finish()) as ArrayRef,
        }).collect();
        self.writer.write(&RecordBatch::try_new(self.schema.clone(), arrays)?)?;
        self.n_rows = 0;
        Ok(())
    }

    fn finish(mut self) -> anyhow::Result<()> {
        self.flush()?;
        self.writer.close()?;
        Ok(())
    }
}

/// Handle for appending one row. Columns must be pushed in schema order.
pub struct RowWriter<'a> {
    tw: &'a mut TableWriter,
    col: usize,
}

impl RowWriter<'_> {
    #[inline(always)]
    pub fn int(&mut self, v: i32) {
        if let Builder::Int(b) = &mut self.tw.builders[self.col] {
            b.append_value(v);
        }
        self.col += 1;
    }

    #[inline(always)]
    pub fn int_opt(&mut self, v: Option<i32>) {
        if let Builder::Int(b) = &mut self.tw.builders[self.col] {
            match v { Some(v) => b.append_value(v), None => b.append_null() }
        }
        self.col += 1;
    }

    #[inline(always)]
    pub fn str(&mut self, v: Option<&str>) {
        if let Builder::Str(b) = &mut self.tw.builders[self.col] {
            match v { Some(s) => b.append_value(s), None => b.append_null() }
        }
        self.col += 1;
    }
}

pub struct WorkerWriter {
    tables: HashMap<&'static str, TableWriter>,
    dir: PathBuf,
    flush_at: usize,
}

impl WorkerWriter {
    /// Create a new writer. Tables are initialized lazily on first row write,
    /// avoiding empty Parquet files for sparse tables (e.g. delete_citation
    /// in baseline files). `flush_threshold` controls how many rows
    /// accumulate before flushing to Parquet. Caller computes this
    /// from available RAM via coren.
    pub fn new(output_dir: &Path, worker_id: usize, flush_threshold: usize) -> anyhow::Result<Self> {
        let dir = output_dir.join(format!("worker_{worker_id}"));
        std::fs::create_dir_all(&dir)?;
        Ok(WorkerWriter {
            tables: HashMap::with_capacity(ALL_TABLES.len()),
            dir,
            flush_at: flush_threshold,
        })
    }

    #[inline]
    pub fn row(&mut self, table: &'static str, f: impl FnOnce(&mut RowWriter)) -> anyhow::Result<()> {
        if !self.tables.contains_key(table) {
            let def = ALL_TABLES.iter().find(|d| d.name == table).expect("unknown table");
            let path = self.dir.join(format!("{}.parquet", table));
            self.tables.insert(table, TableWriter::new(def, &path, self.flush_at)?);
        }
        let tw = self.tables.get_mut(table).unwrap();
        let mut rw = RowWriter { tw, col: 0 };
        f(&mut rw);
        rw.tw.n_rows += 1;
        if rw.tw.n_rows >= rw.tw.flush_at {
            rw.tw.flush()?;
        }
        Ok(())
    }

    pub fn finish(self) -> anyhow::Result<()> {
        for (_, tw) in self.tables { tw.finish()?; }
        Ok(())
    }
}
