Skip to main content

midas_audio/containers/
wav.rs

1use std::io::{self, Read, Seek, SeekFrom, Write};
2
3#[derive(Debug, Clone, Copy, PartialEq)]
4pub enum WavSampleFormat {
5    U8,
6    I16,
7}
8
9use crate::types::{BitDepth, Channels, SampleRate};
10
11pub struct WavHeader {
12    pub channels: Channels,
13    pub sample_rate: SampleRate,
14    pub bits_per_sample: BitDepth,
15    pub data_len: u32,
16}
17
18pub struct WavReader<R> {
19    reader: R,
20    pub header: WavHeader,
21    pub data_start: u64,
22}
23
24impl<R: Read + Seek> WavReader<R> {
25    pub fn new(mut reader: R) -> io::Result<Self> {
26        let mut buf = [0u8; 12];
27        reader.read_exact(&mut buf)?;
28
29        if &buf[0..4] != b"RIFF" || &buf[8..12] != b"WAVE" {
30            return Err(io::Error::new(io::ErrorKind::InvalidData, "Not a WAV file"));
31        }
32
33        // Find fmt chunk
34        let mut fmt_parsed = false;
35
36        // Temp storage for raw values
37        let mut raw_channels = 0u16;
38        let mut raw_rate = 22050;
39        let mut raw_bits = 0u16;
40
41        loop {
42            let mut chunk_head = [0u8; 8];
43            if reader.read_exact(&mut chunk_head).is_err() {
44                break;
45            }
46
47            let len =
48                u32::from_le_bytes([chunk_head[4], chunk_head[5], chunk_head[6], chunk_head[7]]);
49
50            if &chunk_head[0..4] == b"fmt " {
51                let mut fmt_data = vec![0u8; len as usize];
52                reader.read_exact(&mut fmt_data)?;
53
54                let format_tag = u16::from_le_bytes([fmt_data[0], fmt_data[1]]);
55                if format_tag != 1 {
56                    return Err(io::Error::new(
57                        io::ErrorKind::InvalidData,
58                        "Only PCM WAV supported",
59                    ));
60                }
61
62                raw_channels = u16::from_le_bytes([fmt_data[2], fmt_data[3]]);
63                raw_rate = u32::from_le_bytes([fmt_data[4], fmt_data[5], fmt_data[6], fmt_data[7]]);
64                raw_bits = u16::from_le_bytes([fmt_data[14], fmt_data[15]]);
65
66                fmt_parsed = true;
67            } else if &chunk_head[0..4] == b"data" {
68                let data_len = len;
69                let data_start = reader.stream_position()?;
70
71                if !fmt_parsed {
72                    return Err(io::Error::new(
73                        io::ErrorKind::InvalidData,
74                        "fmt chunk not found before data",
75                    ));
76                }
77
78                let channels = Channels::from_u8(raw_channels as u8).ok_or_else(|| {
79                    io::Error::new(
80                        io::ErrorKind::InvalidData,
81                        format!("Unsupported channels: {raw_channels}"),
82                    )
83                })?;
84
85                let sample_rate = SampleRate::from_u32(raw_rate).ok_or_else(|| {
86                    io::Error::new(
87                        io::ErrorKind::InvalidData,
88                        format!("Unsupported sample rate: {raw_rate}"),
89                    )
90                })?;
91
92                let bits_per_sample = BitDepth::from_u16(raw_bits).ok_or_else(|| {
93                    io::Error::new(
94                        io::ErrorKind::InvalidData,
95                        format!("Unsupported bit depth: {raw_bits}"),
96                    )
97                })?;
98
99                return Ok(Self {
100                    reader,
101                    header: WavHeader {
102                        channels,
103                        sample_rate,
104                        bits_per_sample,
105                        data_len,
106                    },
107                    data_start,
108                });
109            } else {
110                reader.seek(SeekFrom::Current(i64::from(len)))?;
111            }
112        }
113
114        Err(io::Error::new(
115            io::ErrorKind::InvalidData,
116            "No data chunk found",
117        ))
118    }
119
120    pub fn read_samples<S: SampleFromBytes>(&mut self) -> io::Result<Vec<S>> {
121        // Read up to data_len bytes
122        // Note: data_len might be huge, reading all at once might be bad for memory.
123        // For CLI converter "simple" usage, maybe ok.
124        // But let's implement just reading a buffer.
125        // This helper reads ALL.
126
127        let mut bytes = vec![0u8; self.header.data_len as usize];
128        self.reader.read_exact(&mut bytes)?;
129
130        let mut samples =
131            Vec::with_capacity(bytes.len() / (self.header.bits_per_sample as usize / 8));
132
133        // Assume 16-bit or 8-bit
134        // Assume 16-bit or 8-bit
135        match self.header.bits_per_sample {
136            BitDepth::SixteenBit => {
137                for chunk in bytes.chunks(2) {
138                    if chunk.len() == 2 {
139                        let val = i16::from_le_bytes([chunk[0], chunk[1]]);
140                        samples.push(S::from_i16(val));
141                    }
142                }
143            }
144            BitDepth::EightBit => {
145                for b in bytes {
146                    // 8-bit WAV is unsigned 0..255, center 128
147                    let val = i16::from(b) - 128;
148                    samples.push(S::from_i16(val * 256)); // Scale to i16 range approx
149                }
150            }
151        }
152
153        Ok(samples)
154    }
155}
156
157pub trait SampleFromBytes {
158    fn from_i16(v: i16) -> Self;
159}
160
161impl SampleFromBytes for i16 {
162    fn from_i16(v: i16) -> Self {
163        v
164    }
165}
166
167pub struct WavWriter<W> {
168    writer: W,
169    _header: WavHeader,
170    data_len_pos: u64,
171    bytes_written: u32,
172}
173
174impl<W: Write + Seek> WavWriter<W> {
175    pub fn new(mut writer: W, channels: Channels, sample_rate: SampleRate) -> io::Result<Self> {
176        writer.write_all(b"RIFF")?;
177        writer.write_all(&[0; 4])?; // Total len placeholder
178        writer.write_all(b"WAVE")?;
179
180        writer.write_all(b"fmt ")?;
181        writer.write_all(&16u32.to_le_bytes())?; // Chunk size 16 for PCM
182        writer.write_all(&1u16.to_le_bytes())?; // PCM
183        writer.write_all(&(channels.count()).to_le_bytes())?;
184        writer.write_all(&sample_rate.value().to_le_bytes())?;
185
186        let bits_per_sample = BitDepth::SixteenBit;
187        let block_align = channels.count() as u16 * bits_per_sample.bits() / 8;
188        let byte_rate = sample_rate.value() * u32::from(block_align);
189
190        writer.write_all(&byte_rate.to_le_bytes())?;
191        writer.write_all(&block_align.to_le_bytes())?;
192        writer.write_all(&bits_per_sample.bits().to_le_bytes())?;
193
194        writer.write_all(b"data")?;
195        let data_len_pos = writer.stream_position()?;
196        writer.write_all(&0u32.to_le_bytes())?; // Data len placeholder
197
198        Ok(Self {
199            writer,
200            _header: WavHeader {
201                channels,
202                sample_rate,
203                bits_per_sample,
204                data_len: 0,
205            },
206            data_len_pos,
207            bytes_written: 0,
208        })
209    }
210
211    pub fn write_samples(&mut self, samples: &[i16]) -> io::Result<()> {
212        for s in samples {
213            self.writer.write_all(&s.to_le_bytes())?;
214        }
215        self.bytes_written += (samples.len() * 2) as u32;
216        Ok(())
217    }
218
219    pub fn finalize(&mut self) -> io::Result<()> {
220        let total_file_len = 4 + (8 + 16) + (8 + self.bytes_written); // WAVE + fmt + data
221
222        self.writer.seek(SeekFrom::Start(4))?;
223        self.writer.write_all(&total_file_len.to_le_bytes())?;
224
225        self.writer.seek(SeekFrom::Start(self.data_len_pos))?;
226        self.writer.write_all(&self.bytes_written.to_le_bytes())?;
227
228        self.writer.seek(SeekFrom::End(0))?;
229        Ok(())
230    }
231}