use crate::Result;
use crate::readers::{
    OnceReader, RecordingBufRead, ReplayBufRead, batch_size::resolve_batch_size,
    dataframe_from_stream, emits_single_batch, finish_stream, stream_record_batches,
};
use autark_dataframe::DataFrame;
use autark_enhanced_reader::autoread;
use std::sync::Arc;

use arrow::datatypes::Schema;
use arrow_csv::{ReaderBuilder, reader::Format};

const DEFAULT_BATCH_SIZE: usize = 1 << 16;

pub struct CsvReader {
    uris: Vec<String>,
    batch_size: usize,
    batch_size_explicit: bool,
    schema: std::sync::OnceLock<Arc<Schema>>,
}

impl CsvReader {
    pub fn new(uri: &str) -> Result<CsvReader> {
        Self::from_uris(vec![uri.to_string()])
    }

    pub fn from_uris(uris: Vec<String>) -> Result<CsvReader> {
        if uris.is_empty() {
            return Err(crate::Error::EmptyReader);
        }

        Ok(CsvReader {
            uris,
            batch_size: DEFAULT_BATCH_SIZE,
            batch_size_explicit: false,
            schema: std::sync::OnceLock::new(),
        })
    }

    pub fn with_batch_size(mut self, batch_size: usize) -> CsvReader {
        self.batch_size = batch_size;
        self.batch_size_explicit = true;
        self
    }

    fn ensure_batch_size(&mut self) -> Result<()> {
        if self.batch_size_explicit {
            return Ok(());
        }

        self.batch_size = resolve_batch_size(DEFAULT_BATCH_SIZE);
        Ok(())
    }

    fn infer_schema(&self) -> Result<Arc<Schema>> {
        let mut reader = autoread(&self.uris[0])?;
        let (schema, _inferred_rows) = Format::default()
            .with_header(true)
            .infer_schema(&mut reader, Some(1e6 as usize))?;
        Ok(Arc::new(schema))
    }

    fn schema_cached(&self) -> Result<Arc<Schema>> {
        if let Some(schema) = self.schema.get() {
            return Ok(schema.clone());
        }
        let schema = self.infer_schema()?;
        let _ = self.schema.set(schema.clone());
        Ok(schema)
    }

    fn cache_schema(&self, schema: Arc<Schema>) -> Arc<Schema> {
        let _ = self.schema.set(schema.clone());
        self.schema.get().cloned().unwrap_or(schema)
    }

    fn open_streaming_reader(
        &self,
        uri: &str,
    ) -> Result<arrow_csv::Reader<ReplayBufRead<Box<dyn std::io::BufRead + Send>>>> {
        let source = autoread(uri)?;
        let mut recording = RecordingBufRead::new(source);
        let (schema, _inferred_rows) = Format::default()
            .with_header(true)
            .infer_schema(&mut recording, Some(1e6 as usize))?;
        let schema = self.cache_schema(Arc::new(schema));
        let (source, prefix) = recording.into_parts();

        ReaderBuilder::new(schema)
            .with_header(true)
            .with_batch_size(self.batch_size.max(1))
            .build(ReplayBufRead::new(prefix, source))
            .map_err(Into::into)
    }
}

impl OnceReader for CsvReader {
    fn read(&mut self) -> Result<DataFrame> {
        dataframe_from_stream(self)
    }

    fn schema(&self) -> Result<Arc<Schema>> {
        self.schema_cached()
    }

    fn read_stream(
        &mut self,
        sink: &mut dyn FnMut(arrow::array::RecordBatch) -> Result<()>,
    ) -> Result<()> {
        self.ensure_batch_size()?;
        let emit_single_batch = emits_single_batch(self.batch_size);
        let mut saw_any = false;
        let mut batch_ix = 0usize;
        if self.schema.get().is_some() {
            let schema = self.schema_cached()?;
            self.uris.iter().try_for_each(|uri| -> Result<()> {
                let mut reader = ReaderBuilder::new(schema.clone())
                    .with_header(true)
                    .with_batch_size(self.batch_size.max(1))
                    .build(autoread(uri)?)?;
                stream_record_batches(
                    "csv",
                    &mut batch_ix,
                    &mut saw_any,
                    emit_single_batch,
                    &mut reader,
                    crate::Error::from,
                    sink,
                )
            })?;
        } else {
            let mut uris = self.uris.iter();
            let first_uri = uris.next().ok_or(crate::Error::EmptyReader)?;
            let mut first_reader = self.open_streaming_reader(first_uri)?;
            stream_record_batches(
                "csv",
                &mut batch_ix,
                &mut saw_any,
                emit_single_batch,
                &mut first_reader,
                crate::Error::from,
                sink,
            )?;

            let schema = self.schema_cached()?;
            uris.try_for_each(|uri| -> Result<()> {
                let mut reader = ReaderBuilder::new(schema.clone())
                    .with_header(true)
                    .with_batch_size(self.batch_size.max(1))
                    .build(autoread(uri)?)?;
                stream_record_batches(
                    "csv",
                    &mut batch_ix,
                    &mut saw_any,
                    emit_single_batch,
                    &mut reader,
                    crate::Error::from,
                    sink,
                )
            })?;
        }

        finish_stream(saw_any)
    }
}
