use anyhow::Result;
use parquet::file::properties::WriterProperties;
use parquet::schema::parser::parse_message_type;
use parquet::record::RecordWriter;
use std::fs::File;
use std::sync::Arc;

#[derive(parquet_derive::ParquetRecordWriter)]
pub struct ScopeRecord {
    pub timestamp: f64,
    pub channel_1: f32,
}

pub struct ParquetWriter {
    writer: parquet::file::writer::SerializedFileWriter<File>,
}

impl ParquetWriter {
    pub fn new(path: &str) -> Result<Self> {
        let file = File::create(path)?;
        let schema = "
            message schema {
                REQUIRED DOUBLE timestamp;
                REQUIRED FLOAT channel_1;
            }
        ";
        let schema = Arc::new(parse_message_type(schema)?);
        let props = Arc::new(WriterProperties::builder().build());
        let writer = parquet::file::writer::SerializedFileWriter::new(file, schema, props)?;
        Ok(Self { writer })
    }

    pub fn write_batch(&mut self, records: &[ScopeRecord]) -> Result<()> {
        let mut row_group_writer = self.writer.next_row_group()?;
        records.write_to_row_group(&mut row_group_writer)?;
        row_group_writer.close()?;
        Ok(())
    }

    pub fn close(self) -> Result<()> {
        self.writer.close()?;
        Ok(())
    }
}