use std::{
    path::{Path, PathBuf},
    sync::Arc,
};

use anyhow::{Context, Result, anyhow};
use arrow::{
    array::{Array, ArrayRef, BinaryArray, BooleanArray, UInt64Array},
    record_batch::RecordBatch,
};
use parquet::{
    arrow::{
        ArrowWriter,
        arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder},
    },
    file::properties::WriterProperties,
};

const ROW_ENTRY_MAX_ROW_GROUP_SIZE: usize = 65_536;

#[derive(PartialEq, Debug, Clone)]
pub struct RowEntry {
    pub key: Vec<u8>,
    pub del_mark: bool,
    pub version: u64,
}

impl RowEntry {
    pub fn new(key: Vec<u8>, del_mark: bool, version: u64) -> Self {
        Self {
            key,
            del_mark,
            version,
        }
    }

    pub fn write_all<P: AsRef<Path> + std::fmt::Debug>(
        fname: P,
        entries: &[RowEntry],
    ) -> Result<()> {
        let mut writer = RowEntryWriter::create(&fname)?;
        writer.write_batch(entries)?;
        writer.close()?;
        Ok(())
    }

    #[cfg(test)]
    pub fn read_all<P: AsRef<Path> + std::fmt::Debug>(fname: P) -> Result<Vec<RowEntry>> {
        let mut reader = RowEntryReader::open(fname)?;
        let mut keys = Vec::with_capacity(reader.num_rows());
        while let Some(entry) = reader.next_entry()? {
            keys.push(entry);
        }
        Ok(keys)
    }
}

pub struct RowEntryWriter {
    fname: PathBuf,
    schema: Arc<arrow::datatypes::Schema>,
    writer: ArrowWriter<std::fs::File>,
    last_key: Option<Vec<u8>>,
    require_strictly_increasing_keys: bool,
}

impl RowEntryWriter {
    pub fn create<P: AsRef<Path> + std::fmt::Debug>(fname: P) -> Result<Self> {
        Self::create_with_validation(fname, true)
    }

    pub fn create_unchecked<P: AsRef<Path> + std::fmt::Debug>(fname: P) -> Result<Self> {
        Self::create_with_validation(fname, false)
    }

    fn create_with_validation<P: AsRef<Path> + std::fmt::Debug>(
        fname: P,
        require_strictly_increasing_keys: bool,
    ) -> Result<Self> {
        let fname = fname.as_ref().to_path_buf();
        let schema = row_entry_schema();
        let props = row_entry_writer_props();
        let file = std::fs::File::create(&fname)
            .with_context(|| format!("Create File failed={:?}", fname))?;
        let writer = ArrowWriter::try_new(file, schema.clone(), Some(props))
            .with_context(|| format!("Create ArrowWriter failed, file={:?}", fname))?;

        Ok(Self {
            fname,
            schema,
            writer,
            last_key: None,
            require_strictly_increasing_keys,
        })
    }

    pub fn write_batch(&mut self, entries: &[RowEntry]) -> Result<()> {
        if entries.is_empty() {
            return Ok(());
        }

        if self.require_strictly_increasing_keys {
            self.validate_batch(entries)?;
        }
        let keys: Vec<&[u8]> = entries.iter().map(|e| e.key.as_slice()).collect();
        let del_marks: Vec<bool> = entries.iter().map(|e| e.del_mark).collect();
        let versions: Vec<u64> = entries.iter().map(|e| e.version).collect();
        let binary_array = BinaryArray::from(keys);
        let del_mark_array = BooleanArray::from(del_marks);
        let version_array = UInt64Array::from(versions);
        let batch = RecordBatch::try_new(
            self.schema.clone(),
            vec![
                Arc::new(binary_array) as ArrayRef,
                Arc::new(del_mark_array) as ArrayRef,
                Arc::new(version_array) as ArrayRef,
            ],
        )
        .with_context(|| format!("Create RecordBatch failed, file={:?}", self.fname))?;

        self.writer
            .write(&batch)
            .with_context(|| format!("Write RecordBatch failed, file={:?}", self.fname))?;
        self.last_key = entries.last().map(|entry| entry.key.clone());
        Ok(())
    }

    pub fn close(self) -> Result<()> {
        self.writer
            .close()
            .with_context(|| format!("Close ArrowWriter failed, file={:?}", self.fname))?;
        Ok(())
    }

    fn validate_batch(&self, entries: &[RowEntry]) -> Result<()> {
        let mut prev_key = self.last_key.as_deref();

        for entry in entries {
            if let Some(prev_key) = prev_key {
                if prev_key >= entry.key.as_slice() {
                    anyhow::bail!(
                        "row entries must be strictly increasing, file={:?}, prev_key={:?}, current_key={:?}",
                        self.fname,
                        prev_key,
                        entry.key
                    );
                }
            }
            prev_key = Some(entry.key.as_slice());
        }

        Ok(())
    }
}

fn row_entry_schema() -> Arc<arrow::datatypes::Schema> {
    Arc::new(arrow::datatypes::Schema::new(vec![
        arrow::datatypes::Field::new("key", arrow::datatypes::DataType::Binary, false),
        arrow::datatypes::Field::new("del_mark", arrow::datatypes::DataType::Boolean, false),
        arrow::datatypes::Field::new("version", arrow::datatypes::DataType::UInt64, false),
    ]))
}

fn row_entry_writer_props() -> WriterProperties {
    WriterProperties::builder()
        .set_compression(parquet::basic::Compression::SNAPPY)
        .set_max_row_group_size(ROW_ENTRY_MAX_ROW_GROUP_SIZE)
        .build()
}

pub struct RowEntryReader {
    fname: PathBuf,
    schema: Arc<arrow::datatypes::Schema>,
    reader: ParquetRecordBatchReader,
    key_col: Option<BinaryArray>,
    del_mark_col: Option<BooleanArray>,
    version_col: Option<UInt64Array>,
    row_idx: usize,
    batch_len: usize,
    num_rows: usize,
}

impl RowEntryReader {
    pub fn open<P: AsRef<Path> + std::fmt::Debug>(fname: P) -> Result<Self> {
        let fname = fname.as_ref().to_path_buf();
        let file = std::fs::File::open(&fname)
            .with_context(|| format!("Open file failed: {:?}", fname))?;

        let builder = ParquetRecordBatchReaderBuilder::try_new(file).with_context(|| {
            format!("Create ParquetRecordBatchReaderBuilder failed: {:?}", fname)
        })?;
        let schema = builder.schema().clone();
        let num_rows: usize = builder
            .metadata()
            .file_metadata()
            .num_rows()
            .try_into()
            .map_err(|e| anyhow!("Invalid row count in file ({:?}): {:?}", fname, e))?;
        let reader = builder
            .build()
            .with_context(|| format!("Build reader failed, file={:?}", fname))?;

        Ok(Self {
            fname,
            schema,
            reader,
            key_col: None,
            del_mark_col: None,
            version_col: None,
            row_idx: 0,
            batch_len: 0,
            num_rows,
        })
    }

    pub fn num_rows(&self) -> usize {
        self.num_rows
    }

    pub fn next_entry(&mut self) -> Result<Option<RowEntry>> {
        loop {
            if let (Some(key_col), Some(del_mark_col), Some(version_col)) =
                (&self.key_col, &self.del_mark_col, &self.version_col)
            {
                if self.row_idx < self.batch_len {
                    let i = self.row_idx;
                    self.row_idx += 1;
                    return Ok(Some(RowEntry::new(
                        key_col.value(i).to_vec(),
                        del_mark_col.value(i),
                        version_col.value(i),
                    )));
                }
                self.key_col = None;
                self.del_mark_col = None;
                self.version_col = None;
                self.row_idx = 0;
                self.batch_len = 0;
            }

            let batch_res = match self.reader.next() {
                Some(batch_res) => batch_res,
                None => return Ok(None),
            };
            let batch = batch_res
                .with_context(|| format!("Read RecordBatch failed, file={:?}", self.fname))?;
            self.load_batch(batch)?;
        }
    }

    fn load_batch(&mut self, batch: RecordBatch) -> Result<()> {
        let key_col = batch.column_by_name("key").ok_or_else(|| {
            anyhow!(
                "Key column not found in file ({:?}), schema={:?}",
                self.fname,
                self.schema
            )
        })?;
        let key_col = key_col
            .as_any()
            .downcast_ref::<BinaryArray>()
            .ok_or_else(|| {
                anyhow!(
                    "Unexpected key column type in file ({:?})={:?}",
                    self.fname,
                    key_col.data_type()
                )
            })?;

        let del_mark_col = batch.column_by_name("del_mark").ok_or_else(|| {
            anyhow!(
                "del_mark column not found in file ({:?}), schema={:?}",
                self.fname,
                self.schema
            )
        })?;
        let del_mark_col = del_mark_col
            .as_any()
            .downcast_ref::<BooleanArray>()
            .ok_or_else(|| {
                anyhow!(
                    "Unexpected del_mark column type in Parquet file ({:?})={:?}",
                    self.fname,
                    del_mark_col.data_type()
                )
            })?;

        let version_col = batch.column_by_name("version").ok_or_else(|| {
            anyhow!(
                "version column not found in file ({:?}), schema={:?}",
                self.fname,
                self.schema
            )
        })?;
        let version_col = version_col
            .as_any()
            .downcast_ref::<UInt64Array>()
            .ok_or_else(|| {
                anyhow!(
                    "Unexpected version column type in Parquet file ({:?})={:?}",
                    self.fname,
                    version_col.data_type()
                )
            })?;

        if key_col.len() != del_mark_col.len() {
            anyhow::bail!(
                "Key/del_mark column length mismatch in file ({:?}): key_len={}, del_mark_len={}",
                self.fname,
                key_col.len(),
                del_mark_col.len()
            );
        }
        if key_col.len() != version_col.len() {
            anyhow::bail!(
                "Key/version column length mismatch in file ({:?}): key_len={}, version_len={}",
                self.fname,
                key_col.len(),
                version_col.len()
            );
        }

        self.batch_len = key_col.len();
        self.key_col = Some(key_col.clone());
        self.del_mark_col = Some(del_mark_col.clone());
        self.version_col = Some(version_col.clone());
        self.row_idx = 0;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use anyhow::bail;
    use tempfile::TempDir;

    use super::*;

    #[test]
    fn test_round_trip_basic() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("basic_test.parquet");
        let entries = vec![
            RowEntry::new(b"user_001".to_vec(), false, 11),
            RowEntry::new(b"user_002".to_vec(), true, 22),
            RowEntry::new(vec![0xde, 0xad, 0xbe, 0xef], false, 33),
        ];
        RowEntry::write_all(&file_path, &entries)?;
        let read_entries = RowEntry::read_all(&file_path)?;
        assert_eq!(entries.len(), read_entries.len());
        assert_eq!(entries, read_entries);
        Ok(())
    }

    #[test]
    fn test_writer_streaming_batches() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("streaming_batches.parquet");
        let entries = vec![
            RowEntry::new(b"user_001".to_vec(), false, 1),
            RowEntry::new(b"user_002".to_vec(), true, 2),
            RowEntry::new(b"user_003".to_vec(), false, 3),
            RowEntry::new(b"user_004".to_vec(), true, 4),
        ];

        let mut writer = RowEntryWriter::create(&file_path)?;
        writer.write_batch(&entries[..2])?;
        writer.write_batch(&entries[2..])?;
        writer.close()?;

        let read_entries = RowEntry::read_all(&file_path)?;
        assert_eq!(entries, read_entries);
        Ok(())
    }

    #[test]
    fn test_empty_entries() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("empty.parquet");
        let entries: Vec<RowEntry> = vec![];
        RowEntry::write_all(&file_path, &entries)?;
        let read_entries = RowEntry::read_all(&file_path)?;
        assert!(read_entries.is_empty());
        Ok(())
    }

    #[test]
    fn test_large_batch() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("large.parquet");
        let mut entries = Vec::with_capacity(2000);
        for i in 0..2000 {
            entries.push(RowEntry::new(
                format!("key_{:04}", i).into_bytes(),
                i % 2 == 0,
                0,
            ));
        }
        RowEntry::write_all(&file_path, &entries)?;
        let read_entries = RowEntry::read_all(&file_path)?;
        assert_eq!(entries.len(), read_entries.len());
        assert_eq!(entries[0], read_entries[0]);
        assert_eq!(entries[1999], read_entries[1999]);

        Ok(())
    }

    #[test]
    fn test_reader_stream_round_trip() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("stream.parquet");
        let entries = vec![
            RowEntry::new(b"user_001".to_vec(), false, 7),
            RowEntry::new(b"user_002".to_vec(), true, 8),
            RowEntry::new(vec![0xde, 0xad, 0xbe, 0xef], false, 9),
        ];
        RowEntry::write_all(&file_path, &entries)?;

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

    #[test]
    fn test_writer_rejects_duplicate_keys_in_batch() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("duplicate_in_batch.parquet");
        let entries = vec![
            RowEntry::new(b"user_001".to_vec(), false, 0),
            RowEntry::new(b"user_001".to_vec(), true, 0),
        ];

        let err = RowEntry::write_all(&file_path, &entries).unwrap_err();
        assert!(
            err.to_string()
                .contains("row entries must be strictly increasing")
        );
        Ok(())
    }

    #[test]
    fn test_writer_rejects_unsorted_keys_in_batch() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("unsorted_in_batch.parquet");
        let entries = vec![
            RowEntry::new(b"user_002".to_vec(), false, 0),
            RowEntry::new(b"user_001".to_vec(), true, 0),
        ];

        let err = RowEntry::write_all(&file_path, &entries).unwrap_err();
        assert!(
            err.to_string()
                .contains("row entries must be strictly increasing")
        );
        Ok(())
    }

    #[test]
    fn test_writer_rejects_duplicate_keys_across_batches() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("duplicate_across_batches.parquet");
        let mut writer = RowEntryWriter::create(&file_path)?;

        writer.write_batch(&[RowEntry::new(b"user_001".to_vec(), false, 0)])?;
        let err = writer
            .write_batch(&[RowEntry::new(b"user_001".to_vec(), true, 0)])
            .unwrap_err();
        assert!(
            err.to_string()
                .contains("row entries must be strictly increasing")
        );
        Ok(())
    }

    #[test]
    fn test_writer_rejects_unsorted_keys_across_batches() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("unsorted_across_batches.parquet");
        let mut writer = RowEntryWriter::create(&file_path)?;

        writer.write_batch(&[RowEntry::new(b"user_002".to_vec(), false, 0)])?;
        let err = writer
            .write_batch(&[RowEntry::new(b"user_001".to_vec(), true, 0)])
            .unwrap_err();
        assert!(
            err.to_string()
                .contains("row entries must be strictly increasing")
        );
        Ok(())
    }

    #[test]
    fn test_writer_unchecked_allows_unsorted_keys_across_batches() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir
            .path()
            .join("unsorted_across_batches_unchecked.parquet");
        let entries = vec![
            RowEntry::new(b"user_002".to_vec(), false, 0),
            RowEntry::new(b"user_001".to_vec(), true, 0),
        ];

        let mut writer = RowEntryWriter::create_unchecked(&file_path)?;
        writer.write_batch(&[entries[0].clone()])?;
        writer.write_batch(&[entries[1].clone()])?;
        writer.close()?;

        let read_entries = RowEntry::read_all(&file_path)?;
        assert_eq!(read_entries, entries);
        Ok(())
    }

    #[test]
    fn test_reader_stream_empty() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let file_path = temp_dir.path().join("stream_empty.parquet");
        let entries: Vec<RowEntry> = vec![];
        RowEntry::write_all(&file_path, &entries)?;

        let mut reader = RowEntryReader::open(&file_path)?;
        assert_eq!(reader.num_rows(), 0);
        assert!(reader.next_entry()?.is_none());
        Ok(())
    }

    #[test]
    fn test_file_not_found() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let non_existent_path = temp_dir.path().join("ghost.parquet");
        let result = RowEntry::read_all(&non_existent_path);
        if result.is_ok() {
            bail!("Expected read_all to fail for non-existent file, but it succeeded");
        }
        let err = result.unwrap_err();
        let err_msg = err.to_string();
        assert!(
            err_msg.contains("Open file failed"),
            "Error message was: {}",
            err_msg
        );
        assert!(
            err_msg.contains("ghost.parquet"),
            "Error message was: {}",
            err_msg
        );
        Ok(())
    }
}