/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2026 metamuffin <metamuffin.org>
*/
use crate::{Config, CONF, LOCAL_VIDEO_TRANSCODING_TASKS};
use anyhow::Result;
use jellycache::{cache, HashKey};
use jellyremuxer::{demuxers::create_demuxer, muxers::write_fragment, ContainerFormat};
use jellystream_types::{StreamFormatInfo, TrackKind};
use log::info;
use std::{
    fmt::Write,
    io::{Cursor, Read, Write as W2},
    process::{Command, Stdio},
    thread::spawn,
};
use winter_matroska::{block::Block, Cluster, Segment, TrackEntry as MatroskaTrackEntry};

// TODO odd video resolutions can cause errors when transcoding to YUV42{0,2}
// TODO with an implementation that cant handle it (SVT-AV1 is such an impl).

pub fn transcode(
    kind: TrackKind,
    input_key: &str,
    output_format: &StreamFormatInfo,
    mut input: Segment,
    next_kf: Option<Block>,
) -> Result<Segment> {
    let command = transcode_command(
        kind,
        &input.tracks.as_ref().unwrap().entries[0],
        output_format,
        &CONF,
    )
    .unwrap();

    let input_duration = input.info.duration;
    let had_next_kf = next_kf.is_some();

    let output = cache(
        &format!(
            "transcode/media-fragment/{input_key}-{}.mkv",
            HashKey(&command)
        ),
        || {
            let _permit = LOCAL_VIDEO_TRANSCODING_TASKS.lock().unwrap();
            info!("encoding with {command:?}");
            let mut args = command.split(" ");
            let mut proc = Command::new(args.next().unwrap())
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .args(args)
                .spawn()?;

            let mut stdin = proc.stdin.take().unwrap();
            let mut stdout = proc.stdout.take().unwrap();

            spawn(move || {
                input.clusters.extend(next_kf.map(|kf| Cluster {
                    simple_blocks: vec![kf],
                    ..Default::default()
                }));
                write_fragment(ContainerFormat::Matroska, &mut stdin, input).unwrap(); // TODO
                stdin.flush().unwrap();
                drop(stdin);
            });

            let mut output = Vec::new();
            stdout.read_to_end(&mut output)?;
            proc.wait().unwrap().exit_ok()?;
            info!("done");
            Ok(output)
        },
    )?;

    let mut demuxer = create_demuxer(ContainerFormat::Matroska, Box::new(Cursor::new(output)));

    let mut info = demuxer.info()?;
    info.duration = input_duration;
    let tracks = demuxer.tracks()?;
    let mut clusters = Vec::new();
    while let Some((_, cluster)) = demuxer.read_cluster()? {
        clusters.push(cluster);
    }
    if had_next_kf {
        if let Some(c) = clusters.last_mut() {
            c.simple_blocks.pop().expect("empty last cluster");
        }
        clusters.retain(|c| !c.simple_blocks.is_empty() || !c.block_groups.is_empty());
    }

    Ok(Segment {
        info,
        tracks,
        clusters,
        ..Default::default()
    })
}

fn transcode_command(
    kind: TrackKind,
    orig_metadata: &MatroskaTrackEntry,
    format: &StreamFormatInfo,
    config: &Config,
) -> Result<String> {
    let br = format.bitrate as u64;
    let w = format.width.unwrap_or(0);
    let h = format.height.unwrap_or(0);
    let mut o = String::new();
    write!(o, "ffmpeg -hide_banner ")?;

    if kind == TrackKind::Video {
        if config.enable_rkmpp {
            write!(o, "-hwaccel rkmpp -hwaccel_output_format drm_prime ")?;
        }
        if config.enable_rkrga {
            write!(o, "-afbc rga ")?;
        }

        write!(o, "-f matroska -i pipe:0 -copyts ")?;

        if config.enable_rkrga {
            write!(o, "-vf scale_rkrga=w={w}:h={h}:format=nv12:afbc=1 ")?;
        } else {
            write!(o, "-vf scale={w}:{h} ")?;
        }

        match format.codec.as_str() {
            "V_MPEG4/ISO/AVC" if config.enable_rkmpp => {
                write!(o, "-c:v h264_rkmpp -profile:v high -b:v {br} ")?
            }
            "V_MPEGH/ISO/HEVC" if config.enable_rkmpp => {
                write!(o, "-c:v h265_rkmpp -profile:v high -b:v {br} ")?
            }
            "A_AV1" if config.use_svtav1 => {
                let p = config.svtav1_preset.unwrap_or(8);
                write!(o, "-c:v libsvtav1 -preset {p} -b:v {br} ")?
            }
            "A_AV1" if config.use_rav1e => {
                let p = config.rav1e_preset.unwrap_or(8);
                write!(o, "-c:v librav1e -speed {p} -b:v {br} ")?
            }
            "V_MPEG4/ISO/AVC" => {
                let p = config.x264_preset.clone().unwrap_or("fast".to_string());
                write!(o, "-c:v libx264 -preset {p} -b:v {br} ")?
            }
            "V_MPEGH/ISO/HEVC" => write!(o, "-c:v libx265 -b:v {br} ")?,
            "V_VP8" => write!(o, "-c:v libvpx -b:v {br} ")?,
            "V_VP9" => write!(o, "-c:v libvpx-vp9 -b:v {br} ")?,
            "V_AV1" => {
                let p = config.aom_preset.unwrap_or(1);
                write!(o, "-c:v libaom -cpu-used {p} -row-mt 1 -b:v {br} ")?
            }
            _ => todo!(),
        };
    } else if kind == TrackKind::Audio {
        write!(o, "-f matroska -i pipe:0 -copyts ")?;
        if format.codec == "A_OPUS" && orig_metadata.audio.as_ref().unwrap().channels > 2 {
            write!(o, "-ac 2 ")?;
        }
        match format.codec.as_str() {
            "A_OPUS" => write!(o, "-c:a libopus -b:a {br} ")?,
            _ => todo!(),
        }
    } else {
        write!(o, "-f matroska -i pipe:0 -copyts ")?;
        todo!()
    }

    write!(o, "-f matroska pipe:1")?;
    Ok(o)
}
