// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

//! Avro writer implementation for the `arrow-avro` crate.

//!

//! # Overview

//!

//! *   Use **`AvroWriter`** (Object Container File) when you want a

//!     self‑contained Avro file with header, schema JSON, optional compression,

//!     blocks, and sync markers.

//! *   Use **`AvroStreamWriter`** (raw binary stream) when you already know the

//!     schema out‑of‑band (i.e., via a schema registry) and need a stream

//!     of Avro‑encoded records with minimal framing.

//!

/// Encodes `RecordBatch` into the Avro binary format.

pub mod encoder;

/// Logic for different Avro container file formats.

pub mod format;

use crate::compression::CompressionCodec;

use crate::schema::AvroSchema;

use crate::writer::encoder::{encode_record_batch, write_long};

use crate::writer::format::{AvroBinaryFormat, AvroFormat, AvroOcfFormat};

use arrow_array::RecordBatch;

use arrow_schema::{ArrowError, Schema};

use std::io::{self, Write};

use std::sync::Arc;

/// Builder to configure and create a `Writer`.

#[derive(Debug, Clone)]

pub struct WriterBuilder {

    schema: Schema,

    codec: Option<CompressionCodec>,

}

impl WriterBuilder {

    /// Create a new builder with default settings.

    pub fn new(schema: Schema) -> Self {

        Self {

            schema,

            codec: None,

        }

    }

    /// Change the compression codec.

    pub fn with_compression(mut self, codec: Option<CompressionCodec>) -> Self {

        self.codec = codec;

        self

    }

    /// Create a new `Writer` with specified `AvroFormat` and builder options.

    pub fn build<W, F>(self, writer: W) -> Writer<W, F>

    where

        W: Write,

        F: AvroFormat,

    {

        Writer {

            writer,

            schema: Arc::from(self.schema),

            format: F::default(),

            compression: self.codec,

            started: false,

        }

    }

}

/// Generic Avro writer.

#[derive(Debug)]

pub struct Writer<W: Write, F: AvroFormat> {

    writer: W,

    schema: Arc<Schema>,

    format: F,

    compression: Option<CompressionCodec>,

    started: bool,

}

/// Alias for an Avro **Object Container File** writer.

pub type AvroWriter<W> = Writer<W, AvroOcfFormat>;

/// Alias for a raw Avro **binary stream** writer.

pub type AvroStreamWriter<W> = Writer<W, AvroBinaryFormat>;

impl<W: Write> Writer<W, AvroOcfFormat> {

    /// Convenience constructor – same as

    pub fn new(writer: W, schema: Schema) -> Result<Self, ArrowError> {

        Ok(WriterBuilder::new(schema).build::<W, AvroOcfFormat>(writer))

    }

    /// Change the compression codec after construction.

    pub fn with_compression(mut self, codec: Option<CompressionCodec>) -> Self {

        self.compression = codec;

        self

    }

    /// Return a reference to the 16‑byte sync marker generated for this file.

    pub fn sync_marker(&self) -> Option<&[u8; 16]> {

        self.format.sync_marker()

    }

}

impl<W: Write> Writer<W, AvroBinaryFormat> {

    /// Convenience constructor to create a new [`AvroStreamWriter`].

    pub fn new(writer: W, schema: Schema) -> Result<Self, ArrowError> {

        Ok(WriterBuilder::new(schema).build::<W, AvroBinaryFormat>(writer))

    }

}

impl<W: Write, F: AvroFormat> Writer<W, F> {

    /// Serialize one [`RecordBatch`] to the output.

    pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {

        if !self.started {

            self.format

                .start_stream(&mut self.writer, &self.schema, self.compression)
?0
;

            self.started = true;

        }

        if batch.schema() != self.schema {

            return Err(ArrowError::SchemaError(

                "Schema of RecordBatch differs from Writer schema".to_string(),

            ));

        }

        match self.format.sync_marker() {

            Some(&sync) => self.write_ocf_block(batch, &sync),

            None => self.write_stream(batch),

        }

    }

    /// A convenience method to write a slice of [`RecordBatch`].

    ///

    /// This is equivalent to calling `write` for each batch in the slice.

    pub fn write_batches(&mut self, batches: &[&RecordBatch]) -> Result<(), ArrowError> {

        for 
b2
 in batches {

            self.write(b)
?0
;

        }

        Ok(())

    }

    /// Flush remaining buffered data and (for OCF) ensure the header is present.

    pub fn finish(&mut self) -> Result<(), ArrowError> {

        if !self.started {

            self.format

                .start_stream(&mut self.writer, &self.schema, self.compression)
?0
;

            self.started = true;

        }

        self.writer

            .flush()

            .map_err(|e| ArrowError::IoError(
format!0
(
"Error flushing writer: {e}"0
), 
e0
))

    }

    /// Consume the writer, returning the underlying output object.

    pub fn into_inner(self) -> W {

        self.writer

    }

    fn write_ocf_block(&mut self, batch: &RecordBatch, sync: &[u8; 16]) -> Result<(), ArrowError> {

        let mut buf = Vec::<u8>::with_capacity(1024);

        encode_record_batch(batch, &mut buf)
?0
;

        let encoded = match self.compression {

            Some(codec) => codec.compress(&buf)
?0
,

            None => buf,

        };

        write_long(&mut self.writer, batch.num_rows() as i64)
?0
;

        write_long(&mut self.writer, encoded.len() as i64)
?0
;

        self.writer

            .write_all(&encoded)

            .map_err(|e| ArrowError::IoError(
format!0
(
"Error writing Avro block: {e}"0
), 
e0
))
?0
;

        self.writer

            .write_all(sync)

            .map_err(|e| ArrowError::IoError(
format!0
(
"Error writing Avro sync: {e}"0
), 
e0
))
?0
;

        Ok(())

    }

    fn write_stream(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {

        encode_record_batch(batch, &mut self.writer)

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    use crate::reader::ReaderBuilder;

    use crate::test_util::arrow_test_data;

    use arrow_array::{ArrayRef, BinaryArray, Int32Array, RecordBatch, StringArray};

    use arrow_schema::{DataType, Field, Schema};

    use std::fs::File;

    use std::io::BufReader;

    use std::sync::Arc;

    use tempfile::NamedTempFile;

    fn make_schema() -> Schema {

        Schema::new(vec![

            Field::new("id", DataType::Int32, false),

            Field::new("name", DataType::Binary, false),

        ])

    }

    fn make_batch() -> RecordBatch {

        let ids = Int32Array::from(vec![1, 2, 3]);

        let names = BinaryArray::from_vec(vec![b"a".as_ref(), b"b".as_ref(), b"c".as_ref()]);

        RecordBatch::try_new(

            Arc::new(make_schema()),

            vec![Arc::new(ids) as ArrayRef, Arc::new(names) as ArrayRef],

        )

        .expect("failed to build test RecordBatch")

    }

    fn contains_ascii(haystack: &[u8], needle: &[u8]) -> bool {

        haystack.windows(needle.len()).any(|w| w == needle)

    }

    #[test]

    fn test_ocf_writer_generates_header_and_sync() -> Result<(), ArrowError> {

        let batch = make_batch();

        let buffer: Vec<u8> = Vec::new();

        let mut writer = AvroWriter::new(buffer, make_schema())
?0
;

        writer.write(&batch)
?0
;

        writer.finish()
?0
;

        let out = writer.into_inner();

        assert_eq!(&out[..4], b"Obj\x01", 
"OCF magic bytes missing/incorrect"0
);

        let sync = AvroWriter::new(Vec::new(), make_schema())
?0

            .sync_marker()

            .cloned();

        let trailer = &out[out.len() - 16..];

        assert_eq!(trailer.len(), 16, 
"expected 16‑byte sync marker"0
);

        let _ = sync;

        Ok(())

    }

    #[test]

    fn test_schema_mismatch_yields_error() {

        let batch = make_batch();

        let alt_schema = Schema::new(vec![Field::new("x", DataType::Int32, false)]);

        let buffer = Vec::<u8>::new();

        let mut writer = AvroWriter::new(buffer, alt_schema).unwrap();

        let err = writer.write(&batch).unwrap_err();

        assert!(
matches!0
(err, ArrowError::SchemaError(_)));

    }

    #[test]

    fn test_write_batches_accumulates_multiple() -> Result<(), ArrowError> {

        let batch1 = make_batch();

        let batch2 = make_batch();

        let buffer = Vec::<u8>::new();

        let mut writer = AvroWriter::new(buffer, make_schema())
?0
;

        writer.write_batches(&[&batch1, &batch2])
?0
;

        writer.finish()
?0
;

        let out = writer.into_inner();

        assert!(out.len() > 4, 
"combined batches produced tiny file"0
);

        Ok(())

    }

    #[test]

    fn test_finish_without_write_adds_header() -> Result<(), ArrowError> {

        let buffer = Vec::<u8>::new();

        let mut writer = AvroWriter::new(buffer, make_schema())
?0
;

        writer.finish()
?0
;

        let out = writer.into_inner();

        assert_eq!(&out[..4], b"Obj\x01", 
"finish() should emit OCF header"0
);

        Ok(())

    }

    #[test]

    fn test_write_long_encodes_zigzag_varint() -> Result<(), ArrowError> {

        let mut buf = Vec::new();

        write_long(&mut buf, 0)
?0
;

        write_long(&mut buf, -1)
?0
;

        write_long(&mut buf, 1)
?0
;

        write_long(&mut buf, -2)
?0
;

        write_long(&mut buf, 2147483647)
?0
;

        assert!(

            buf.starts_with(&[0x00, 0x01, 0x02, 0x03]),

            "zig‑zag varint encodings incorrect: {buf:?}"

        );

        Ok(())

    }

    #[test]

    fn test_roundtrip_alltypes_roundtrip_writer() -> Result<(), ArrowError> {

        let files = [

            "avro/alltypes_plain.avro",

            "avro/alltypes_plain.snappy.avro",

            "avro/alltypes_plain.zstandard.avro",

            "avro/alltypes_plain.bzip2.avro",

            "avro/alltypes_plain.xz.avro",

        ];

        for 
rel5
 in files {

            let path = arrow_test_data(rel);

            let rdr_file = File::open(&path).expect("open input avro");

            let mut reader = ReaderBuilder::new()

                .build(BufReader::new(rdr_file))

                .expect("build reader");

            let schema = reader.schema();

            let input_batches = reader.collect::<Result<Vec<_>, _>>()
?0
;

            let original =

                arrow::compute::concat_batches(&schema, &input_batches).expect("concat input");

            let tmp = NamedTempFile::new().expect("create temp file");

            let out_path = tmp.into_temp_path();

            let out_file = File::create(&out_path).expect("create temp avro");

            let mut writer = AvroWriter::new(out_file, original.schema().as_ref().clone())
?0
;

            if rel.contains(".snappy.") {

                writer = writer.with_compression(Some(CompressionCodec::Snappy));

            } else if rel.contains(".zstandard.") {

                writer = writer.with_compression(Some(CompressionCodec::ZStandard));

            } else if rel.contains(".bzip2.") {

                writer = writer.with_compression(Some(CompressionCodec::Bzip2));

            } else if rel.contains(".xz.") {

                writer = writer.with_compression(Some(CompressionCodec::Xz));

            }

            writer.write(&original)
?0
;

            writer.finish()
?0
;

            drop(writer);

            let rt_file = File::open(&out_path).expect("open roundtrip avro");

            let mut rt_reader = ReaderBuilder::new()

                .build(BufReader::new(rt_file))

                .expect("build roundtrip reader");

            let rt_schema = rt_reader.schema();

            let rt_batches = rt_reader.collect::<Result<Vec<_>, _>>()
?0
;

            let roundtrip =

                arrow::compute::concat_batches(&rt_schema, &rt_batches).expect("concat roundtrip");

            assert_eq!(

                roundtrip, original,

                "Round-trip batch mismatch for file: {}",

                rel

            );

        }

        Ok(())

    }

}