use std::{mem, path::Path};

use anyhow::{Context, Result, bail};
use tempfile::NamedTempFile;

use crate::utils::row_entry::{RowEntry, RowEntryReader, RowEntryWriter};

const OBSOLETE_ROW_ENTRY_BATCH_SIZE: usize = 4_096;

pub(super) struct ObsoleteRowEntryTempFile {
    file: NamedTempFile,
    writer: Option<RowEntryWriter>,
    buffered_entries: Vec<RowEntry>,
}

impl ObsoleteRowEntryTempFile {
    /// Create a temp file in the system default temp directory.
    #[cfg(test)]
    pub(super) fn new() -> Result<Self> {
        let file = NamedTempFile::new().context("create obsolete row temp file failed")?;
        // Obsolete entries are appended in callback order, not key order.
        // During dedup we may prefetch a later out-of-range key from the same
        // source before deciding that the current in-range key is obsolete.
        let writer = RowEntryWriter::create_unchecked(file.path())
            .context("create obsolete row temp file writer failed")?;
        Ok(Self {
            file,
            writer: Some(writer),
            buffered_entries: Vec::with_capacity(OBSOLETE_ROW_ENTRY_BATCH_SIZE),
        })
    }

    /// Create a temp file in the given directory.
    pub(super) fn new_in(dir: impl AsRef<Path>) -> Result<Self> {
        let dir = dir.as_ref();
        std::fs::create_dir_all(dir)
            .with_context(|| format!("create obsolete row temp dir failed, dir={:?}", dir))?;
        let file =
            NamedTempFile::new_in(dir).context("create obsolete row temp file in dir failed")?;
        // Keep obsolete temp files unchecked for the same reason as `new()`.
        let writer = RowEntryWriter::create_unchecked(file.path())
            .context("create obsolete row temp file writer failed")?;
        Ok(Self {
            file,
            writer: Some(writer),
            buffered_entries: Vec::with_capacity(OBSOLETE_ROW_ENTRY_BATCH_SIZE),
        })
    }

    pub(super) fn push(&mut self, entry: &RowEntry) -> Result<()> {
        if self.writer.is_none() {
            bail!(
                "obsolete row temp file is already finalized, file={:?}",
                self.file.path()
            );
        }

        self.buffered_entries.push(entry.clone());
        if self.buffered_entries.len() >= OBSOLETE_ROW_ENTRY_BATCH_SIZE {
            self.flush_buffered_entries()?;
        }
        Ok(())
    }

    /// Finalize the parquet file. No more entries can be pushed afterwards.
    pub(super) fn finalize(&mut self) -> Result<()> {
        self.finalize_writer()
    }

    /// Open a reader after implicitly finalizing the parquet writer.
    pub(super) fn reader(&mut self) -> Result<ObsoleteRowEntryReader> {
        self.finalize_writer()?;
        let path = self.file.path().to_path_buf();
        Ok(ObsoleteRowEntryReader {
            reader: RowEntryReader::open(&path)
                .with_context(|| format!("open obsolete row temp file failed, path={:?}", path))?,
        })
    }

    fn flush_buffered_entries(&mut self) -> Result<()> {
        if self.buffered_entries.is_empty() {
            return Ok(());
        }

        let entries = mem::take(&mut self.buffered_entries);
        let writer = self.writer.as_mut().with_context(|| {
            format!(
                "obsolete row temp file writer is missing, file={:?}",
                self.file.path()
            )
        })?;
        writer.write_batch(&entries).with_context(|| {
            format!(
                "write obsolete row entries failed, file={:?}, batch_size={}",
                self.file.path(),
                entries.len()
            )
        })?;
        Ok(())
    }

    fn finalize_writer(&mut self) -> Result<()> {
        self.flush_buffered_entries()?;
        if let Some(writer) = self.writer.take() {
            writer.close().with_context(|| {
                format!(
                    "close obsolete row temp file failed, file={:?}",
                    self.file.path()
                )
            })?;
        }
        Ok(())
    }
}

pub(super) struct ObsoleteRowEntryReader {
    reader: RowEntryReader,
}

impl ObsoleteRowEntryReader {
    pub(super) fn next_entry(&mut self) -> Result<Option<RowEntry>> {
        self.reader.next_entry()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_obsolete_key_temp_file_round_trip() -> Result<()> {
        let mut file = ObsoleteRowEntryTempFile::new()?;
        let entries = vec![
            RowEntry::new(Vec::new(), false, 1),
            RowEntry::new(vec![0, 1, 2], false, 2),
            RowEntry::new(b"a".to_vec(), true, 3),
        ];
        for entry in &entries {
            file.push(entry)?;
        }
        file.finalize()?;

        let mut reader = file.reader()?;
        let mut read_entries = Vec::new();
        while let Some(entry) = reader.next_entry()? {
            read_entries.push(entry);
        }
        assert_eq!(read_entries, entries);
        Ok(())
    }

    #[test]
    fn test_obsolete_key_reader_empty_file_returns_none() -> Result<()> {
        let mut file = ObsoleteRowEntryTempFile::new()?;
        let mut reader = file.reader()?;
        assert_eq!(reader.next_entry()?, None);
        Ok(())
    }

    #[test]
    fn test_obsolete_key_reader_finalizes_pending_entries() -> Result<()> {
        let mut file = ObsoleteRowEntryTempFile::new()?;
        let entries = vec![
            RowEntry::new(b"a".to_vec(), false, 11),
            RowEntry::new(b"b".to_vec(), true, 12),
        ];
        for entry in &entries {
            file.push(entry)?;
        }

        let mut reader = file.reader()?;
        let mut read_entries = Vec::new();
        while let Some(entry) = reader.next_entry()? {
            read_entries.push(entry);
        }
        assert_eq!(read_entries, entries);
        Ok(())
    }

    #[test]
    fn test_obsolete_key_push_after_finalize_returns_error() -> Result<()> {
        let mut file = ObsoleteRowEntryTempFile::new()?;
        file.push(&RowEntry::new(b"a".to_vec(), false, 0))?;
        file.finalize()?;

        let err = file
            .push(&RowEntry::new(b"b".to_vec(), true, 0))
            .unwrap_err();
        assert!(
            err.to_string()
                .contains("obsolete row temp file is already finalized")
        );
        Ok(())
    }
}