pub mod arrow;
pub mod batch_size;
pub mod csv;
pub mod json;
pub mod parquet;
pub mod transcend;

use crate::Result;
use ::arrow::{array::RecordBatch, compute::concat_batches, datatypes::Schema};
// use arrow::datatypes::Schema;
use autark_dataframe::DataFrame;
use std::io::{BufRead, Cursor, Read};
use std::sync::Arc;

pub(crate) fn log_stream_batch(reader: &str, batch_ix: usize, batch: &RecordBatch) {
    if autark_telemetry::enabled(autark_telemetry::log::Level::Info) {
        autark_telemetry::log::info!(
            target: "autark_reader::stream",
            "reader={reader} batch={} rows={} cols={}",
            batch_ix,
            batch.num_rows(),
            batch.num_columns(),
        );
    }
}

pub trait OnceReader {
    fn read(&mut self) -> Result<DataFrame>;
    fn schema(&self) -> Result<Arc<Schema>>;
    fn read_stream(&mut self, sink: &mut dyn FnMut(RecordBatch) -> Result<()>) -> Result<()>;
}

pub(crate) fn dataframe_from_stream(reader: &mut dyn OnceReader) -> Result<DataFrame> {
    let mut batches = Vec::new();
    reader.read_stream(&mut |batch| {
        batches.push(batch);
        Ok(())
    })?;

    let batch = merge_record_batches(&batches)?;
    Ok(DataFrame::try_from(batch)?)
}

pub(crate) fn merge_record_batches(batches: &[RecordBatch]) -> Result<RecordBatch> {
    let first = batches.first().ok_or(crate::Error::EmptyReader)?;
    if batches.len() == 1 {
        Ok(first.clone())
    } else {
        Ok(concat_batches(&first.schema(), batches)?)
    }
}

pub(crate) fn emits_single_batch(batch_size: usize) -> bool {
    batch_size == 0
}

pub(crate) fn stream_record_batches<I, E, F>(
    reader: &str,
    batch_ix: &mut usize,
    saw_any: &mut bool,
    emit_single_batch: bool,
    batches: &mut I,
    map_error: F,
    sink: &mut dyn FnMut(RecordBatch) -> Result<()>,
) -> Result<()>
where
    I: Iterator<Item = std::result::Result<RecordBatch, E>>,
    F: Fn(E) -> crate::Error,
{
    if emit_single_batch {
        let mut collected = Vec::new();
        batches.try_for_each(|batch| -> Result<()> {
            let batch = batch.map_err(&map_error)?;
            *saw_any = true;
            collected.push(batch);
            Ok(())
        })?;

        if !collected.is_empty() {
            let batch = merge_record_batches(&collected)?;
            log_stream_batch(reader, *batch_ix, &batch);
            sink(batch)?;
            *batch_ix += 1;
        }

        return Ok(());
    }

    batches.try_for_each(|batch| {
        let batch = batch.map_err(&map_error)?;
        *saw_any = true;
        log_stream_batch(reader, *batch_ix, &batch);
        sink(batch)?;
        *batch_ix += 1;
        Ok(())
    })
}

pub(crate) fn finish_stream(saw_any: bool) -> Result<()> {
    if saw_any {
        Ok(())
    } else {
        Err(crate::Error::EmptyReader)
    }
}

pub(crate) struct RecordingBufRead<R> {
    inner: R,
    prefix: Vec<u8>,
}

impl<R> RecordingBufRead<R> {
    pub(crate) fn new(inner: R) -> Self {
        Self {
            inner,
            prefix: Vec::new(),
        }
    }

    pub(crate) fn into_parts(self) -> (R, Vec<u8>) {
        (self.inner, self.prefix)
    }
}

impl<R: Read> Read for RecordingBufRead<R> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let n = self.inner.read(buf)?;
        self.prefix.extend_from_slice(&buf[..n]);
        Ok(n)
    }
}

impl<R: BufRead> BufRead for RecordingBufRead<R> {
    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
        self.inner.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        if amt == 0 {
            return;
        }

        if let Ok(buf) = self.inner.fill_buf() {
            self.prefix.extend_from_slice(&buf[..amt.min(buf.len())]);
        }
        self.inner.consume(amt);
    }
}

pub(crate) struct ReplayBufRead<R> {
    prefix: Cursor<Vec<u8>>,
    inner: R,
}

impl<R> ReplayBufRead<R> {
    pub(crate) fn new(prefix: Vec<u8>, inner: R) -> Self {
        Self {
            prefix: Cursor::new(prefix),
            inner,
        }
    }
}

impl<R: Read> Read for ReplayBufRead<R> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let n = self.prefix.read(buf)?;
        if n != 0 {
            return Ok(n);
        }
        self.inner.read(buf)
    }
}

impl<R: BufRead> BufRead for ReplayBufRead<R> {
    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
        let buf = self.prefix.fill_buf()?;
        if !buf.is_empty() {
            return Ok(buf);
        }
        self.inner.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        let prefix_remaining = self.prefix.fill_buf().map(|buf| buf.len()).unwrap_or(0);
        if prefix_remaining != 0 {
            let consumed = amt.min(prefix_remaining);
            self.prefix.consume(consumed);
            if consumed == amt {
                return;
            }
            self.inner.consume(amt - consumed);
            return;
        }
        self.inner.consume(amt);
    }
}
