use std::{
    io::{BufWriter, IntoInnerError, Seek},
    path::{Path, PathBuf},
};

use arrow_array::RecordBatch;
use arrow_ipc::writer::FileWriter;
use arrow_schema::{Schema, SchemaRef};
use arrow_select::concat::concat_batches;
use tempfile::NamedTempFile;
use tracing::{error, warn};

use crate::dataset::storage::{error::DatasetFsError, layout::chunk_path};

const MAX_BATCH_BYTE_SIZE: usize = 64 * 1024 * 1024;
const MAX_CHUNK_BYTE_SIZE: u64 = 256 * 1024 * 1024;

pub(crate) struct ChunkWriter {
    dir_path: PathBuf,
    schema: SchemaRef,
    next_chunk_index: usize,
    current_writer: Option<InnerWriter>,
}

impl ChunkWriter {
    pub(crate) fn new(schema: SchemaRef, dir_path: PathBuf) -> Self {
        Self {
            dir_path,
            schema,
            next_chunk_index: 0,
            current_writer: None,
        }
    }

    /// Write a [`RecordBatch`], return true if current chunk file is completed.
    pub(crate) fn write(&mut self, batch: RecordBatch) -> Result<bool, DatasetFsError> {
        let writer = self.current_writer()?;
        writer.write(batch)?;
        if writer.written_size >= MAX_CHUNK_BYTE_SIZE {
            self.finish_current_writer()?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    pub(crate) fn finish(mut self) -> Result<(), DatasetFsError> {
        self.finish_current_writer()
    }

    fn current_writer(&mut self) -> Result<&mut InnerWriter, DatasetFsError> {
        if self.current_writer.is_none() {
            self.current_writer = Some(self.create_writer()?);
        }
        Ok(self.current_writer.as_mut().expect("Not none here."))
    }

    fn create_writer(&mut self) -> Result<InnerWriter, DatasetFsError> {
        let writer = InnerWriter::new(&self.dir_path, self.next_chunk_index, &self.schema)?;
        self.next_chunk_index += 1;
        Ok(writer)
    }

    fn finish_current_writer(&mut self) -> Result<(), DatasetFsError> {
        if let Some(writer) = self.current_writer.take() {
            writer.finish()?;
        }
        Ok(())
    }
}

impl Drop for ChunkWriter {
    fn drop(&mut self) {
        if self.current_writer.is_some() {
            warn!("ChunkWriter dropped with an active writer.");
            if let Err(e) = self.finish_current_writer() {
                error!("Error dropping ChunkWriter: {e}");
            }
        }
    }
}

struct InnerWriter {
    inner: FileWriter<BufWriter<NamedTempFile>>,
    buffered_batches: Vec<RecordBatch>,
    buffered_size: usize,
    written_size: u64,
    final_path: PathBuf,
}

impl InnerWriter {
    fn new(
        dir_path: &Path,
        chunk_index: usize,
        schema: &Schema,
    ) -> Result<InnerWriter, DatasetFsError> {
        let file = NamedTempFile::new_in(dir_path)?;
        let writer = FileWriter::try_new(BufWriter::new(file), schema)?;
        Ok(InnerWriter {
            inner: writer,
            buffered_batches: vec![],
            buffered_size: 0,
            written_size: 0,
            final_path: chunk_path(dir_path, chunk_index),
        })
    }

    fn write(&mut self, batch: RecordBatch) -> Result<(), DatasetFsError> {
        self.push_to_buffer(batch)?;
        if self.buffered_size >= MAX_BATCH_BYTE_SIZE {
            self.flush()?;
        }
        Ok(())
    }

    fn finish(mut self) -> Result<(), DatasetFsError> {
        self.flush()?;
        let file = self
            .inner
            .into_inner()?
            .into_inner()
            .map_err(IntoInnerError::into_error)?;
        file.persist(self.final_path).map_err(|e| e.error)?;
        Ok(())
    }

    fn flush(&mut self) -> Result<(), DatasetFsError> {
        if !self.buffered_batches.is_empty() {
            let batch = self.drain_buffer()?;
            self.inner.write(&batch)?;
            self.written_size = self.inner.get_mut().stream_position()?;
        }
        Ok(())
    }

    fn push_to_buffer(&mut self, batch: RecordBatch) -> Result<(), DatasetFsError> {
        if batch.schema() == *self.inner.schema() {
            self.buffered_size += batch.get_array_memory_size();
            self.buffered_batches.push(batch);
            Ok(())
        } else {
            Err(DatasetFsError::SchemaMismatch)
        }
    }

    fn drain_buffer(&mut self) -> Result<RecordBatch, DatasetFsError> {
        let group = concat_batches(self.inner.schema(), &self.buffered_batches)?;
        self.buffered_batches.clear();
        self.buffered_size = 0;
        Ok(group)
    }
}
