use std::collections::HashMap;
use std::time::Instant;

use arrow::array::RecordBatch;
use autark_dataframe::DataFrame;
use autark_reader::OnceReader;

use crate::Result;

pub(crate) fn stream_input_batches(
    reader: &mut dyn OnceReader,
    reader_ix: usize,
    rename_map: Option<&HashMap<String, String>>,
    mut on_batch: impl FnMut(RecordBatch) -> Result<()>,
) -> Result<()> {
    let started = Instant::now();
    let mut batches = 0usize;
    let mut rows = 0usize;
    reader.read_stream(&mut |batch: RecordBatch| {
        batches += 1;
        let batch_rows = batch.num_rows();
        rows += batch_rows;
        let batch = maybe_rename_batch(batch, rename_map)?;
        on_batch(batch)?;
        log_batch_runtime(reader_ix, batches - 1, batch_rows, 0);
        Ok(())
    })?;
    log_reader_runtime(reader_ix, batches, rows, started.elapsed().as_millis());
    if batches == 0 {
        return Err(autark_error::Error::Reader(format!(
            "reader `{reader_ix}` produced no record batches"
        )));
    }
    Ok(())
}

fn maybe_rename_batch(
    batch: RecordBatch,
    rename_map: Option<&HashMap<String, String>>,
) -> Result<RecordBatch> {
    match rename_map {
        Some(rename_map) if !rename_map.is_empty() => Ok(DataFrame::try_from(batch)?
            .rename_columns(rename_map)?
            .record_batch()),
        _ => Ok(batch),
    }
}

fn log_batch_runtime(reader_ix: usize, batch_ix: usize, rows: usize, runtime_ms: u128) {
    if autark_telemetry::enabled(autark_telemetry::log::Level::Info) {
        autark_telemetry::log::info!(
            target: "autark_interface_onceframe::stream",
            "reader={} batch={} rows={} runtime_ms={}",
            reader_ix,
            batch_ix,
            rows,
            runtime_ms,
        );
    }
}

fn log_reader_runtime(reader_ix: usize, batches: usize, rows: usize, elapsed_ms: u128) {
    if autark_telemetry::enabled(autark_telemetry::log::Level::Info) {
        autark_telemetry::log::info!(
            target: "autark_interface_onceframe::stream",
            "reader={} batches={} rows={} elapsed_ms={}",
            reader_ix,
            batches,
            rows,
            elapsed_ms,
        );
    }
}
