From 2e63c5e9aa3d1e2cc24180b0e3a91d273184b1cd Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Mon, 16 Oct 2023 13:40:00 +0500
Subject: [PATCH] attempt at av1 nalu reading

---
 src/encode/codecs/audio.rs |  4 ++++
 src/encode/codecs/mod.rs   |  9 ++++++++-
 src/encode/codecs/video.rs | 25 ++++++++++++++++++++-----
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/src/encode/codecs/audio.rs b/src/encode/codecs/audio.rs
index eb0e1b2..9a364de 100644
--- a/src/encode/codecs/audio.rs
+++ b/src/encode/codecs/audio.rs
@@ -15,6 +15,10 @@ impl Encoder for OpusEncoder {
         self.cmd.stdin.take().unwrap()
     }
 
+    fn preprocess_stream(&mut self) -> Result<(), Box<dyn Error + Send + Sync>> {
+        return Ok(());
+    }
+
     fn read_nalu(&mut self) -> Result<util::NALUPacket, Box<dyn Error + Send + Sync>> {
         todo!();
     }
diff --git a/src/encode/codecs/mod.rs b/src/encode/codecs/mod.rs
index 21feb39..5146823 100644
--- a/src/encode/codecs/mod.rs
+++ b/src/encode/codecs/mod.rs
@@ -6,6 +6,7 @@ use std::process::{ChildStdin, ChildStdout};
 use std::sync::mpsc;
 use std::io::Write;
 use std::thread;
+
 use crate::util;
 
 pub trait Encoder {
@@ -13,11 +14,17 @@ pub trait Encoder {
     
     fn read_nalu(&mut self) -> Result<util::NALUPacket, Box<dyn Error + Send + Sync>>;
 
+    fn preprocess_stream(&mut self) -> Result<(), Box<dyn Error + Send + Sync>>;
+
     fn write_nalu_loop(&mut self, c_out: mpsc::Sender<util::NALUPacket>, c_err: mpsc::Sender<Box<dyn Error + Send + Sync>>) {
+        match self.preprocess_stream() {
+            Ok(_) => (),
+            Err(err) => {util::thread_freeze(c_err, err); return}
+        }
         loop {
             let nalu = match self.read_nalu() {
                 Ok(x) => x,
-                Err(err) => {util::thread_freeze(c_err, err); return}
+                Err(err) => {drop(c_out); thread::park(); return}
             };
             match c_out.send(nalu) {
                 Ok(_) => (),
diff --git a/src/encode/codecs/video.rs b/src/encode/codecs/video.rs
index d094086..016f8e4 100644
--- a/src/encode/codecs/video.rs
+++ b/src/encode/codecs/video.rs
@@ -2,25 +2,40 @@ use std::error::Error;
 use std::process::{Child, Command, Stdio, ChildStdin};
 use std::sync::Arc;
 use std::fs::File;
+use std::io::Read;
 
 use crate::util;
 use crate::encode::codecs::Encoder;
 
-pub struct VP9Encoder {
+pub struct AV1Encoder {
     cmd: Child,
 }
 
-impl Encoder for VP9Encoder {
+impl Encoder for AV1Encoder {
     fn stdin(&mut self) -> ChildStdin {
         self.cmd.stdin.take().unwrap()
     }
 
+    fn preprocess_stream(&mut self) -> Result<(), Box<dyn Error + Send + Sync>> {
+        let mut buff = [0u8; 8];
+        self.cmd.stdout.as_mut().unwrap().read_exact(&mut buff)?;
+        let head_len = u16::from_le_bytes(buff[6..].try_into().unwrap()) as usize;
+        let mut buff = vec![0u8; head_len - 8];
+        self.cmd.stdout.as_mut().unwrap().read_exact(buff.as_mut_slice())?;
+        return Ok(());
+    }
+
     fn read_nalu(&mut self) -> Result<util::NALUPacket, Box<dyn Error + Send + Sync>> {
-        todo!();
+        let mut buff = [0u8; 11];
+        self.cmd.stdout.as_mut().unwrap().read_exact(&mut buff)?;
+        let nalu_len = (u32::from_le_bytes(buff[..4].try_into().unwrap()) & 0xffffff) as usize;
+        let mut data = vec![0u8; nalu_len];
+        self.cmd.stdout.as_mut().unwrap().read_exact(data.as_mut_slice())?;
+        return Ok(util::NALUPacket {packet_type: util::NALUPacketType::Video(util::VideoCodec::AV1), packet_data: data});
     }
 }
 
-pub fn new_vp9(metadata: Arc<util::Metadata>) -> Result<VP9Encoder, Box<dyn Error>> {
+pub fn new_av1(metadata: Arc<util::Metadata>) -> Result<AV1Encoder, Box<dyn Error>> {
     let f = File::create("dump.av1")?;
     let width = metadata.video.width.to_string();
     let height = metadata.video.height.to_string();
@@ -36,5 +51,5 @@ pub fn new_vp9(metadata: Arc<util::Metadata>) -> Result<VP9Encoder, Box<dyn Erro
             "-i", "-",
             "-b", "-"
         ]).stdin(Stdio::piped()).stdout(f).spawn()?;
-    return Ok(VP9Encoder {cmd: cmd});
+    return Ok(AV1Encoder {cmd: cmd});
 }