netris-nestri/packages/server/src/enc_helper.rs

use crate::args::encoding_args::RateControl;
use crate::gpu::{self, GPUInfo, get_gpu_by_card_path, get_gpus_by_vendor};
use gst::prelude::*;

#[derive(Debug, Eq, PartialEq, Clone)]
pub enum VideoCodec {
    H264,
    H265,
    AV1,
    UNKNOWN,
}

impl VideoCodec {
    pub fn to_str(&self) -> &'static str {
        match self {
            Self::H264 => "H.264",
            Self::H265 => "H.265",
            Self::AV1 => "AV1",
            Self::UNKNOWN => "Unknown",
        }
    }

    pub fn from_str(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "h264" | "h.264" | "avc" => Self::H264,
            "h265" | "h.265" | "hevc" | "hev1" => Self::H265,
            "av1" => Self::AV1,
            _ => Self::UNKNOWN,
        }
    }
}

#[derive(Debug, Eq, PartialEq, Clone)]
pub enum EncoderAPI {
    QSV,
    VAAPI,
    NVENC,
    AMF,
    SOFTWARE,
    UNKNOWN,
}

impl EncoderAPI {
    pub fn to_str(&self) -> &'static str {
        match self {
            Self::QSV => "Intel QuickSync Video",
            Self::VAAPI => "Video Acceleration API",
            Self::NVENC => "NVIDIA NVENC",
            Self::AMF => "AMD Media Framework",
            Self::SOFTWARE => "Software",
            Self::UNKNOWN => "Unknown",
        }
    }
}

#[derive(Debug, Eq, PartialEq, Clone)]
pub enum EncoderType {
    SOFTWARE,
    HARDWARE,
    UNKNOWN,
}

impl EncoderType {
    pub fn to_str(&self) -> &'static str {
        match self {
            Self::SOFTWARE => "Software",
            Self::HARDWARE => "Hardware",
            Self::UNKNOWN => "Unknown",
        }
    }

    pub fn from_str(s: &str) -> Self {
        match s.to_lowercase().as_str() {
            "software" => Self::SOFTWARE,
            "hardware" => Self::HARDWARE,
            _ => Self::UNKNOWN,
        }
    }
}

#[derive(Debug, Clone)]
pub struct VideoEncoderInfo {
    pub name: String,
    pub codec: VideoCodec,
    pub encoder_type: EncoderType,
    pub encoder_api: EncoderAPI,
    pub parameters: Vec<(String, String)>,
    pub gpu_info: Option<GPUInfo>,
}

impl VideoEncoderInfo {
    pub fn new(
        name: String,
        codec: VideoCodec,
        encoder_type: EncoderType,
        encoder_api: EncoderAPI,
    ) -> Self {
        Self {
            name,
            codec,
            encoder_type,
            encoder_api,
            parameters: Vec::new(),
            gpu_info: None,
        }
    }

    pub fn get_parameters_string(&self) -> String {
        self.parameters
            .iter()
            .map(|(k, v)| format!("{}={}", k, v))
            .collect::<Vec<_>>()
            .join(" ")
    }

    pub fn set_parameter(&mut self, key: &str, value: &str) {
        self.parameters.push((key.into(), value.into()));
    }

    pub fn apply_parameters(&self, element: &gst::Element, verbose: bool) {
        for (key, value) in &self.parameters {
            if element.has_property(key) {
                if verbose {
                    println!("Setting property {} to {}", key, value);
                }
                element.set_property_from_str(key, value);
            }
        }
    }
}

fn get_encoder_api(encoder: &str, encoder_type: &EncoderType) -> EncoderAPI {
    match encoder_type {
        EncoderType::HARDWARE => {
            if encoder.starts_with("qsv") {
                EncoderAPI::QSV
            } else if encoder.starts_with("va") {
                EncoderAPI::VAAPI
            } else if encoder.starts_with("nv") {
                EncoderAPI::NVENC
            } else if encoder.starts_with("amf") {
                EncoderAPI::AMF
            } else {
                EncoderAPI::UNKNOWN
            }
        }
        EncoderType::SOFTWARE => EncoderAPI::SOFTWARE,
        _ => EncoderAPI::UNKNOWN,
    }
}

fn codec_from_encoder_name(name: &str) -> Option<VideoCodec> {
    if name.contains("h264") {
        Some(VideoCodec::H264)
    } else if name.contains("h265") {
        Some(VideoCodec::H265)
    } else if name.contains("av1") {
        Some(VideoCodec::AV1)
    } else {
        None
    }
}

fn modify_encoder_params<F>(encoder: &VideoEncoderInfo, mut param_check: F) -> VideoEncoderInfo
where
    F: FnMut(&str) -> Option<(String, String)>,
{
    let mut encoder_optz = encoder.clone();
    let element = match gst::ElementFactory::make(&encoder_optz.name).build() {
        Ok(e) => e,
        Err(_) => return encoder_optz, // Return original if element creation fails
    };

    element.list_properties().iter().for_each(|prop| {
        let prop_name = prop.name();
        if let Some((key, value)) = param_check(prop_name) {
            encoder_optz.set_parameter(&key, &value);
        }
    });

    encoder_optz
}

// Parameter setting helpers
pub fn encoder_cqp_params(encoder: &VideoEncoderInfo, quality: u32) -> VideoEncoderInfo {
    modify_encoder_params(encoder, |prop| {
        let pl = prop.to_lowercase();
        if !pl.contains("qp") {
            return None;
        }

        if pl.contains("i") || pl.contains("min") {
            Some((prop.into(), quality.to_string()))
        } else if pl.contains("p") || pl.contains("max") {
            Some((prop.into(), (quality + 2).to_string()))
        } else {
            None
        }
    })
}

pub fn encoder_vbr_params(
    encoder: &VideoEncoderInfo,
    bitrate: u32,
    max_bitrate: u32,
) -> VideoEncoderInfo {
    modify_encoder_params(encoder, |prop| {
        let pl = prop.to_lowercase();
        if !pl.contains("bitrate") {
            return None;
        }

        if !pl.contains("max") {
            Some((prop.into(), bitrate.to_string()))
        } else if encoder.name != "svtav1enc" {
            Some((prop.into(), max_bitrate.to_string()))
        } else {
            None
        }
    })
}

pub fn encoder_cbr_params(encoder: &VideoEncoderInfo, bitrate: u32) -> VideoEncoderInfo {
    modify_encoder_params(encoder, |prop| {
        let pl = prop.to_lowercase();
        if pl.contains("bitrate") && !pl.contains("max") {
            Some((prop.into(), bitrate.to_string()))
        } else {
            None
        }
    })
}

pub fn encoder_gop_params(encoder: &VideoEncoderInfo, gop_size: u32) -> VideoEncoderInfo {
    modify_encoder_params(encoder, |prop| {
        let pl = prop.to_lowercase();
        if pl.contains("gop-size")
            || pl.contains("int-max")
            || pl.contains("max-dist")
            || pl.contains("intra-period-length")
        {
            Some((prop.into(), gop_size.to_string()))
        } else {
            None
        }
    })
}

pub fn encoder_low_latency_params(
    encoder: &VideoEncoderInfo,
    rate_control: &RateControl,
) -> VideoEncoderInfo {
    let mut encoder_optz = encoder_gop_params(encoder, 30);

    match encoder_optz.encoder_api {
        EncoderAPI::QSV => {
            encoder_optz.set_parameter("low-latency", "true");
            encoder_optz.set_parameter("target-usage", "7");
        }
        EncoderAPI::VAAPI => {
            encoder_optz.set_parameter("target-usage", "7");
        }
        EncoderAPI::NVENC => {
            encoder_optz.set_parameter("multi-pass", "disabled");
            encoder_optz.set_parameter("preset", "p1");
            encoder_optz.set_parameter("tune", "ultra-low-latency");
        }
        EncoderAPI::AMF => {
            encoder_optz.set_parameter("preset", "speed");
            let usage = match encoder_optz.codec {
                VideoCodec::H264 | VideoCodec::H265 => "ultra-low-latency",
                VideoCodec::AV1 => "low-latency",
                _ => "",
            };
            if !usage.is_empty() {
                encoder_optz.set_parameter("usage", usage);
            }
        }
        EncoderAPI::SOFTWARE => match encoder_optz.name.as_str() {
            "openh264enc" => {
                encoder_optz.set_parameter("complexity", "low");
                encoder_optz.set_parameter("usage-type", "screen");
            }
            "x264enc" => {
                encoder_optz.set_parameter("rc-lookahead", "0");
                encoder_optz.set_parameter("speed-preset", "ultrafast");
                encoder_optz.set_parameter("tune", "zerolatency");
            }
            "svtav1enc" => {
                encoder_optz.set_parameter("preset", "11");
                encoder_optz.set_parameter("parameters-string", "lookahead=0");
            }
            "av1enc" => {
                encoder_optz.set_parameter("usage-profile", "realtime");
                encoder_optz.set_parameter("cpu-used", "10");
                encoder_optz.set_parameter("lag-in-frames", "0");
            }
            _ => {}
        },
        _ => {}
    }

    encoder_optz
}

pub fn get_compatible_encoders() -> Vec<VideoEncoderInfo> {
    let mut encoders = Vec::new();
    let registry = gst::Registry::get();
    let gpus = gpu::get_gpus();

    for plugin in registry.plugins() {
        for feature in registry.features_by_plugin(plugin.plugin_name().as_str()) {
            let encoder_name = feature.name();

            let factory = match gst::ElementFactory::find(encoder_name.as_str()) {
                Some(f) => f,
                None => continue,
            };

            let klass = match factory.metadata("klass") {
                Some(k) => k.to_lowercase(),
                None => continue,
            };

            if !klass.contains("encoder/video") {
                continue;
            }

            let encoder_type = if klass.contains("/hardware") {
                EncoderType::HARDWARE
            } else {
                EncoderType::SOFTWARE
            };

            let api = get_encoder_api(encoder_name.as_str(), &encoder_type);
            let codec = match codec_from_encoder_name(encoder_name.as_str()) {
                Some(c) => c,
                None => continue,
            };

            let element = match factory.create().build() {
                Ok(e) => e,
                Err(_) => continue,
            };

            let mut gpu_info = None;

            if encoder_type == EncoderType::HARDWARE {
                gpu_info = std::panic::catch_unwind(|| {
                    match api {
                        EncoderAPI::QSV | EncoderAPI::VAAPI => {
                            // Safe property access with panic protection, gstreamer-rs is fun
                            let path = if element.has_property("device-path") {
                                Some(element.property::<String>("device-path"))
                            } else if element.has_property("device") {
                                Some(element.property::<String>("device"))
                            } else {
                                None
                            };

                            path.and_then(|p| get_gpu_by_card_path(&gpus, &p))
                        }
                        EncoderAPI::NVENC if element.has_property("cuda-device-id") => {
                            let cuda_id = element.property::<u32>("cuda-device-id");
                            get_gpus_by_vendor(&gpus, "nvidia")
                                .get(cuda_id as usize)
                                .cloned()
                        }
                        EncoderAPI::AMF if element.has_property("device") => {
                            let device_id = element.property::<u32>("device");
                            get_gpus_by_vendor(&gpus, "amd")
                                .get(device_id as usize)
                                .cloned()
                        }
                        _ => None,
                    }
                })
                .unwrap_or_else(|_| {
                    log::error!(
                        "Panic occurred while querying properties for {}",
                        encoder_name
                    );
                    None
                });
            }

            let mut encoder_info =
                VideoEncoderInfo::new(encoder_name.into(), codec, encoder_type.clone(), api);
            encoder_info.gpu_info = gpu_info;
            encoders.push(encoder_info);
        }
    }

    encoders
}

/// Helper to return encoder from vector by name (case-insensitive).
/// # Arguments
/// * `encoders` - A vector containing information about each encoder.
/// * `name` - A string slice that holds the encoder name.
/// # Returns
/// * `Option<EncoderInfo>` - A reference to an EncoderInfo struct if found.
pub fn get_encoder_by_name(
    encoders: &Vec<VideoEncoderInfo>,
    name: &str,
) -> Option<VideoEncoderInfo> {
    let name = name.to_lowercase();
    encoders
        .iter()
        .find(|encoder| encoder.name.to_lowercase() == name)
        .cloned()
}

/// Helper to get encoders from vector by video codec.
/// # Arguments
/// * `encoders` - A vector containing information about each encoder.
/// * `codec` - The codec of the encoder.
/// # Returns
/// * `Vec<EncoderInfo>` - A vector containing EncoderInfo structs if found.
pub fn get_encoders_by_videocodec(
    encoders: &Vec<VideoEncoderInfo>,
    codec: &VideoCodec,
) -> Vec<VideoEncoderInfo> {
    encoders
        .iter()
        .filter(|encoder| encoder.codec == *codec)
        .cloned()
        .collect()
}

/// Helper to get encoders from vector by encoder type.
/// # Arguments
/// * `encoders` - A vector containing information about each encoder.
/// * `encoder_type` - The type of the encoder.
/// # Returns
/// * `Vec<EncoderInfo>` - A vector containing EncoderInfo structs if found.
pub fn get_encoders_by_type(
    encoders: &Vec<VideoEncoderInfo>,
    encoder_type: &EncoderType,
) -> Vec<VideoEncoderInfo> {
    encoders
        .iter()
        .filter(|encoder| encoder.encoder_type == *encoder_type)
        .cloned()
        .collect()
}

/// Returns best-case compatible encoder given desired codec and encoder type.
/// # Arguments
/// * `encoders` - List of encoders to pick from.
/// * `codec` - Desired codec.
/// * `encoder_type` - Desired encoder type.
/// # Returns
/// * `Option<EncoderInfo>` - Best-case compatible encoder.
pub fn get_best_compatible_encoder(
    encoders: &Vec<VideoEncoderInfo>,
    codec: VideoCodec,
    encoder_type: EncoderType,
) -> Option<VideoEncoderInfo> {
    let mut best_encoder: Option<VideoEncoderInfo> = None;
    let mut best_score: i32 = 0;

    // Filter by codec and type first
    let encoders = get_encoders_by_videocodec(encoders, &codec);
    let encoders = get_encoders_by_type(&encoders, &encoder_type);

    for encoder in encoders {
        // Local score
        let mut score = 0;

        // API score
        score += match encoder.encoder_api {
            EncoderAPI::NVENC => 3,
            EncoderAPI::QSV => 3,
            EncoderAPI::AMF => 3,
            EncoderAPI::VAAPI => 2,
            EncoderAPI::SOFTWARE => 1,
            EncoderAPI::UNKNOWN => 0,
        };

        // If software, score also based on name to get most compatible software encoder for low latency
        if encoder.encoder_type == EncoderType::SOFTWARE {
            score += match encoder.name.as_str() {
                "openh264enc" => 2,
                "x264enc" => 1,
                "svtav1enc" => 2,
                "av1enc" => 1,
                _ => 0,
            };
        }

        // Update best encoder based on score
        if score > best_score {
            best_encoder = Some(encoder.clone());
            best_score = score;
        }
    }

    best_encoder
}