alpha2-native 0.5.2

//! `probe-encoders`' view of the D3D12VA encoders: each rate-control mode opened on a DX12 adapter
//! and held to one encoded frame.
//!
//! Adapters are asked in the host's order, and the first whose D3D12 Video Encode reports the
//! encoder answers for it. A mode that device does not list is refused and never opened.

use std::sync::Mutex;

use alpha2_media::ColorMetadata;
use alpha2_native_gpu::SharedD3D12Context;
use rsmpeg::avcodec::AVCodecContext;
use rsmpeg::avutil::{AVFrame, ra};
use rsmpeg::error::RsmpegError;
use rsmpeg::ffi;

use super::{DeviceModes, Encoder, d3d12va_encode_frames};
use crate::NativeVideoCodec;
use crate::rate_control::RateControlMechanism;

const CODECS: [NativeVideoCodec; 3] = [
    NativeVideoCodec::H264,
    NativeVideoCodec::Hevc,
    NativeVideoCodec::Av1,
];

/// One open of a D3D12VA encoder with one rate-control request, on the adapter its LUID names,
/// possibly in another process via a [`D3D12RateControlRunner`].
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct D3D12RateControlOpen {
    /// `DXGI_ADAPTER_DESC1::AdapterLuid` of the adapter the open runs on.
    pub adapter: u64,
    /// The codec whose D3D12VA encoder opens.
    pub codec: NativeVideoCodec,
    /// The output size.
    pub size: (u32, u32),
    /// `AVCodecContext::bit_rate`; zero for a quantizer mode.
    pub bit_rate: i64,
    /// The private and generic options that carry the mode, in order.
    pub options: Vec<(String, String)>,
}

impl D3D12RateControlOpen {
    /// The open as command-line arguments for [`Self::parse`]: `--adapter <LUID> --codec <codec>
    /// --size <W>x<H> --bit-rate <b>`, then one `--option <key>=<value>` per option.
    #[must_use]
    pub fn arguments(&self) -> Vec<String> {
        let mut arguments = vec![
            "--adapter".to_owned(),
            format!("{:016X}", self.adapter),
            "--codec".to_owned(),
            self.codec.as_str().to_owned(),
            "--size".to_owned(),
            format!("{}x{}", self.size.0, self.size.1),
            "--bit-rate".to_owned(),
            self.bit_rate.to_string(),
        ];
        for (key, value) in &self.options {
            arguments.push("--option".to_owned());
            arguments.push(format!("{key}={value}"));
        }
        arguments
    }

    /// The open [`Self::arguments`] wrote.
    ///
    /// # Errors
    /// Names the argument that is missing or unreadable.
    pub fn parse(arguments: &[String]) -> Result<Self, String> {
        if !arguments.len().is_multiple_of(2) {
            return Err("an argument has no value".to_owned());
        }
        let mut pairs = arguments.chunks(2);
        let mut value = |flag: &str| match pairs.next() {
            Some([given, value]) if given == flag => Ok(value.clone()),
            _ => Err(format!("expected {flag} <value>")),
        };
        let adapter = u64::from_str_radix(&value("--adapter")?, 16)
            .map_err(|_| "--adapter takes a LUID in hexadecimal".to_owned())?;
        let codec = value("--codec")?;
        let codec = CODECS
            .into_iter()
            .find(|known| known.as_str() == codec)
            .ok_or_else(|| format!("D3D12 Video Encode encodes h264, hevc or av1, not {codec}"))?;
        let size = value("--size")?;
        let size = size
            .split_once('x')
            .and_then(|(width, height)| Some((width.parse().ok()?, height.parse().ok()?)))
            .ok_or_else(|| format!("--size takes <width>x<height>, not {size}"))?;
        let bit_rate = value("--bit-rate")?
            .parse()
            .map_err(|_| "--bit-rate takes an integer".to_owned())?;
        let mut options = Vec::new();
        for _ in 0..arguments.len().saturating_sub(8) / 2 {
            let option = value("--option")?;
            let (key, value) = option
                .split_once('=')
                .ok_or_else(|| format!("--option takes <key>=<value>, not {option}"))?;
            options.push((key.to_owned(), value.to_owned()));
        }
        Ok(Self {
            adapter,
            codec,
            size,
            bit_rate,
            options,
        })
    }
}

/// Runs one [`D3D12RateControlOpen`] and returns the packets drained, or why it gave none.
pub type D3D12RateControlRunner = fn(&D3D12RateControlOpen) -> Result<u32, String>;

static RUNNER: std::sync::OnceLock<D3D12RateControlRunner> = std::sync::OnceLock::new();

/// Installs `runner` to run every [`probe_d3d12va_rate_control`] open out of process, since the
/// driver can crash here as in a test encode. The runner bounds its own wait.
///
/// # Errors
/// Returns `runner` when one is already installed; the first one stays.
pub fn set_d3d12va_rate_control_runner(
    runner: D3D12RateControlRunner,
) -> Result<(), D3D12RateControlRunner> {
    RUNNER.set(runner)
}

/// Makes every rate-control open in this test binary answer one packet without encoding,
/// because a driver crash would take the session tests with it; `probe-encoders`' own test runs
/// the real opens in children.
#[cfg(test)]
pub(crate) fn encode_no_rate_control_frames() {
    #[allow(clippy::unnecessary_wraps)]
    fn answered(open: &D3D12RateControlOpen) -> Result<u32, String> {
        assert_ne!(open.adapter, 0, "an adapter another process can find");
        Ok(1)
    }
    let _ = RUNNER.set(answered);
}

struct ProbedAdapter {
    /// Its position as `ALPHA2_GPU_ADAPTER` counts (`DX12 adapter <index>`).
    position: String,
    /// Position and name, for rows that opened.
    label: String,
    /// Lets another process find the adapter again.
    luid: Option<u64>,
    /// The renderer's device on it, shared with `FFmpeg`, or why there is none.
    shared: Result<SharedD3D12Context, String>,
}

/// Every DX12 adapter, opened once per process on first use.
static ADAPTERS: Mutex<Option<Vec<ProbedAdapter>>> = Mutex::new(None);

/// Opens every DX12 adapter in `wgpu`'s preference order and shares each device with `FFmpeg`.
fn probed_adapters() -> Vec<ProbedAdapter> {
    pollster::block_on(alpha2_gpu::GpuContext::adapters_by_preference_on(
        wgpu::Backends::DX12,
    ))
    .into_iter()
    .map(|(index, adapter)| {
        let name = adapter.get_info().name;
        let position = format!("DX12 adapter {index}");
        let luid =
            alpha2_native_gpu::d3d12_adapter_identity(&adapter).map(|identity| identity.luid);
        let shared = pollster::block_on(alpha2_gpu::GpuContext::open_adapter(adapter))
            .map_err(|error| format!("{name}'s device did not open: {error}"))
            .and_then(|gpu| {
                SharedD3D12Context::from_gpu(gpu)
                    .map_err(|error| format!("{name}'s D3D12VA device did not open: {error}"))
            });
        ProbedAdapter {
            label: format!("{position} ({name})"),
            position,
            luid,
            shared,
        }
    })
    .collect()
}

/// The D3D12 Video Encode mode name for a mechanism.
pub(super) const fn device_mode(mechanism: RateControlMechanism) -> &'static str {
    match mechanism {
        RateControlMechanism::ConstantBitrate => "cbr",
        RateControlMechanism::VariableBitrate => "vbr",
        RateControlMechanism::ConstantQuantizer => "disabled",
    }
}

/// The first of `adapters` whose D3D12 Video Encode reports the encoder; it opens `mode` at `size`
/// or refuses with its own reason. Unlisted modes are never opened.
///
/// # Errors
/// The answering device's refusal, or every adapter's answer when none reports the encoder.
pub(super) fn answering_adapter(
    adapters: &[(&str, Result<DeviceModes, String>)],
    mode: &'static str,
    size: (u32, u32),
) -> Result<usize, String> {
    let mut reasons = Vec::new();
    for (at, (label, modes)) in adapters.iter().enumerate() {
        let modes = match modes {
            Ok(modes) => modes,
            Err(reason) => {
                reasons.push(format!("{label}: {reason}"));
                continue;
            }
        };
        if let Some(reason) = modes.refusal(mode) {
            return Err(format!("on {label}: {reason}"));
        }
        if !modes.extent.accepts(size.0, size.1) {
            let coded = modes.extent.coded(size.0, size.1);
            return Err(format!(
                "on {label}: {} codes {}x{} as {}x{}, outside the device's {}x{} to {}x{}",
                modes.encoder.name,
                size.0,
                size.1,
                coded.0,
                coded.1,
                modes.extent.minimum.0,
                modes.extent.minimum.1,
                modes.extent.maximum.0,
                modes.extent.maximum.1
            ));
        }
        return Ok(at);
    }
    Err(if reasons.is_empty() {
        "no DX12 adapter".to_owned()
    } else {
        reasons.join("; ")
    })
}

/// Opens D3D12VA encoder `name` with one rate-control request and encodes one pool frame on the
/// answering adapter, returning the packets and that adapter.
///
/// It passes if it opened, consumed every option and produced a packet. Each open is one frame, at
/// most one per encoder and mechanism, run by the installed [`D3D12RateControlRunner`] if any,
/// because some drivers occasionally crash during encode.
pub(crate) fn probe_d3d12va_rate_control(
    name: &str,
    width: u32,
    height: u32,
    mechanism: RateControlMechanism,
    request: &crate::rate_control::MechanismRequest,
) -> Result<(u32, String), String> {
    let encoder = CODECS
        .into_iter()
        .filter_map(|codec| Encoder::for_codec(codec).ok())
        .find(|encoder| encoder.name == name)
        .ok_or_else(|| format!("{name} is not a D3D12VA encoder"))?;
    let mut adapters = ADAPTERS
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    let adapters = adapters.get_or_insert_with(probed_adapters);
    let modes: Vec<_> = adapters
        .iter()
        .map(|adapter| {
            let modes = adapter
                .shared
                .as_ref()
                .map_err(Clone::clone)
                .and_then(|shared| DeviceModes::of(shared, encoder));
            (adapter.position.as_str(), modes)
        })
        .collect();
    let adapter = &adapters[answering_adapter(&modes, device_mode(mechanism), (width, height))?];
    let shared = adapter.shared.as_ref().map_err(Clone::clone)?;
    let open = D3D12RateControlOpen {
        adapter: adapter.luid.unwrap_or(0),
        codec: encoder.codec,
        size: (width, height),
        bit_rate: request.bit_rate,
        options: request
            .options
            .iter()
            .map(|(key, value)| (key.to_string_lossy().into_owned(), value.clone()))
            .collect(),
    };
    // An adapter DXGI does not describe cannot be found again by another process.
    let packets = match (RUNNER.get(), adapter.luid) {
        (Some(runner), Some(_)) => runner(&open),
        _ => encode_one_frame(shared, &open),
    };
    packets
        .map(|packets| (packets, adapter.label.clone()))
        .map_err(|error| format!("on {}: {error}", adapter.label))
}

/// Runs `open` in this process (what a [`D3D12RateControlRunner`]'s child runs): opens the adapter
/// by LUID, shares its device with `FFmpeg`, opens the encoder and encodes one pool frame.
///
/// # Errors
/// Why the adapter, device or encoder did not open, or the frame did not encode.
pub fn open_d3d12va_rate_control(open: &D3D12RateControlOpen) -> Result<u32, String> {
    let adapter = pollster::block_on(alpha2_gpu::GpuContext::adapters_by_preference_on(
        wgpu::Backends::DX12,
    ))
    .into_iter()
    .map(|(_, adapter)| adapter)
    .find(|adapter| {
        alpha2_native_gpu::d3d12_adapter_identity(adapter)
            .is_some_and(|identity| identity.luid == open.adapter)
    })
    .ok_or_else(|| format!("no DX12 adapter has the LUID {:016X}", open.adapter))?;
    let gpu = pollster::block_on(alpha2_gpu::GpuContext::open_adapter(adapter))
        .map_err(|error| format!("the adapter's device did not open: {error}"))?;
    let shared = SharedD3D12Context::from_gpu(gpu)
        .map_err(|error| format!("the D3D12VA device did not open: {error}"))?;
    encode_one_frame(&shared, open)
}

/// Opens `open`'s encoder on `shared` with exactly its options, at the export profile,
/// and encodes one pool frame.
fn encode_one_frame(
    shared: &SharedD3D12Context,
    open: &D3D12RateControlOpen,
) -> Result<u32, String> {
    let encoder = Encoder::for_codec(open.codec)?;
    let codec = encoder.find()?;
    let coded_width = i32::try_from(open.size.0).map_err(|_| "width is out of range")?;
    let coded_height = i32::try_from(open.size.1).map_err(|_| "height is out of range")?;
    let mut frames = d3d12va_encode_frames(&shared.ffmpeg_device(), coded_width, coded_height)
        .map_err(|error| error.to_string())?;
    let mut context = AVCodecContext::new(&codec);
    context.set_width(coded_width);
    context.set_height(coded_height);
    context.set_pix_fmt(ffi::AV_PIX_FMT_D3D12);
    context.set_time_base(ra(1, 30));
    context.set_framerate(ra(30, 1));
    context.set_bit_rate(open.bit_rate);
    context.set_gop_size(12);
    context.set_max_b_frames(0);
    context.set_hw_frames_ctx(frames.clone());
    let mut options =
        crate::convert::color_options(ColorMetadata::EXPORT).set(c"profile", encoder.profile, 0);
    for (key, value) in &open.options {
        let nul = |_| format!("the option {key} has a NUL byte");
        options = options.set(
            &std::ffi::CString::new(key.as_str()).map_err(nul)?,
            &std::ffi::CString::new(value.as_str()).map_err(nul)?,
            0,
        );
    }
    let unused = alpha2_ffmpeg_bridge::open_codec_context(&mut context, options)
        .map_err(|error| error.to_string())?;
    crate::convert::reject_unused_encode_options(unused).map_err(|error| error.to_string())?;
    let mut frame = AVFrame::new();
    frames
        .get_buffer(&mut frame)
        .map_err(|error| error.to_string())?;
    frame.set_pts(0);
    context
        .send_frame(Some(&frame))
        .map_err(|error| error.to_string())?;
    context
        .send_frame(None)
        .map_err(|error| error.to_string())?;
    let mut packets = 0;
    loop {
        match context.receive_packet() {
            Ok(_) => packets += 1,
            Err(RsmpegError::EncoderDrainError | RsmpegError::EncoderFlushedError) => break,
            Err(error) => return Err(error.to_string()),
        }
    }
    if packets == 0 {
        return Err("opened but produced no packet".to_owned());
    }
    Ok(packets)
}