//! vmaf logic
use crate::process::{Chunks, CommandExt, FfmpegOut, cmd_err, exit_ok_stderr};
use anyhow::Context;
use log::{debug, info};
use std::{path::Path, process::Stdio};
use tokio::process::Command;
use tokio_process_stream::{Item, ProcessChunkStream};
use tokio_stream::{Stream, StreamExt};

/// Calculate VMAF score using ffmpeg.
pub fn run(
    reference: &Path,
    distorted: &Path,
    filter_complex: &str,
    fps: Option<f32>,
) -> anyhow::Result<impl Stream<Item = VmafOut> + use<>> {
    info!(
        "vmaf {} vs reference {}",
        distorted.file_name().and_then(|n| n.to_str()).unwrap_or(""),
        reference.file_name().and_then(|n| n.to_str()).unwrap_or(""),
    );

    let mut cmd = Command::new("ffmpeg");
    cmd.kill_on_drop(true)
        .arg2_opt("-r", fps)
        .arg2("-i", distorted)
        .arg2_opt("-r", fps)
        .arg2("-i", reference)
        .arg2("-filter_complex", filter_complex)
        // Workaround unused streams causing ffmpeg memory leaks
        // See https://github.com/alexheretic/ab-av1/issues/189
        .arg("-an")
        .arg("-sn")
        .arg("-dn")
        .arg2("-f", "null")
        .arg("-")
        .stdin(Stdio::null());

    let cmd_str = cmd.to_cmd_str();
    debug!("cmd `{cmd_str}`");
    let mut vmaf = crate::process::child::AddOnDropChunkStream::from(
        ProcessChunkStream::try_from(cmd).context("ffmpeg vmaf")?,
    );

    Ok(async_stream::stream! {
        let mut chunks = Chunks::default();
        let mut parsed_done = false;
        while let Some(next) = vmaf.next().await {
            match next {
                Item::Stderr(chunk) => {
                    if let Some(out) = VmafOut::try_from_chunk(&chunk, &mut chunks) {
                        if matches!(out, VmafOut::Done(_)) {
                            parsed_done = true;
                        }
                        yield out;
                    }
                }
                Item::Stdout(_) => {}
                Item::Done(code) => {
                    if let Err(err) = exit_ok_stderr("ffmpeg vmaf", code, &cmd_str, &chunks) {
                        yield VmafOut::Err(err);
                    }
                }
            }
        }
        if !parsed_done {
            yield VmafOut::Err(cmd_err(
                "could not parse ffmpeg vmaf score",
                &cmd_str,
                &chunks,
            ));
        }
    })
}

#[derive(Debug)]
pub enum VmafOut {
    Progress(FfmpegOut),
    Done(f32),
    Err(anyhow::Error),
}

impl VmafOut {
    fn try_from_chunk(chunk: &[u8], chunks: &mut Chunks) -> Option<Self> {
        const SCORE_PREFIX: &str = "VMAF score: ";

        chunks.push(chunk);

        if let Some(line) = chunks.rfind_line(|l| l.contains(SCORE_PREFIX)) {
            let idx = line.find(SCORE_PREFIX).unwrap();
            return Some(Self::Done(
                line[idx + SCORE_PREFIX.len()..].trim().parse().ok()?,
            ));
        }
        if let Some(progress) = FfmpegOut::try_parse(chunks.last_line()) {
            return Some(Self::Progress(progress));
        }
        None
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn parse_vmaf_score_207() {
        const FFMPEG_OUT: &str = r#"ffmpeg version n7.0.1 Copyright (c) 2000-2024 the FFmpeg developers
  built with gcc 14.1.1 (GCC) 20240522
  configuration: --prefix=/usr --disable-debug --disable-static --disable-stripping --enable-amf --enable-avisynth --enable-cuda-llvm --enable-lto --enable-fontconfig --enable-frei0r --enable-gmp --enable-gpl --enable-ladspa --enable-libaom --enable-libass --enable-libbluray --enable-libbs2b --enable-libdav1d --enable-libdrm --enable-libdvdnav --enable-libdvdread --enable-libfreetype --enable-libfribidi --enable-libgsm --enable-libharfbuzz --enable-libiec61883 --enable-libjack --enable-libjxl --enable-libmodplug --enable-libmp3lame --enable-libopencore_amrnb --enable-libopencore_amrwb --enable-libopenjpeg --enable-libopenmpt --enable-libopus --enable-libplacebo --enable-libpulse --enable-librav1e --enable-librsvg --enable-librubberband --enable-libsnappy --enable-libsoxr --enable-libspeex --enable-libsrt --enable-libssh --enable-libsvtav1 --enable-libtheora --enable-libv4l2 --enable-libvidstab --enable-libvmaf --enable-libvorbis --enable-libvpl --enable-libvpx --enable-libwebp --enable-libx264 --enable-libx265 --enable-libxcb --enable-libxml2 --enable-libxvid --enable-libzimg --enable-mbedtls --enable-nvdec --enable-nvenc --enable-opencl --enable-opengl --enable-shared --enable-vapoursynth --enable-version3 --enable-vulkan
  libavutil      59.  8.100 / 59.  8.100
  libavcodec     61.  3.100 / 61.  3.100
  libavformat    61.  1.100 / 61.  1.100
  libavdevice    61.  1.100 / 61.  1.100
  libavfilter    10.  1.100 / 10.  1.100
  libswscale      8.  1.100 /  8.  1.100
  libswresample   5.  1.100 /  5.  1.100
  libpostproc    58.  1.100 / 58.  1.100

  libavutil      59.  8.100 / 59.  8.100
  libavcodec     61.  3.100 / 61.  3.100
  libavformat    61.  1.100 / 61.  1.100
  libavdevice    61.  1.100 / 61.  1.100
  libavfilter    10.  1.100 / 10.  1.100
  libswscale      8.  1.100 /  8.  1.100
  libswresample   5.  1.100 /  5.  1.100
  libpostproc    58.  1.100 / 58.  1.100
Input #0, mov,mp4,m4a,3gp,3g2,mj2, from 'C:\Users\Administrator\Personal_scripts\Python\PythonScripts\PythonScripts\src\.ab-av1-RM46M2PZOVjb\A11 崩三 黑曼巴之影_1.sample2+600f.av1.crf37.5.mp4':
  Metadata:
    major_brand     : isom
    minor_version   : 512
    compatible_brands: isomav01iso2mp41
    title           : Project 1
    date            : 2019-07-11
    encoder         : Lavf61.1.100
  Duration: 00:00:20.00, start: 0.000000, bitrate: 1562 kb/s
  Stream #0:0[0x1](und): Video: av1 (libdav1d) (Main) (av01 / 0x31307661), yuv420p10le(tv, progressive), 1000x696, 1560 kb/s, SAR 1:1 DAR 125:87, 30 fps, 30 tbr, 15360 tbn (default)
      Metadata:
        handler_name    : VideoHandler
        vendor_id       : [0][0][0][0]
        encoder         : Lavc61.3.100 libsvtav1
Input #1, matroska,webm, from 'C:\Users\Administrator\Personal_scripts\Python\PythonScripts\PythonScripts\src\.ab-av1-RM46M2PZOVjb\A11 崩三 黑曼巴之影_1.sample2+600f.mkv':
  Metadata:
    title           : Project 1
    DATE            : 2019-07-11
    MAJOR_BRAND     : isom
    MINOR_VERSION   : 512
    COMPATIBLE_BRANDS: isomiso2mp41
    ENCODER         : Lavf61.1.100
  Duration: 00:00:20.00, start: 0.000000, bitrate: 6114 kb/s
  Stream #1:0: Video: mpeg4 (Simple Profile), yuv420p, 1000x696 [SAR 1:1 DAR 125:87], 30 fps, 30 tbr, 1k tbn (default)
      Metadata:
        HANDLER_NAME    : VideoHandler
        VENDOR_ID       : [0][0][0][0]
        DURATION        : 00:00:20.000000000
Stream mapping:
  Stream #0:0 (libdav1d) -> format:default
  Stream #1:0 (mpeg4) -> format:default
  libvmaf:default -> Stream #0:0 (wrapped_avframe)
Press [q] to stop, [?] for help
Output #0, null, to 'pipe:':
  Metadata:
    major_brand     : isom
    minor_version   : 512
    compatible_brands: isomav01iso2mp41
    title           : Project 1
    date            : 2019-07-11
    encoder         : Lavf61.1.100
  Stream #0:0: Video: wrapped_avframe, yuv420p10le(tv, progressive), 1552x1080 [SAR 5625:5626 DAR 125:87], q=2-31, 200 kb/s, 24 tbn
      Metadata:
        encoder         : Lavc61.3.100 wrapped_avframe
frame=   48 fps=0.0 q=-0.0 size=N/A time=00:00:01.95 bitrate=N/A speed=3.79x    
frame=  101 fps= 97 q=-0.0 size=N/A time=00:00:04.16 bitrate=N/A speed=   4x    
frame=  156 fps=100 q=-0.0 size=N/A time=00:00:06.45 bitrate=N/A speed=4.14x    
frame=  209 fps=101 q=-0.0 size=N/A time=00:00:08.66 bitrate=N/A speed= 4.2x    
frame=  264 fps=102 q=-0.0 size=N/A time=00:00:10.95 bitrate=N/A speed=4.23x    
frame=  319 fps=103 q=-0.0 size=N/A time=00:00:13.25 bitrate=N/A speed=4.26x    
frame=  373 fps=103 q=-0.0 size=N/A time=00:00:15.50 bitrate=N/A speed=4.27x    
frame=  429 fps=103 q=-0.0 size=N/A time=00:00:17.83 bitrate=N/A speed= 4.3x    
frame=  482 fps=103 q=-0.0 size=N/A time=00:00:20.04 bitrate=N/A speed=4.29x    
frame=  536 fps=104 q=-0.0 size=N/A time=00:00:22.29 bitrate=N/A speed=4.31x    
frame=  589 fps=103 q=-0.0 size=N/A time=00:00:24.50 bitrate=N/A speed= 4.3x    
[Parsed_libvmaf_6 @ 000002b296bac480] VMAF score: 94.826380
[out#0/null @ 000002b2916f8b80] video:258KiB audio:0KiB subtitle:0KiB other streams:0KiB global headers:0KiB muxing overhead: unknown
frame=  600 fps=102 q=-0.0 Lsize=N/A time=00:00:24.95 bitrate=N/A speed=4.24x"#;

        const CHUNK_SIZE: usize = 64;

        let ffmpeg = FFMPEG_OUT.as_bytes();

        let mut chunks = Chunks::default();
        let mut start_idx = 0;
        let mut vmaf_score = None;
        while start_idx < ffmpeg.len() {
            let chunk = &ffmpeg[start_idx..(start_idx + CHUNK_SIZE).min(FFMPEG_OUT.len())];
            println!("* {}", String::from_utf8_lossy(chunk).trim());

            if let Some(vmaf) = VmafOut::try_from_chunk(chunk, &mut chunks) {
                println!("{vmaf:?}");
                if let VmafOut::Done(score) = vmaf {
                    vmaf_score = Some(score);
                }
            }

            start_idx += CHUNK_SIZE;
        }

        assert_eq!(vmaf_score, Some(94.82638), "failed to parse vmaf score");
    }
}
