use std::fs::File;
use std::path::PathBuf;

use nu_plugin::EvaluatedCall;
use nu_protocol::{ShellError, shell_error::generic::GenericError};
use polars_io::SerWriter;
use polars_io::avro::{AvroCompression, AvroWriter};

use crate::command::core::resource::Resource;
use crate::values::NuDataFrame;

fn get_compression(call: &EvaluatedCall) -> Result<Option<AvroCompression>, ShellError> {
    if let Some((compression, span)) = call
        .get_flag_value("avro-compression")
        .map(|e| e.as_str().map(|s| (s.to_owned(), e.span())))
        .transpose()?
    {
        match compression.as_ref() {
            "snappy" => Ok(Some(AvroCompression::Snappy)),
            "deflate" => Ok(Some(AvroCompression::Deflate)),
            _ => Err(ShellError::IncorrectValue {
                msg: "compression must be one of deflate or snappy".to_string(),
                val_span: span,
                call_span: span,
            }),
        }
    } else {
        Ok(None)
    }
}

pub(crate) fn command_eager(
    call: &EvaluatedCall,
    df: &NuDataFrame,
    resource: Resource,
) -> Result<(), ShellError> {
    let file_span = resource.span;
    let compression = get_compression(call)?;
    let path: PathBuf = resource.as_path_buf();
    let file = File::create(&path).map_err(|e| {
        ShellError::Generic(GenericError::new(
            format!("Error with file name: {e}"),
            "",
            file_span,
        ))
    })?;

    AvroWriter::new(file)
        .with_compression(compression)
        .finish(&mut df.to_polars())
        .map_err(|e| {
            ShellError::Generic(GenericError::new(
                "Error saving file",
                e.to_string(),
                file_span,
            ))
        })?;

    Ok(())
}

#[cfg(test)]
pub mod test {
    use crate::command::core::save::test::{test_eager_save, test_lazy_save};

    #[test]
    pub fn test_avro_eager_save() -> Result<(), Box<dyn std::error::Error>> {
        test_eager_save("avro")
    }

    #[test]
    pub fn test_avro_lazy_save() -> Result<(), Box<dyn std::error::Error>> {
        test_lazy_save("avro")
    }
}
