//! AVC and HEVC through the encoders a device already has ([platform-encode]).
//!
//! WAI writes no AVC or HEVC encoder. This module hands pictures to an encoder that the
//! device's operating system, its GPU driver or a separately obtained runtime library
//! provides, reads back the elementary stream, and checks what came back: an ISO/IEC
//! 14496-10 or ISO/IEC 23008-2 Annex B byte stream with one access unit per picture,
//! an IDR access unit exactly where the caller asked for one, and parameter sets at
//! every IDR ([`annexb`]). It records who made the bytes ([`Identity`]) and signs that
//! statement ([`record`]). Nothing here is a deterministic step: a verifier holds the
//! bytes by their digest and never re-encodes them.
//!
//! The encoder interfaces, by role ([platform-encode] §5):
//!
//! | [`Api`] | what it drives | where |
//! |---|---|---|
//! | `platform` | the operating system's media framework, its hardware media engine or its software encoder | macOS, iOS, tvOS |
//! | `va` | the Video Acceleration API, whose driver drives the GPU's encode engine | Linux |
//! | `gpu-driver` | the encode interface a GPU's driver exports from its own library | Linux |
//! | `software` | a software AVC encoder library loaded at run time from a path the user names | any Unix |
//!
//! The adapter files name each interface's real symbols (`adapter_*.rs`); everything
//! else names the interface by role.
//!
//! [platform-encode]: ../../../extensions/platform-encode.md

pub mod annexb;
pub mod corpus;
pub mod energy;
pub mod record;
pub mod scale;
pub mod y4m;

#[cfg(unix)]
mod dl;

#[cfg(target_vendor = "apple")]
mod adapter_videotoolbox;

#[cfg(target_os = "linux")]
mod adapter_vaapi;

#[cfg(target_os = "linux")]
mod adapter_nvenc;

#[cfg(unix)]
mod adapter_openh264;

/// The media framework's decoder, for conformance and quality measurement (§10.2).
#[cfg(target_vendor = "apple")]
pub use adapter_videotoolbox::decode as platform_decode;

use std::time::Instant;

/// The coded format.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Codec {
    /// ISO/IEC 14496-10, `wai.video.avc`.
    Avc,
    /// ISO/IEC 23008-2, `wai.video.hevc`.
    Hevc,
}

impl Codec {
    /// The capability string the payload is registered under (SPEC §5).
    pub fn capability(self) -> &'static str {
        match self {
            Codec::Avc => "wai.video.avc",
            Codec::Hevc => "wai.video.hevc",
        }
    }

    /// The codec from its capability string.
    pub fn from_capability(c: &str) -> Option<Codec> {
        match c {
            "wai.video.avc" => Some(Codec::Avc),
            "wai.video.hevc" => Some(Codec::Hevc),
            _ => None,
        }
    }

    /// The short name the command line takes: `avc` or `hevc`.
    pub fn name(self) -> &'static str {
        match self {
            Codec::Avc => "avc",
            Codec::Hevc => "hevc",
        }
    }

    /// The file extension of an elementary stream of this codec: `h264` or `h265`.
    pub fn extension(self) -> &'static str {
        match self {
            Codec::Avc => "h264",
            Codec::Hevc => "h265",
        }
    }
}

/// The encoder interface, named by role ([platform-encode] §5).
///
/// [platform-encode]: ../../../extensions/platform-encode.md
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Api {
    Platform,
    Va,
    GpuDriver,
    Software,
}

impl Api {
    pub const ALL: [Api; 4] = [Api::Platform, Api::Va, Api::GpuDriver, Api::Software];

    /// The interface's name in a record and on the command line.
    pub fn name(self) -> &'static str {
        match self {
            Api::Platform => "platform",
            Api::Va => "va",
            Api::GpuDriver => "gpu-driver",
            Api::Software => "software",
        }
    }

    pub fn from_name(s: &str) -> Option<Api> {
        Api::ALL.into_iter().find(|a| a.name() == s)
    }
}

/// What an encode is asked for. Every encoder is held to the same shape of output
/// ([platform-encode] §4): 8-bit 4:2:0, progressive, no reordering (each picture's access
/// unit follows the one before it in presentation order), an IDR access unit exactly at
/// the pictures the caller marks and nowhere else, a target average bit rate.
///
/// [platform-encode]: ../../../extensions/platform-encode.md
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Config {
    pub codec: Codec,
    /// Coded width and height in luma samples, each even and at least 16.
    pub width: u32,
    pub height: u32,
    /// Frame rate `fps_num / fps_den` frames per second.
    pub fps_num: u32,
    pub fps_den: u32,
    /// Target average bit rate, kilobits (1000 bits) per second.
    pub kbps: u32,
    /// Whether the samples use the full range (`true`) or the video range (`false`).
    pub full_range: bool,
}

impl Config {
    /// Check the shape every encoder is held to.
    pub fn check(&self) -> Result<(), EncodeError> {
        let bad = |d: String| Err(EncodeError::new("encode-config", d));
        if self.width < 16 || self.height < 16 || !self.width.is_multiple_of(2) || !self.height.is_multiple_of(2) {
            return bad(format!("{}×{}: each side must be even and at least 16", self.width, self.height));
        }
        if self.width > 8192 || self.height > 8192 {
            return bad(format!("{}×{}: each side at most 8192", self.width, self.height));
        }
        if self.fps_num == 0 || self.fps_den == 0 {
            return bad("a frame rate of 0".into());
        }
        if self.kbps == 0 || self.kbps > 1_000_000 {
            return bad(format!("{} kb/s: from 1 to 1000000", self.kbps));
        }
        Ok(())
    }

    /// Bytes of one 4:2:0 picture.
    pub fn picture_len(&self) -> usize {
        let (w, h) = (self.width as usize, self.height as usize);
        w * h + 2 * (w / 2) * (h / 2)
    }
}

/// One 8-bit 4:2:0 picture: the luma plane then the two chroma planes, each row-major
/// with no padding (stride = width, and width / 2 for chroma).
#[derive(Debug, Clone, Copy)]
pub struct Picture<'a> {
    pub y: &'a [u8],
    pub u: &'a [u8],
    pub v: &'a [u8],
}

impl<'a> Picture<'a> {
    /// The picture held in `i420` (the three planes one after another) for `cfg`.
    pub fn from_i420(cfg: &Config, i420: &'a [u8]) -> Picture<'a> {
        let (w, h) = (cfg.width as usize, cfg.height as usize);
        let c = (w / 2) * (h / 2);
        Picture { y: &i420[..w * h], u: &i420[w * h..w * h + c], v: &i420[w * h + c..w * h + 2 * c] }
    }
}

/// Who made the bytes ([platform-encode] §5). Every string is what the interface itself
/// reports at run time, read as the adapter documents.
///
/// [platform-encode]: ../../../extensions/platform-encode.md
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Identity {
    /// The interface ([`Api::name`]).
    pub api: &'static str,
    /// The encoder, as the interface names it.
    pub encoder: String,
    /// The encoder's version, as the interface or its library reports it.
    pub version: String,
    /// Whether a fixed-function encode engine coded the pictures.
    pub hardware: bool,
    /// The device the encoder ran on.
    pub device: String,
    /// The driver, or for a platform encoder the operating system build, that ran it.
    pub driver: String,
    /// For a runtime-loaded library: the SHA-256 of the file loaded, 64 lowercase hex.
    pub library_sha256: Option<String>,
}

/// Why an encode did not happen or was refused.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EncodeError {
    pub code: &'static str,
    pub detail: String,
}

impl EncodeError {
    pub fn new(code: &'static str, detail: impl Into<String>) -> EncodeError {
        EncodeError { code, detail: detail.into() }
    }
}

impl std::fmt::Display for EncodeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}: {}", self.code, self.detail)
    }
}

impl std::error::Error for EncodeError {}

/// One open encoder. An adapter implements it.
///
/// Contract, which [`encode`] checks on the output and refuses where it is broken:
/// - `push` is called once per picture in presentation order; `idr` is true for the
///   first picture and wherever the caller asks for one;
/// - the bytes returned, concatenated across every `push` and `finish`, are one Annex B
///   byte stream of exactly one access unit per picture, in the order pushed;
/// - an access unit is an IDR exactly when its picture was pushed with `idr`, and every
///   IDR access unit carries the parameter sets (SPS and PPS, and the VPS for HEVC)
///   before its first slice.
pub(crate) trait Session {
    /// Encode one picture; returns the bytes of the access units completed so far.
    fn push(&mut self, pic: &Picture<'_>, idr: bool) -> Result<Vec<u8>, EncodeError>;
    /// Drain the encoder; returns the bytes of every access unit not yet returned.
    fn finish(&mut self) -> Result<Vec<u8>, EncodeError>;
}

/// An adapter's open session and the identity it read.
pub(crate) struct Opened {
    pub identity: Identity,
    pub session: Box<dyn Session>,
}

/// What a caller may say about which encoder to use, beyond its interface.
#[derive(Debug, Clone, Default)]
pub struct Options {
    /// `software`: the path of the library to load. Required for that interface.
    pub library: Option<std::path::PathBuf>,
    /// `va`: the DRM render node to open (default: the first that offers the profile).
    /// `gpu-driver`: the device ordinal, in decimal (default 0).
    pub device: Option<String>,
    /// `platform`: require (`Some(true)`) or refuse (`Some(false)`) the hardware engine.
    pub hardware: Option<bool>,
}

/// Open an encoder of interface `api` for `cfg`.
pub(crate) fn open(api: Api, cfg: &Config, opts: &Options) -> Result<Opened, EncodeError> {
    cfg.check()?;
    let unavailable = || EncodeError::new("encoder-unavailable", format!("the {} interface is not built for this platform", api.name()));
    match api {
        Api::Platform => {
            #[cfg(target_vendor = "apple")]
            {
                adapter_videotoolbox::open(cfg, opts)
            }
            #[cfg(not(target_vendor = "apple"))]
            {
                Err(unavailable())
            }
        }
        Api::Va => {
            #[cfg(target_os = "linux")]
            {
                adapter_vaapi::open(cfg, opts)
            }
            #[cfg(not(target_os = "linux"))]
            {
                Err(unavailable())
            }
        }
        Api::GpuDriver => {
            #[cfg(target_os = "linux")]
            {
                adapter_nvenc::open(cfg, opts)
            }
            #[cfg(not(target_os = "linux"))]
            {
                Err(unavailable())
            }
        }
        Api::Software => {
            #[cfg(unix)]
            {
                adapter_openh264::open(cfg, opts)
            }
            #[cfg(not(unix))]
            {
                Err(unavailable())
            }
        }
    }
}

/// One interface's answer to a probe: the identity of the encoder it would use for a
/// small encode of `codec`, or why it has none.
#[derive(Debug, Clone)]
pub struct Probe {
    pub api: Api,
    pub codec: Codec,
    pub result: Result<Identity, EncodeError>,
}

/// Ask every interface built here whether it can encode each codec, by opening a
/// session of 256×144 at 30 fps and closing it again (some engines refuse a picture
/// narrower than 146 samples). Nothing is encoded.
pub fn probe(opts: &Options) -> Vec<Probe> {
    let mut out = Vec::new();
    for api in Api::ALL {
        for codec in [Codec::Avc, Codec::Hevc] {
            let cfg = Config { codec, width: 256, height: 144, fps_num: 30, fps_den: 1, kbps: 500, full_range: false };
            let result = open(api, &cfg, opts).map(|o| o.identity);
            out.push(Probe { api, codec, result });
        }
    }
    out
}

/// An encode, checked.
#[derive(Debug, Clone)]
pub struct Encoded {
    pub identity: Identity,
    /// The Annex B byte stream.
    pub stream: Vec<u8>,
    /// Each access unit's byte range in `stream`, in decode (= presentation) order.
    pub units: Vec<std::ops::Range<usize>>,
    /// The stream's sequence parameter set, read (ISO/IEC 14496-10 §7.3.2.1.1,
    /// ISO/IEC 23008-2 §7.3.2.2).
    pub sps: annexb::SpsInfo,
    /// Wall time of the encode from the first picture handed in to the last byte back.
    pub wall_us: u64,
}

/// Encode `frames` (each one 4:2:0 picture of `cfg`, its planes one after another) with
/// an IDR at each frame index in `idr_at` (sorted, holding 0), and check the result
/// ([platform-encode] §4.3). Any encoder output that breaks the shape is refused here,
/// with the reason, rather than returned.
///
/// [platform-encode]: ../../../extensions/platform-encode.md
pub fn encode<'f>(api: Api, cfg: &Config, opts: &Options, frames: impl IntoIterator<Item = &'f [u8]>, idr_at: &[u64]) -> Result<Encoded, EncodeError> {
    if idr_at.first() != Some(&0) || idr_at.windows(2).any(|w| w[0] >= w[1]) {
        return Err(EncodeError::new("encode-config", "the IDR plan must start at frame 0 and increase"));
    }
    let Opened { identity, mut session } = open(api, cfg, opts)?;
    let t0 = Instant::now();
    let mut stream = Vec::new();
    let mut n = 0u64;
    let mut next = 0usize;
    for f in frames {
        if f.len() != cfg.picture_len() {
            return Err(EncodeError::new("encode-input", format!("frame {n} holds {} bytes, a picture of {}×{} holds {}", f.len(), cfg.width, cfg.height, cfg.picture_len())));
        }
        let idr = idr_at.get(next) == Some(&n);
        if idr {
            next += 1;
        }
        stream.extend(session.push(&Picture::from_i420(cfg, f), idr)?);
        n += 1;
    }
    if next != idr_at.len() {
        return Err(EncodeError::new("encode-config", format!("the IDR plan names frame {} of {n}", idr_at[next])));
    }
    stream.extend(session.finish()?);
    let wall_us = t0.elapsed().as_micros() as u64;
    drop(session);
    let checked = annexb::check_stream(cfg.codec, &stream, n, idr_at).map_err(|r| EncodeError::new(r.code, format!("the {} encoder's output: {}", api.name(), r.detail)))?;
    if (checked.sps.width, checked.sps.height) != (cfg.width, cfg.height) {
        return Err(EncodeError::new("encode-sps-mismatch", format!("the stream codes {}×{}, the encode asked for {}×{}", checked.sps.width, checked.sps.height, cfg.width, cfg.height)));
    }
    Ok(Encoded { identity, stream, units: checked.units, sps: checked.sps, wall_us })
}

/// A source read and made ready for one rung: the frames the encoder is handed, the
/// encode's settings, and what the record says of the source (§4.1, §4.2).
pub struct Prepared {
    pub cfg: Config,
    pub source: record::SourceInfo,
    pub frames: Vec<Vec<u8>>,
}

/// Read a YUV4MPEG2 source (8-bit 4:2:0, each side even) and make the frames of a rung of
/// `size` (the source's size when `None`) at `kbps`, resampled by [`scale`] where the
/// sizes differ.
pub fn prepare(y4m_bytes: &[u8], codec: Codec, size: Option<(u32, u32)>, kbps: u32) -> Result<Prepared, EncodeError> {
    use sha2::Digest as _;
    let (hd, frames) = y4m::read_all(y4m_bytes)?;
    if frames.is_empty() {
        return Err(EncodeError::new("encode-input", "a source of no frames"));
    }
    if !hd.width.is_multiple_of(2) || !hd.height.is_multiple_of(2) {
        return Err(EncodeError::new("encode-input", format!("a source of {}×{}: each side must be even", hd.width, hd.height)));
    }
    let mut h = sha2::Sha256::new();
    for f in &frames {
        h.update(f);
    }
    let source = record::SourceInfo { frames: frames.len() as u64, width: hd.width, height: hd.height, sha256: hex(&h.finalize()) };
    let (w, ht) = size.unwrap_or((hd.width, hd.height));
    let cfg = Config { codec, width: w, height: ht, fps_num: hd.fps_num, fps_den: hd.fps_den, kbps, full_range: hd.full_range };
    cfg.check()?;
    let frames = if (w, ht) == (hd.width, hd.height) {
        frames
    } else {
        frames.iter().map(|f| scale::i420(f, hd.width as usize, hd.height as usize, w as usize, ht as usize)).collect()
    };
    Ok(Prepared { cfg, source, frames })
}

/// Read an IDR plan: one frame index in decimal per line, each line ended by a line
/// feed (packaging §7.3's keyframe plan has this form).
pub fn read_plan(text: &str) -> Result<Vec<u64>, EncodeError> {
    let bad = |d: String| EncodeError::new("encode-config", d);
    if !text.ends_with('\n') {
        return Err(bad("an IDR plan whose last line has no line feed".into()));
    }
    text.lines().map(|l| if l.bytes().all(|c| c.is_ascii_digit()) && !l.is_empty() && (l == "0" || !l.starts_with('0')) { l.parse::<u64>().map_err(|e| bad(e.to_string())) } else { Err(bad(format!("an IDR plan line {l:?}"))) }).collect()
}

/// The IDR plan of a fixed key-frame interval: 0, `every`, 2 · `every`, … below `frames`.
pub fn idr_every(frames: u64, every: u64) -> Vec<u64> {
    let every = every.max(1);
    (0..frames).step_by(every as usize).collect()
}

/// The host's processor, as its operating system names it.
pub fn cpu_name() -> String {
    #[cfg(target_os = "linux")]
    {
        std::fs::read_to_string("/proc/cpuinfo")
            .ok()
            .and_then(|t| t.lines().find(|l| l.starts_with("model name")).and_then(|l| l.split_once(':')).map(|(_, v)| v.trim().to_string()))
            .unwrap_or_else(|| "unknown processor".into())
    }
    #[cfg(target_vendor = "apple")]
    {
        sysctl(c"machdep.cpu.brand_string")
    }
    #[cfg(not(any(target_os = "linux", target_vendor = "apple")))]
    {
        "unknown processor".into()
    }
}

/// The operating system's release, as `os <release>`.
pub fn os_release() -> String {
    #[cfg(target_os = "linux")]
    {
        format!("os {}", std::fs::read_to_string("/proc/sys/kernel/osrelease").map(|s| s.trim().to_string()).unwrap_or_default())
    }
    #[cfg(target_vendor = "apple")]
    {
        format!("os {} build {}", sysctl(c"kern.osproductversion"), sysctl(c"kern.osversion"))
    }
    #[cfg(not(any(target_os = "linux", target_vendor = "apple")))]
    {
        "os unknown".into()
    }
}

#[cfg(target_vendor = "apple")]
fn sysctl(name: &std::ffi::CStr) -> String {
    unsafe extern "C" {
        fn sysctlbyname(name: *const std::ffi::c_char, old: *mut std::ffi::c_void, oldlen: *mut usize, new: *const std::ffi::c_void, newlen: usize) -> std::ffi::c_int;
    }
    let mut len = 0usize;
    // SAFETY: a size query, then a read into a buffer of that size, with a NUL-terminated name.
    unsafe {
        if sysctlbyname(name.as_ptr(), std::ptr::null_mut(), &mut len, std::ptr::null(), 0) != 0 || len == 0 {
            return String::new();
        }
        let mut buf = vec![0u8; len];
        if sysctlbyname(name.as_ptr(), buf.as_mut_ptr().cast(), &mut len, std::ptr::null(), 0) != 0 {
            return String::new();
        }
        buf.truncate(len);
        std::ffi::CStr::from_bytes_until_nul(&buf).map(|c| c.to_string_lossy().into_owned()).unwrap_or_default()
    }
}

/// Lowercase hexadecimal of `b`.
pub fn hex(b: &[u8]) -> String {
    use std::fmt::Write as _;
    b.iter().fold(String::with_capacity(2 * b.len()), |mut s, x| {
        let _ = write!(s, "{x:02x}");
        s
    })
}

/// SHA-256 of `b`, lowercase hex.
pub fn sha256_hex(b: &[u8]) -> String {
    use sha2::Digest as _;
    hex(&sha2::Sha256::digest(b))
}
