moq-mux 0.6.0

use crate::catalog::hang::CatalogExt;
use crate::codec::annexb::{NalIterator, START_CODE};
use crate::container::jitter::Jitter;

use anyhow::Context;
use bytes::{Buf, Bytes, BytesMut};
use scuffle_h265::{NALUnitType, SpsNALUnit};

/// A decoder for H.265 with inline SPS/PPS.
/// Only supports single layer streams (VPS is cached but not parsed).
pub struct Import<E: CatalogExt = ()> {
	// Where new media tracks come from.
	tracks: crate::track_provider::TrackProvider,

	// The catalog being produced.
	catalog: crate::catalog::Producer<E>,

	// The track being produced.
	track: Option<crate::container::Producer<crate::catalog::hang::Container>>,

	// Whether the track has been initialized.
	// If it changes, then we'll reinitialize with a new track.
	config: Option<hang::catalog::VideoConfig>,

	// The current frame being built.
	current: Frame,

	// Used to compute wall clock timestamps if needed.
	zero: Option<tokio::time::Instant>,

	// Retained parameter set NALs from the latest keyframe that carried them,
	// re-injected before bare keyframes. A keyframe may define several of each
	// (slices reference them by id); all are kept, but a new GOP's set supersedes
	// them (replace, not accumulate) so a mid-stream reinit drops stale entries.
	vps: Vec<Bytes>,
	sps: Vec<Bytes>,
	pps: Vec<Bytes>,

	// Tracks the minimum frame duration and updates the catalog `jitter` field.
	jitter: Jitter,
}

impl<E: CatalogExt> Import<E> {
	pub fn new(broadcast: moq_net::BroadcastProducer, catalog: crate::catalog::Producer<E>) -> Self {
		Self {
			tracks: crate::track_provider::TrackProvider::unique(broadcast, ".hev1"),
			catalog,
			track: None,
			config: None,
			current: Default::default(),
			zero: None,
			vps: Vec::new(),
			sps: Vec::new(),
			pps: Vec::new(),
			jitter: Jitter::new(),
		}
	}

	pub fn new_with_track(track: moq_net::TrackProducer, catalog: crate::catalog::Producer<E>) -> Self {
		Self {
			tracks: crate::track_provider::TrackProvider::fixed(track),
			catalog,
			track: None,
			config: None,
			current: Default::default(),
			zero: None,
			vps: Vec::new(),
			sps: Vec::new(),
			pps: Vec::new(),
			jitter: Jitter::new(),
		}
	}

	/// Publish (or republish) the catalog rendition for this SPS. Returns true if
	/// it changed an existing config (a reconfiguration), so the caller drops the
	/// parameter sets tied to the old config. The first SPS is not a reconfiguration.
	fn init(&mut self, sps: &SpsNALUnit) -> anyhow::Result<bool> {
		let profile = &sps.rbsp.profile_tier_level.general_profile;
		let vui_data = sps.rbsp.vui_parameters.as_ref().map(VuiData::new).unwrap_or_default();

		let mut config = hang::catalog::VideoConfig::new(hang::catalog::H265 {
			in_band: true, // We only support `hev1` with inline SPS/PPS for now
			profile_space: profile.profile_space,
			profile_idc: profile.profile_idc,
			profile_compatibility_flags: profile.profile_compatibility_flag.bits().to_be_bytes(),
			tier_flag: profile.tier_flag,
			level_idc: profile.level_idc.context("missing level_idc in SPS")?,
			constraint_flags: crate::codec::h265::pack_constraint_flags(profile),
		});
		config.coded_width = Some(sps.rbsp.cropped_width() as u32);
		config.coded_height = Some(sps.rbsp.cropped_height() as u32);
		config.framerate = vui_data.framerate;
		config.display_ratio_width = vui_data.display_ratio_width;
		config.display_ratio_height = vui_data.display_ratio_height;
		config.container = hang::catalog::Container::Legacy;

		if let Some(old) = &self.config
			&& old == &config
		{
			return Ok(false);
		}

		let reconfigured = self.config.is_some();
		// Seed jitter from whatever has accumulated: a dirty start feeds frames before this
		// first rendition exists, so those per-frame updates would otherwise be lost. The
		// cached `config` stays jitter-free so a later jitter change is not mistaken for a
		// codec reconfiguration.
		let jitter = self.jitter.current();
		let mut catalog = self.catalog.lock();

		if self.track.is_some() && self.tracks.is_fixed() {
			anyhow::bail!("fixed track cannot be reconfigured");
		}

		if let Some(track) = self.track.take() {
			tracing::debug!(name = ?track.name, "reinitializing track");
			catalog.video.renditions.remove(&track.name);
		}

		let track = self.tracks.create()?;
		tracing::debug!(name = ?track.name, ?config, "starting track");
		let mut published = config.clone();
		published.jitter = jitter;
		catalog.video.renditions.insert(track.name.clone(), published);

		self.config = Some(config);
		self.track =
			Some(crate::container::Producer::new(track, crate::catalog::hang::Container::Legacy).with_lenient_start());

		Ok(reconfigured)
	}

	/// Initialize the decoder with SPS/PPS and other non-slice NALs.
	pub fn initialize<T: Buf + AsRef<[u8]>>(&mut self, buf: &mut T) -> anyhow::Result<()> {
		let mut nals = NalIterator::new(buf);

		while let Some(nal) = nals.next().transpose()? {
			self.decode_nal(nal, None)?;
		}

		if let Some(nal) = nals.flush()? {
			self.decode_nal(nal, None)?;
		}

		Ok(())
	}

	/// Returns a reference to the underlying track producer.
	pub fn track(&self) -> anyhow::Result<&moq_net::TrackProducer> {
		Ok(self.track.as_ref().context("not initialized")?.track())
	}

	/// Decode as much data as possible from the given buffer.
	///
	/// Unlike [Self::decode_frame], this method needs the start code for the next frame.
	/// This means it works for streaming media (ex. stdin) but adds a frame of latency.
	///
	/// TODO: This currently associates PTS with the *previous* frame, as part of `maybe_start_frame`.
	pub fn decode_stream<T: Buf + AsRef<[u8]>>(
		&mut self,
		buf: &mut T,
		pts: Option<crate::container::Timestamp>,
	) -> anyhow::Result<()> {
		let pts = self.pts(pts)?;

		// Iterate over the NAL units in the buffer based on start codes.
		let nals = NalIterator::new(buf);

		for nal in nals {
			self.decode_nal(nal?, Some(pts))?;
		}

		Ok(())
	}

	/// Decode all data in the buffer, assuming the buffer contains (the rest of) a frame.
	///
	/// Unlike [Self::decode_stream], this is called when we know NAL boundaries.
	/// This can avoid a frame of latency just waiting for the next frame's start code.
	/// This can also be used when EOF is detected to flush the final frame.
	///
	/// NOTE: The next decode will fail if it doesn't begin with a start code.
	/// Record a frame's reorder delay (`PTS - DTS`) so the catalog `jitter` reflects the
	/// B-frame reorder depth (the decode buffer a transmuxer/player must hold). The container
	/// supplies this since the elementary stream alone carries no decode time. No-op until the
	/// track exists.
	pub fn observe_reorder(&mut self, reorder: crate::container::Timestamp) {
		let Some(jitter) = self.jitter.observe_reorder(reorder) else {
			return;
		};
		let Some(track) = self.track.as_ref() else {
			return;
		};
		if let Some(c) = self.catalog.lock().video.renditions.get_mut(&track.name) {
			c.jitter = Some(jitter);
		}
	}

	pub fn decode_frame<T: Buf + AsRef<[u8]>>(
		&mut self,
		buf: &mut T,
		pts: Option<crate::container::Timestamp>,
	) -> anyhow::Result<()> {
		let pts = self.pts(pts)?;
		// Iterate over the NAL units in the buffer based on start codes.
		let mut nals = NalIterator::new(buf);

		// Iterate over each NAL that is followed by a start code.
		while let Some(nal) = nals.next().transpose()? {
			self.decode_nal(nal, Some(pts))?;
		}

		// Assume the rest of the buffer is a single NAL.
		if let Some(nal) = nals.flush()? {
			self.decode_nal(nal, Some(pts))?;
		}

		// Flush the frame if we read a slice.
		self.maybe_start_frame(Some(pts))?;

		Ok(())
	}

	/// Decode a single NAL unit. Only reads the first header byte to extract nal_unit_type,
	/// Ignores nuh_layer_id and nuh_temporal_id_plus1.
	fn decode_nal(&mut self, nal: Bytes, pts: Option<crate::container::Timestamp>) -> anyhow::Result<()> {
		anyhow::ensure!(nal.len() >= 2, "NAL unit is too short");
		// u16 header: [forbidden_zero_bit(1) | nal_unit_type(6) | nuh_layer_id(6) | nuh_temporal_id_plus1(3)]
		let header = nal.first().context("NAL unit is too short")?;

		let forbidden_zero_bit = (header >> 7) & 1;
		anyhow::ensure!(forbidden_zero_bit == 0, "forbidden zero bit is not zero");

		// Bits 1-6: nal_unit_type
		let nal_unit_type = (header >> 1) & 0b111111;
		let nal_type = NALUnitType::from(nal_unit_type);

		match nal_type {
			NALUnitType::VpsNut => {
				self.maybe_start_frame(pts)?;

				// Track only what this AU carries; the retained set is reconciled at
				// the keyframe so a new GOP's set replaces (not accumulates onto) it.
				crate::codec::annexb::push_distinct(&mut self.current.vps_seen, &nal);
			}
			NALUnitType::SpsNut => {
				self.maybe_start_frame(pts)?;

				// Try to reinitialize the track if the SPS has changed.
				let sps = SpsNALUnit::parse(&mut &nal[..]).context("failed to parse SPS NAL unit")?;
				let reconfigured = self.init(&sps)?;

				// A changed config means the retained VPS/SPS/PPS no longer apply; they
				// may already have been appended to current.chunks earlier in this AU,
				// so reset the sets and AU so only the new parameter sets emit.
				if reconfigured {
					self.vps.clear();
					self.sps.clear();
					self.pps.clear();
					self.current.chunks.clear();
					// Keep vps_seen: in H.265 the VPS precedes the reconfiguring SPS, so
					// any VPS already seen this AU belongs to the new config. Re-append
					// it to the cleared chunks so the keyframe still carries it.
					for nal in &self.current.vps_seen {
						self.current.chunks.extend_from_slice(&START_CODE);
						self.current.chunks.extend_from_slice(nal);
					}
					self.current.sps_seen.clear();
					self.current.pps_seen.clear();
				}

				crate::codec::annexb::push_distinct(&mut self.current.sps_seen, &nal);
			}
			NALUnitType::PpsNut => {
				self.maybe_start_frame(pts)?;

				crate::codec::annexb::push_distinct(&mut self.current.pps_seen, &nal);
			}
			NALUnitType::AudNut | NALUnitType::PrefixSeiNut | NALUnitType::SuffixSeiNut => {
				self.maybe_start_frame(pts)?;
			}
			// Keyframe containing slices
			NALUnitType::IdrWRadl
			| NALUnitType::IdrNLp
			| NALUnitType::BlaNLp
			| NALUnitType::BlaWRadl
			| NALUnitType::BlaWLp
			| NALUnitType::CraNut => {
				// Adopt this keyframe's inline set (dropping any the new GOP no longer
				// uses), or re-inject the retained set if the keyframe carried none.
				crate::codec::annexb::reconcile_keyframe_params(
					&mut self.current.chunks,
					&mut self.vps,
					&mut self.current.vps_seen,
				);
				crate::codec::annexb::reconcile_keyframe_params(
					&mut self.current.chunks,
					&mut self.sps,
					&mut self.current.sps_seen,
				);
				crate::codec::annexb::reconcile_keyframe_params(
					&mut self.current.chunks,
					&mut self.pps,
					&mut self.current.pps_seen,
				);

				self.current.contains_idr = true;
				self.current.contains_slice = true;
			}
			// All other slice types (both N and R variants)
			NALUnitType::TrailN
			| NALUnitType::TrailR
			| NALUnitType::TsaN
			| NALUnitType::TsaR
			| NALUnitType::StsaN
			| NALUnitType::StsaR
			| NALUnitType::RadlN
			| NALUnitType::RadlR
			| NALUnitType::RaslN
			| NALUnitType::RaslR => {
				// Check first_slice_segment_in_pic_flag (bit 7 of third byte, after 2-byte header)
				if nal.get(2).context("NAL unit is too short")? & 0x80 != 0 {
					self.maybe_start_frame(pts)?;
				}
				self.current.contains_slice = true;
			}
			_ => {}
		}

		// Replace the original start code with a canonical 4-byte start code (marginally easier
		// for downstream players, e.g. MSE).
		self.current.chunks.extend_from_slice(&START_CODE);
		self.current.chunks.extend_from_slice(&nal);

		Ok(())
	}

	fn maybe_start_frame(&mut self, pts: Option<crate::container::Timestamp>) -> anyhow::Result<()> {
		// If we haven't seen any slices, we shouldn't flush yet.
		if !self.current.contains_slice {
			return Ok(());
		}

		let track = self.track.as_mut().context("expected SPS before any frames")?;
		let pts = pts.context("missing timestamp")?;

		let payload = std::mem::take(&mut self.current.chunks).freeze();

		let frame = crate::container::Frame {
			timestamp: pts,
			payload,
			keyframe: self.current.contains_idr,
		};

		track.write(frame)?;

		if let Some(jitter) = self.jitter.observe(pts)
			&& let Some(c) = self.catalog.lock().video.renditions.get_mut(&track.name)
		{
			c.jitter = Some(jitter);
		}

		self.current.contains_idr = false;
		self.current.contains_slice = false;
		self.current.vps_seen.clear();
		self.current.sps_seen.clear();
		self.current.pps_seen.clear();

		Ok(())
	}

	/// Finish the track, flushing the current group.
	pub fn finish(&mut self) -> anyhow::Result<()> {
		let track = self.track.as_mut().context("not initialized")?;
		track.finish()?;
		Ok(())
	}

	/// Close the current group and open the next one at `sequence`.
	pub fn seek(&mut self, sequence: u64) -> anyhow::Result<()> {
		let track = self.track.as_mut().context("not initialized")?;
		track.seek(sequence)?;
		Ok(())
	}

	pub fn is_initialized(&self) -> bool {
		self.track.is_some()
	}

	fn pts(&mut self, hint: Option<crate::container::Timestamp>) -> anyhow::Result<crate::container::Timestamp> {
		if let Some(pts) = hint {
			return Ok(pts);
		}

		let zero = self.zero.get_or_insert_with(tokio::time::Instant::now);
		Ok(crate::container::Timestamp::from_micros(
			zero.elapsed().as_micros() as u64
		)?)
	}
}

impl<E: CatalogExt> Drop for Import<E> {
	fn drop(&mut self) {
		if let Some(track) = &self.track {
			tracing::debug!(name = ?track.name, "ending track");
			self.catalog.lock().video.renditions.remove(&track.name);
		}
	}
}

#[derive(Default)]
struct Frame {
	chunks: BytesMut,
	contains_idr: bool,
	contains_slice: bool,
	/// VPS/SPS/PPS NALs already inline in this access unit, so re-injection skips them.
	vps_seen: Vec<Bytes>,
	sps_seen: Vec<Bytes>,
	pps_seen: Vec<Bytes>,
}

#[derive(Default)]
struct VuiData {
	framerate: Option<f64>,
	display_ratio_width: Option<u32>,
	display_ratio_height: Option<u32>,
}

impl VuiData {
	fn new(vui: &scuffle_h265::VuiParameters) -> Self {
		// FPS = time_scale / num_units_in_tick
		let framerate = vui
			.vui_timing_info
			.as_ref()
			.map(|t| t.time_scale.get() as f64 / t.num_units_in_tick.get() as f64);

		let (display_ratio_width, display_ratio_height) = match &vui.aspect_ratio_info {
			// Extended SAR has explicit arbitrary values for width and height.
			scuffle_h265::AspectRatioInfo::ExtendedSar { sar_width, sar_height } => {
				(Some(*sar_width as u32), Some(*sar_height as u32))
			}
			// Predefined map to known values.
			scuffle_h265::AspectRatioInfo::Predefined(idc) => aspect_ratio_from_idc(*idc)
				.map(|(w, h)| (Some(w), Some(h)))
				.unwrap_or((None, None)),
		};

		VuiData {
			framerate,
			display_ratio_width,
			display_ratio_height,
		}
	}
}

fn aspect_ratio_from_idc(idc: scuffle_h265::AspectRatioIdc) -> Option<(u32, u32)> {
	match idc {
		scuffle_h265::AspectRatioIdc::Unspecified => None,
		scuffle_h265::AspectRatioIdc::Square => Some((1, 1)),
		scuffle_h265::AspectRatioIdc::Aspect12_11 => Some((12, 11)),
		scuffle_h265::AspectRatioIdc::Aspect10_11 => Some((10, 11)),
		scuffle_h265::AspectRatioIdc::Aspect16_11 => Some((16, 11)),
		scuffle_h265::AspectRatioIdc::Aspect40_33 => Some((40, 33)),
		scuffle_h265::AspectRatioIdc::Aspect24_11 => Some((24, 11)),
		scuffle_h265::AspectRatioIdc::Aspect20_11 => Some((20, 11)),
		scuffle_h265::AspectRatioIdc::Aspect32_11 => Some((32, 11)),
		scuffle_h265::AspectRatioIdc::Aspect80_33 => Some((80, 33)),
		scuffle_h265::AspectRatioIdc::Aspect18_11 => Some((18, 11)),
		scuffle_h265::AspectRatioIdc::Aspect15_11 => Some((15, 11)),
		scuffle_h265::AspectRatioIdc::Aspect64_33 => Some((64, 33)),
		scuffle_h265::AspectRatioIdc::Aspect160_99 => Some((160, 99)),
		scuffle_h265::AspectRatioIdc::Aspect4_3 => Some((4, 3)),
		scuffle_h265::AspectRatioIdc::Aspect3_2 => Some((3, 2)),
		scuffle_h265::AspectRatioIdc::Aspect2_1 => Some((2, 1)),
		scuffle_h265::AspectRatioIdc::ExtendedSar => None,
		_ => None, // Reserved
	}
}