chdlady-av 0.1.0

//! AVI RIFF container writer.

use super::riff::{write_fourcc, write_u16le, write_u32le, AviError, AviInfo};
use std::io::{Seek, SeekFrom, Write};

/// Writes a minimal AVI file with one YUY2 video stream and optional audio stream.
pub struct AviWriter<W: Write + Seek> {
    inner: W,
    info: AviInfo,
    /// Byte offset of the RIFF chunk size field.
    riff_size_pos: u64,
    /// Byte offset of the movi LIST chunk size field.
    movi_size_pos: u64,
    /// Byte offset of the avih dwTotalFrames field.
    avih_frames_pos: u64,
    /// Byte offset of the video strh dwLength field.
    strh_video_length_pos: u64,
    /// Byte offset of the audio strh dwLength field.
    strh_audio_length_pos: Option<u64>,
    /// Number of video frames written.
    frames_written: u32,
    /// Number of audio samples written per channel.
    audio_samples_written: u32,
    /// Index entries accumulated for the idx1 chunk.
    index_entries: Vec<IndexEntry>,
    /// Byte offset of the movi data start (after the 'movi' fourcc+size).
    movi_data_start: u64,
    /// Pending per-channel audio samples waiting to be flushed as interleaved chunks.
    audio_buffers: Vec<Vec<i16>>,
}

struct IndexEntry {
    fourcc: [u8; 4],
    flags: u32,
    offset: u32,
    size: u32,
}

const AVIIF_KEYFRAME: u32 = 0x10;

impl<W: Write + Seek> AviWriter<W> {
    /// Create a new AVI writer.
    pub fn new(mut inner: W, info: AviInfo) -> Result<Self, AviError> {
        let channels = info.channels as usize;
        let fps_us = 1_000_000_000_000u64 / info.fps_times_1million as u64;

        // --- RIFF AVI header ---
        write_fourcc(&mut inner, b"RIFF")?;
        let riff_size_pos = inner.stream_position()?;
        write_u32le(&mut inner, 0)?; // placeholder
        write_fourcc(&mut inner, b"AVI ")?;

        // --- LIST hdrl ---
        write_fourcc(&mut inner, b"LIST")?;
        let hdrl_size_pos = inner.stream_position()?;
        write_u32le(&mut inner, 0)?; // placeholder
        write_fourcc(&mut inner, b"hdrl")?;

        // avih (main AVI header, 56 bytes)
        write_fourcc(&mut inner, b"avih")?;
        write_u32le(&mut inner, 56)?;
        write_u32le(&mut inner, fps_us as u32)?; // dwMicroSecPerFrame
        write_u32le(&mut inner, 0)?; // dwMaxBytesPerSec
        write_u32le(&mut inner, 0)?; // dwPaddingGranularity
        write_u32le(&mut inner, 0x10)?; // dwFlags = AVIF_HASINDEX
        let avih_frames_pos = inner.stream_position()?;
        write_u32le(&mut inner, 0)?; // dwTotalFrames placeholder
        write_u32le(&mut inner, 0)?; // dwInitialFrames
        write_u32le(&mut inner, 1 + if channels > 0 { 1 } else { 0 })?; // dwStreams
        write_u32le(&mut inner, 0)?; // dwSuggestedBufferSize
        write_u32le(&mut inner, info.width)?; // dwWidth
        write_u32le(&mut inner, info.height)?; // dwHeight
        write_u32le(&mut inner, 0)?; // reserved[4]
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;

        // --- Video stream LIST strl ---
        write_fourcc(&mut inner, b"LIST")?;
        let strl_video_size_pos = inner.stream_position()?;
        write_u32le(&mut inner, 0)?;
        write_fourcc(&mut inner, b"strl")?;

        // strh (video stream header, 56 bytes)
        write_fourcc(&mut inner, b"strh")?;
        write_u32le(&mut inner, 56)?;
        write_fourcc(&mut inner, b"vids")?;
        write_fourcc(&mut inner, b"YUY2")?;
        write_u32le(&mut inner, 0)?;
        write_u16le(&mut inner, 0)?;
        write_u16le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 1_000_000)?;
        write_u32le(&mut inner, info.fps_times_1million)?;
        write_u32le(&mut inner, 0)?;
        let strh_video_length_pos = inner.stream_position()?;
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;
        write_u16le(&mut inner, 0)?;
        write_u16le(&mut inner, 0)?;
        write_u16le(&mut inner, info.width as u16)?;
        write_u16le(&mut inner, info.height as u16)?;

        // strf: BITMAPINFOHEADER (40 bytes)
        write_fourcc(&mut inner, b"strf")?;
        write_u32le(&mut inner, 40)?;
        write_u32le(&mut inner, 40)?; // biSize
        write_u32le(&mut inner, info.width)?;
        write_u32le(&mut inner, info.height)?;
        write_u16le(&mut inner, 1)?;
        write_u16le(&mut inner, 16)?;
        write_fourcc(&mut inner, b"YUY2")?;
        write_u32le(&mut inner, info.width * info.height * 2)?;
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;
        write_u32le(&mut inner, 0)?;

        let after_strl_video = inner.stream_position()?;
        inner.seek(SeekFrom::Start(strl_video_size_pos))?;
        write_u32le(
            &mut inner,
            (after_strl_video - strl_video_size_pos - 4) as u32,
        )?;
        inner.seek(SeekFrom::Start(after_strl_video))?;

        // --- Audio stream LIST strl ---
        let mut strh_audio_length_pos = None;
        if channels > 0 {
            write_fourcc(&mut inner, b"LIST")?;
            let strl_audio_size_pos = inner.stream_position()?;
            write_u32le(&mut inner, 0)?;
            write_fourcc(&mut inner, b"strl")?;

            // strh (audio, 56 bytes)
            write_fourcc(&mut inner, b"strh")?;
            write_u32le(&mut inner, 56)?;
            write_fourcc(&mut inner, b"auds")?;
            write_u32le(&mut inner, 0)?;
            write_u32le(&mut inner, 0)?;
            write_u16le(&mut inner, 0)?;
            write_u16le(&mut inner, 0)?;
            write_u32le(&mut inner, 0)?;
            write_u32le(&mut inner, 1)?;
            write_u32le(&mut inner, info.sample_rate)?;
            write_u32le(&mut inner, 0)?;
            let audio_len_pos = inner.stream_position()?;
            write_u32le(&mut inner, 0)?;
            write_u32le(&mut inner, 0)?;
            write_u32le(&mut inner, 0)?;
            write_u32le(&mut inner, 2 * channels as u32)?;
            write_u16le(&mut inner, 0)?;
            write_u16le(&mut inner, 0)?;
            write_u16le(&mut inner, 0)?;
            write_u16le(&mut inner, 0)?;
            strh_audio_length_pos = Some(audio_len_pos);

            // strf: WAVEFORMATEX (18 bytes)
            write_fourcc(&mut inner, b"strf")?;
            write_u32le(&mut inner, 18)?;
            write_u16le(&mut inner, 1)?; // WAVE_FORMAT_PCM
            write_u16le(&mut inner, channels as u16)?;
            write_u32le(&mut inner, info.sample_rate)?;
            write_u32le(&mut inner, info.sample_rate * channels as u32 * 2)?;
            write_u16le(&mut inner, channels as u16 * 2)?;
            write_u16le(&mut inner, 16)?;
            write_u16le(&mut inner, 0)?;

            let after_strl_audio = inner.stream_position()?;
            inner.seek(SeekFrom::Start(strl_audio_size_pos))?;
            write_u32le(
                &mut inner,
                (after_strl_audio - strl_audio_size_pos - 4) as u32,
            )?;
            inner.seek(SeekFrom::Start(after_strl_audio))?;
        }

        // patch hdrl size
        let after_hdrl = inner.stream_position()?;
        inner.seek(SeekFrom::Start(hdrl_size_pos))?;
        write_u32le(&mut inner, (after_hdrl - hdrl_size_pos - 4) as u32)?;
        inner.seek(SeekFrom::Start(after_hdrl))?;

        // --- LIST movi ---
        write_fourcc(&mut inner, b"LIST")?;
        let movi_size_pos = inner.stream_position()?;
        write_u32le(&mut inner, 0)?;
        write_fourcc(&mut inner, b"movi")?;
        let movi_data_start = inner.stream_position()?;

        Ok(Self {
            inner,
            info,
            riff_size_pos,
            movi_size_pos,
            avih_frames_pos,
            strh_video_length_pos,
            strh_audio_length_pos,
            frames_written: 0,
            audio_samples_written: 0,
            index_entries: Vec::new(),
            movi_data_start,
            audio_buffers: vec![Vec::new(); channels],
        })
    }

    /// Append a YUY2 video frame. `data` is raw YUY2 bytes (width*height*2).
    pub fn append_video_frame(&mut self, data: &[u8]) -> Result<(), AviError> {
        let expected = (self.info.width * self.info.height * 2) as usize;
        if data.len() != expected {
            return Err(AviError::Invalid(format!(
                "video frame size mismatch: expected {}, got {}",
                expected,
                data.len()
            )));
        }

        let pos = self.inner.stream_position()? - self.movi_data_start;
        let size = data.len() as u32;

        write_fourcc(&mut self.inner, b"00dc")?;
        write_u32le(&mut self.inner, size)?;
        self.inner.write_all(data)?;
        if !size.is_multiple_of(2) {
            self.inner.write_all(&[0])?;
        }

        self.index_entries.push(IndexEntry {
            fourcc: *b"00dc",
            flags: AVIIF_KEYFRAME,
            offset: pos as u32,
            size,
        });
        self.frames_written += 1;
        Ok(())
    }

    /// Append mono audio samples for one channel.
    pub fn append_audio_samples(
        &mut self,
        channel: usize,
        samples: &[i16],
    ) -> Result<(), AviError> {
        let channels = self.info.channels as usize;
        if channel >= channels {
            return Err(AviError::Invalid(format!(
                "channel {} out of range (max {})",
                channel, channels
            )));
        }

        self.audio_buffers[channel].extend_from_slice(samples);

        let min_len = self
            .audio_buffers
            .iter()
            .map(|b| b.len())
            .min()
            .unwrap_or(0);
        if min_len >= 512 {
            self.flush_audio_chunk(min_len)?;
        }
        Ok(())
    }

    fn flush_audio_chunk(&mut self, samples_to_flush: usize) -> Result<(), AviError> {
        if samples_to_flush == 0 {
            return Ok(());
        }
        let channels = self.info.channels as usize;
        let pos = self.inner.stream_position()? - self.movi_data_start;
        let size = (samples_to_flush * channels * 2) as u32;

        write_fourcc(&mut self.inner, b"01wb")?;
        write_u32le(&mut self.inner, size)?;
        for s in 0..samples_to_flush {
            for ch in 0..channels {
                self.inner
                    .write_all(&self.audio_buffers[ch][s].to_le_bytes())?;
            }
        }
        if !size.is_multiple_of(2) {
            self.inner.write_all(&[0])?;
        }

        self.index_entries.push(IndexEntry {
            fourcc: *b"01wb",
            flags: 0,
            offset: pos as u32,
            size,
        });
        self.audio_samples_written += samples_to_flush as u32;

        for ch in 0..channels {
            self.audio_buffers[ch].drain(0..samples_to_flush);
        }
        Ok(())
    }

    /// Finish writing: patch all size fields and write the idx1 chunk.
    pub fn finish(mut self) -> Result<(), AviError> {
        let min_len = self
            .audio_buffers
            .iter()
            .map(|b| b.len())
            .min()
            .unwrap_or(0);
        self.flush_audio_chunk(min_len)?;

        let movi_end = self.inner.stream_position()?;
        let movi_data_size = movi_end - self.movi_data_start;

        let idx1_size = (self.index_entries.len() * 16) as u32;
        write_fourcc(&mut self.inner, b"idx1")?;
        write_u32le(&mut self.inner, idx1_size)?;
        for entry in &self.index_entries {
            write_fourcc(&mut self.inner, &entry.fourcc)?;
            write_u32le(&mut self.inner, entry.flags)?;
            write_u32le(&mut self.inner, 4 + entry.offset)?;
            write_u32le(&mut self.inner, entry.size)?;
        }

        let file_end = self.inner.stream_position()?;

        self.inner.seek(SeekFrom::Start(self.movi_size_pos))?;
        write_u32le(&mut self.inner, (movi_data_size + 4) as u32)?;

        self.inner.seek(SeekFrom::Start(self.riff_size_pos))?;
        write_u32le(&mut self.inner, (file_end - self.riff_size_pos - 4) as u32)?;

        self.inner.seek(SeekFrom::Start(self.avih_frames_pos))?;
        write_u32le(&mut self.inner, self.frames_written)?;

        self.inner
            .seek(SeekFrom::Start(self.strh_video_length_pos))?;
        write_u32le(&mut self.inner, self.frames_written)?;

        if let Some(pos) = self.strh_audio_length_pos {
            self.inner.seek(SeekFrom::Start(pos))?;
            write_u32le(&mut self.inner, self.audio_samples_written)?;
        }

        Ok(())
    }
}