From 6a8042f1ed7ed3b70d9bc04165347b1cf84f0c2b Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Thu, 10 Oct 2024 14:59:56 +0530
Subject: [PATCH] Function to generate WAV header for PCM data

---
 src/calibre/utils/ffmpeg.c | 75 ++++++++++++++++++++++++++++++++++----
 1 file changed, 67 insertions(+), 8 deletions(-)

diff --git a/src/calibre/utils/ffmpeg.c b/src/calibre/utils/ffmpeg.c
index ef3580d9eb..d9258b6b4b 100644
--- a/src/calibre/utils/ffmpeg.c
+++ b/src/calibre/utils/ffmpeg.c
@@ -10,6 +10,7 @@
 #define PY_SSIZE_T_CLEAN
 
 #include <Python.h>
+#include <stdbool.h>
 
 #include <libavutil/audio_fifo.h>
 #include <libavcodec/avcodec.h>
@@ -19,16 +20,19 @@
 #include <libavutil/opt.h>
 #include <libavutil/mem.h>
 #include <libavutil/fifo.h>
+#include <libavformat/avformat.h>
+#include <libavformat/avio.h>
 #include <libswresample/swresample.h>
 
 static PyObject*
-averror_as_python(int errnum) {
-    char buf[4096];
-    av_strerror(errnum, buf, sizeof(buf));
-    PyErr_SetString(PyExc_Exception, buf);
+averror_as_python(int errnum, int line) {
+    char avbuf[4096];
+    av_strerror(errnum, avbuf, sizeof(avbuf));
+    PyErr_Format(PyExc_Exception, "%s:%d:%s", __FILE__, line, avbuf);
     return NULL;
 }
 
+// resample_raw_audio_16bit {{{
 static PyObject*
 resample_raw_audio_16bit(PyObject *self, PyObject *args) {
     int input_sample_rate, output_sample_rate, input_num_channels = 1, output_num_channels = 1;
@@ -47,9 +51,9 @@ resample_raw_audio_16bit(PyObject *self, PyObject *args) {
             &output_layout, fmt, output_sample_rate,
             &input_layout, fmt, input_sample_rate,
             0, NULL);
-    if (ret != 0) { av_free(output); PyBuffer_Release(&inb); return averror_as_python(ret); }
+    if (ret != 0) { av_free(output); PyBuffer_Release(&inb); return averror_as_python(ret, __LINE__); }
 #define free_resources av_free(output); PyBuffer_Release(&inb); swr_free(&swr_ctx);
-    if ((ret = swr_init(swr_ctx)) < 0) { free_resources; return averror_as_python(ret); }
+    if ((ret = swr_init(swr_ctx)) < 0) { free_resources; return averror_as_python(ret, __LINE__); }
     const uint8_t *input = inb.buf;
     Py_BEGIN_ALLOW_THREADS
     ret = swr_convert(swr_ctx,
@@ -57,14 +61,66 @@ resample_raw_audio_16bit(PyObject *self, PyObject *args) {
         &input,  inb.len / (input_num_channels * bytes_per_sample)
     );
     Py_END_ALLOW_THREADS
-    if (ret < 0) { free_resources; return averror_as_python(ret); }
+    if (ret < 0) { free_resources; return averror_as_python(ret, __LINE__); }
     output_size = ret * output_num_channels * bytes_per_sample;
     PyObject *ans = PyBytes_FromStringAndSize((char*)output, output_size);
     free_resources;
 #undef free_resources
     return ans;
-}
+} // }}}
 
+// wav_header_for_pcm_data {{{
+static PyObject*
+wav_header_for_pcm_data(PyObject *self, PyObject *args) {
+    unsigned int sample_rate = 22050, num_channels=1, audio_data_size=0;
+    if (!PyArg_ParseTuple(args, "|III", &audio_data_size, &sample_rate, &num_channels)) return NULL;
+    struct {
+        char riff[4];
+        uint32_t file_size;
+        char wave[4];
+        char fmt[4];
+        uint32_t block_size;
+        uint16_t audio_format;
+        uint16_t num_channels;
+        uint32_t sample_rate;
+        uint32_t byte_rate;
+        uint16_t bytes_per_block;
+        uint16_t bits_per_sample;
+        char data[4];
+        uint32_t subchunk2_size;
+    } wav_header;
+    wav_header.riff[0] = 'R';
+    wav_header.riff[1] = 'I';
+    wav_header.riff[2] = 'F';
+    wav_header.riff[3] = 'F';
+
+    wav_header.wave[0] = 'W';
+    wav_header.wave[1] = 'A';
+    wav_header.wave[2] = 'V';
+    wav_header.wave[3] = 'E';
+
+    wav_header.fmt[0] = 'f';
+    wav_header.fmt[1] = 'm';
+    wav_header.fmt[2] = 't';
+    wav_header.fmt[3] = ' ';
+
+    wav_header.data[0] = 'd';
+    wav_header.data[1] = 'a';
+    wav_header.data[2] = 't';
+    wav_header.data[3] = 'a';
+
+    wav_header.file_size = audio_data_size + sizeof(wav_header) - 8;
+    wav_header.bits_per_sample = 16; // number of bits per sample per channel
+    wav_header.block_size = wav_header.bits_per_sample;
+    wav_header.audio_format = 1; // 1 for PCM 3 for float32
+    wav_header.num_channels = num_channels; // Mono
+    wav_header.sample_rate = sample_rate;
+    wav_header.bytes_per_block = num_channels * wav_header.bits_per_sample / 8;
+    wav_header.byte_rate = sample_rate * wav_header.bytes_per_block;
+    wav_header.subchunk2_size = audio_data_size;
+    return PyBytes_FromStringAndSize((void*)&wav_header, sizeof(wav_header));
+}
+// }}}
 
 // Boilerplate {{{
 static int
@@ -76,6 +132,9 @@ CALIBRE_MODINIT_FUNC PyInit_ffmpeg(void) {
         {"resample_raw_audio_16bit", (PyCFunction)resample_raw_audio_16bit, METH_VARARGS,
         "resample_raw_audio(input_data, input_sample_rate, output_sample_rate, input_num_channels=1, output_num_channels=1) -> Return resampled raw audio data."
         },
+        {"wav_header_for_pcm_data", (PyCFunction)wav_header_for_pcm_data, METH_VARARGS,
+        "wav_header_for_pcm_data(audio_data_size=0, sample_rate=22050, num_channels=1) -> WAV header for specified amount of PCM data as bytestring"
+        },
         {0}  /* Sentinel */
     };
     static struct PyModuleDef module_def = {