use crate::Result;
use crate::readers::{
    OnceReader, 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_to_bytes, local_path};
use std::sync::Arc;

use arrow::datatypes::Schema;
use bytes::Bytes;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

const DEFAULT_BATCH_SIZE: usize = 1 << 16;

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

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

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

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

    pub fn with_batch_size(mut self, batch_size: usize) -> ParquetReader {
        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 bytes_cached(&self) -> Result<Bytes> {
        if let Some(bytes) = self.bytes.get() {
            return Ok(bytes.clone());
        }
        let bytes = Bytes::from(autoread_to_bytes(&self.uris[0])?);
        let _ = self.bytes.set(bytes.clone());
        Ok(bytes)
    }

    fn infer_schema(&self) -> Result<Arc<Schema>> {
        if let Some(path) = local_path(&self.uris[0])? {
            let file = std::fs::File::open(path)?;
            let builder = ParquetRecordBatchReaderBuilder::try_new(file)
                .map_err(|err| crate::Error::Reader(format!("parquet error: {err}")))?;
            return Ok(builder.schema().clone());
        }

        let bytes = self.bytes_cached()?;
        let builder = ParquetRecordBatchReaderBuilder::try_new(bytes)
            .map_err(|err| crate::Error::Reader(format!("parquet error: {err}")))?;
        Ok(builder.schema().clone())
    }

    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 open_reader_for_uri(
        &self,
        uri: &str,
    ) -> Result<parquet::arrow::arrow_reader::ParquetRecordBatchReader> {
        if let Some(path) = local_path(uri)? {
            let file = std::fs::File::open(path)?;
            return ParquetRecordBatchReaderBuilder::try_new(file)
                .map_err(|err| crate::Error::Reader(format!("parquet error: {err}")))?
                .with_batch_size(self.batch_size.max(1))
                .build()
                .map_err(|err| crate::Error::Reader(format!("parquet error: {err}")));
        }

        let bytes = if uri == self.uris[0] {
            self.bytes_cached()?
        } else {
            Bytes::from(autoread_to_bytes(uri)?)
        };

        ParquetRecordBatchReaderBuilder::try_new(bytes)
            .map_err(|err| crate::Error::Reader(format!("parquet error: {err}")))?
            .with_batch_size(self.batch_size.max(1))
            .build()
            .map_err(|err| crate::Error::Reader(format!("parquet error: {err}")))
    }
}

impl OnceReader for ParquetReader {
    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;
        self.uris.iter().try_for_each(|uri| -> Result<()> {
            let mut reader = self.open_reader_for_uri(uri)?;
            stream_record_batches(
                "parquet",
                &mut batch_ix,
                &mut saw_any,
                emit_single_batch,
                &mut reader,
                |err| crate::Error::Reader(format!("parquet error: {err}")),
                sink,
            )
        })?;

        finish_stream(saw_any)
    }
}
