kithara-decode 0.0.1-alpha2

#![allow(unsafe_code)]

use std::{ffi::c_void, ptr::NonNull, time::Duration};

use kithara_bufpool::PcmBuf;
use kithara_stream::AudioCodec;

use super::{
    aformat::OwnedFormat,
    ensure_current_thread_attached,
    error::AndroidBackendError,
    ffi::{
        self, KEY_CHANNEL_COUNT, KEY_CSD_0, KEY_MIME, KEY_PCM_ENCODING, KEY_SAMPLE_RATE, MIME_AAC,
        MIME_ALAC, MIME_FLAC, MIME_MP3, MIME_RAW, PCM_ENCODING_16BIT, PCM_ENCODING_FLOAT,
    },
    media_codec::{AndroidPcmEncoding, DequeueOutput, OwnedCodec, QueueInput},
};
use crate::{
    codec::FrameCodec,
    demuxer::TrackInfo,
    error::{DecodeError, DecodeResult},
    types::{DecoderTrackInfo, PcmSpec},
};

struct Consts;

impl Consts {
    const INPUT_DEQUEUE_TIMEOUT_US: i64 = 10_000;
    const OUTPUT_DEQUEUE_TIMEOUT_US: i64 = 10_000;
    const PCM16_SCALE: f32 = 32_768.0;
}

/// Frame-level codec wrapping Android's `AMediaCodec`.
///
/// `MediaCodec` does not surface encoder priming, so this backend has
/// no `decoder_algo_delay` (default 0). Gapless metadata, when needed,
/// comes from the upstream demuxer (`AndroidMediaExtractorDemuxer`
/// parses MP4 `udta`/`iTunSMPB` and stamps it into `TrackInfo.gapless`).
pub(crate) struct AndroidCodec {
    pcm_encoding: AndroidPcmEncoding,
    track_info: DecoderTrackInfo,
    codec: OwnedCodec,
    spec: PcmSpec,
}

impl AndroidCodec {
    /// Build an [`AndroidCodec`] from `TrackInfo`. Captures
    /// `track.gapless` verbatim — `MediaCodec` has no algorithmic
    /// delay of its own, so no per-backend adjustment is applied.
    ///
    /// # Errors
    ///
    /// Returns [`DecodeError::UnsupportedCodec`] for codecs the
    /// `MediaCodec` codec layer doesn't accept; any FFI failure
    /// surfaces as [`DecodeError::Backend`] via [`DecodeError::from`].
    pub(crate) fn open_with_config(track: &TrackInfo) -> DecodeResult<Self> {
        ensure_current_thread_attached().map_err(DecodeError::from)?;

        let mime = match track.codec {
            AudioCodec::AacLc => MIME_AAC,
            AudioCodec::Flac => MIME_FLAC,
            AudioCodec::Pcm => MIME_RAW,
            AudioCodec::Mp3 => MIME_MP3,
            AudioCodec::Alac => MIME_ALAC,
            other => return Err(DecodeError::UnsupportedCodec(other)),
        };

        let format = build_format(mime, track).map_err(DecodeError::from)?;
        let codec = OwnedCodec::create_with_format(mime, &format).map_err(DecodeError::from)?;
        let (spec, pcm_encoding) = read_output_format(&codec).map_err(DecodeError::from)?;

        Ok(Self {
            codec,
            spec,
            pcm_encoding,
            track_info: DecoderTrackInfo {
                gapless: track.gapless,
                ..DecoderTrackInfo::default()
            },
        })
    }

    /// Whether `MediaCodec` accepts this codec at the codec layer alone
    /// (i.e. without an extractor providing per-track metadata).
    ///
    /// Scope: AAC-LC and FLAC over fMP4 (HLS), plus standalone WAV/PCM,
    /// MP3, and ALAC paired with [`super::AndroidMediaExtractorDemuxer`].
    pub(crate) fn supports(codec: AudioCodec) -> bool {
        matches!(
            codec,
            AudioCodec::AacLc
                | AudioCodec::Flac
                | AudioCodec::Pcm
                | AudioCodec::Mp3
                | AudioCodec::Alac
        )
    }
}

impl FrameCodec for AndroidCodec {
    fn decode_frame(
        &mut self,
        frame_data: &[u8],
        pts: Duration,
        _packet_desc: &[u8],
        out: &mut PcmBuf,
    ) -> DecodeResult<u32> {
        if frame_data.is_empty() {
            out.clear();
            return Ok(0);
        }

        match self
            .codec
            .dequeue_input_buffer(Consts::INPUT_DEQUEUE_TIMEOUT_US)
            .map_err(DecodeError::from)?
        {
            Some(mut buf) => {
                let dst = buf.data_mut();
                let copy_len = dst.len().min(frame_data.len());
                dst[..copy_len].copy_from_slice(&frame_data[..copy_len]);
                let pts_us = i64::try_from(pts.as_micros()).unwrap_or(i64::MAX);
                self.codec
                    .queue_input_buffer(QueueInput {
                        index: buf.index,
                        size: copy_len,
                        presentation_time_us: pts_us,
                        flags: 0,
                    })
                    .map_err(DecodeError::from)?;
            }
            None => {
                out.clear();
                return Ok(0);
            }
        }

        match self
            .codec
            .dequeue_output_buffer(Consts::OUTPUT_DEQUEUE_TIMEOUT_US)
            .map_err(DecodeError::from)?
        {
            DequeueOutput::Output(buffer) => {
                let bytes = buffer.data();
                match self.pcm_encoding {
                    AndroidPcmEncoding::Pcm16 => decode_pcm16_into(bytes, out)?,
                    AndroidPcmEncoding::Float => decode_pcm_float_into(bytes, out)?,
                };
                let index = buffer.index;
                self.codec
                    .release_output_buffer(index)
                    .map_err(DecodeError::from)?;
                let channels = self.spec.channels as usize;
                let frames = if channels == 0 {
                    0
                } else {
                    u32::try_from(out.len() / channels).unwrap_or(u32::MAX)
                };
                Ok(frames)
            }
            DequeueOutput::OutputFormatChanged(new_format) => {
                self.spec = new_format.spec;
                self.pcm_encoding = new_format.pcm_encoding;
                out.clear();
                Ok(0)
            }
            DequeueOutput::TryAgainLater => {
                out.clear();
                Ok(0)
            }
        }
    }

    fn flush(&mut self) -> DecodeResult<()> {
        self.codec.flush().map_err(DecodeError::from)
    }

    fn spec(&self) -> PcmSpec {
        self.spec
    }

    fn track_info(&self) -> DecoderTrackInfo {
        self.track_info.clone()
    }
}

fn build_format(
    mime: &std::ffi::CStr,
    track: &TrackInfo,
) -> Result<OwnedFormat, AndroidBackendError> {
    // SAFETY: AMediaFormat_new returns a freshly allocated AMediaFormat
    let raw = NonNull::new(unsafe { ffi::AMediaFormat_new() })
        .ok_or_else(|| AndroidBackendError::operation("media-format-new", "returned null"))?;
    let mut format = OwnedFormat::from(raw);

    // SAFETY: format is live; key/value are static null-terminated CStrs.
    unsafe {
        ffi::AMediaFormat_setString(format.raw(), KEY_MIME.as_ptr(), mime.as_ptr());
    }

    let sample_rate = i32::try_from(track.sample_rate).map_err(|_| {
        AndroidBackendError::operation(
            "media-format-sample-rate",
            format!("rate={} out of range", track.sample_rate),
        )
    })?;
    let channels = i32::from(track.channels);
    if !format.set_i32(KEY_SAMPLE_RATE, sample_rate) {
        return Err(AndroidBackendError::operation(
            "media-format-set-sample-rate",
            format!("rate={sample_rate}"),
        ));
    }
    if !format.set_i32(KEY_CHANNEL_COUNT, channels) {
        return Err(AndroidBackendError::operation(
            "media-format-set-channels",
            format!("channels={channels}"),
        ));
    }
    let _ = format.set_i32(KEY_PCM_ENCODING, PCM_ENCODING_16BIT);

    if !track.extra_data.is_empty() {
        // SAFETY: format is live; extra_data is a readable byte slice.
        unsafe {
            ffi::AMediaFormat_setBuffer(
                format.raw(),
                KEY_CSD_0.as_ptr(),
                track.extra_data.as_ptr() as *const c_void,
                track.extra_data.len(),
            );
        }
    }

    Ok(format)
}

fn read_output_format(
    codec: &OwnedCodec,
) -> Result<(PcmSpec, AndroidPcmEncoding), AndroidBackendError> {
    let format = codec.output_format()?;
    let sample_rate = format.get_u32(KEY_SAMPLE_RATE)?.ok_or_else(|| {
        AndroidBackendError::operation("codec-output-format", "missing sample-rate")
    })?;
    let channels = format.get_u16(KEY_CHANNEL_COUNT)?.ok_or_else(|| {
        AndroidBackendError::operation("codec-output-format", "missing channel-count")
    })?;
    let pcm_encoding = match format.get_i32(KEY_PCM_ENCODING) {
        None | Some(PCM_ENCODING_16BIT) => AndroidPcmEncoding::Pcm16,
        Some(PCM_ENCODING_FLOAT) => AndroidPcmEncoding::Float,
        Some(other) => return Err(AndroidBackendError::UnsupportedPcmEncoding { encoding: other }),
    };
    Ok((
        PcmSpec {
            sample_rate,
            channels,
        },
        pcm_encoding,
    ))
}

fn decode_pcm16_into(bytes: &[u8], out: &mut PcmBuf) -> DecodeResult<()> {
    let count = bytes.len() / 2;
    out.ensure_len(count)?;
    for (dst, chunk) in out.iter_mut().zip(bytes.chunks_exact(2)) {
        let s = i16::from_le_bytes([chunk[0], chunk[1]]);
        *dst = f32::from(s) / Consts::PCM16_SCALE;
    }
    out.truncate(count);
    Ok(())
}

fn decode_pcm_float_into(bytes: &[u8], out: &mut PcmBuf) -> DecodeResult<()> {
    let count = bytes.len() / 4;
    out.ensure_len(count)?;
    for (dst, chunk) in out.iter_mut().zip(bytes.chunks_exact(4)) {
        *dst = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
    }
    out.truncate(count);
    Ok(())
}