/*
    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, LOCAL_VIDEO_TRANSCODING_TASKS};
use anyhow::Result;
use jellycache::{Cache, HashKey};
use jellyremuxer::{ContainerFormat, demuxers::create_demuxer, muxers::write_init_frag};
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::Segment;

// 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_init(
    cache: &Cache,
    config: &Config,
    kind: TrackKind,
    input_format: &StreamFormatInfo,
    output_format: &StreamFormatInfo,
) -> Result<Segment> {
    let command = transcode_command(kind, input_format, output_format, true, config)?;
    let output = cache.cache(
        &format!(
            "transcode/media-init/{}-{}.mkv",
            input_format.metadata_str(),
            HashKey(&command)
        ),
        || {
            info!("init encode with {command:?}");
            let mut args = command.split(" ");
            let proc = Command::new(args.next().unwrap())
                .stdout(Stdio::piped())
                .args(args)
                .spawn()?;
            Ok(proc.wait_with_output()?.exit_ok()?.stdout)
        },
    )?;

    let mut demuxer = create_demuxer(ContainerFormat::Matroska, Box::new(Cursor::new(output)));
    let info = demuxer.info()?;
    let tracks = demuxer.tracks()?;
    Ok(Segment {
        info,
        tracks,
        ..Default::default()
    })
}

pub fn transcode(
    cache: &Cache,
    config: &Config,
    kind: TrackKind,
    input_key: &str,
    input_format: &StreamFormatInfo,
    output_format: &StreamFormatInfo,
    segment: impl FnOnce() -> Result<Segment>,
) -> Result<Segment> {
    let command = transcode_command(kind, &input_format, output_format, false, config).unwrap();

    let output = cache.cache(
        &format!(
            "transcode/media-fragment/{input_key}-{}.mkv",
            HashKey(&command)
        ),
        || {
            let _permit = LOCAL_VIDEO_TRANSCODING_TASKS.lock().unwrap();
            let input = segment()?;
            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();

            let mut buf = Vec::new();
            write_init_frag(ContainerFormat::Matroska, &mut buf, input)?;

            spawn(move || {
                stdin.write_all(&buf).unwrap();
                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 info = demuxer.info()?;
    let tracks = demuxer.tracks()?;
    let mut clusters = Vec::new();
    while let Some((_, cluster)) = demuxer.read_cluster()? {
        clusters.push(cluster);
    }

    //? Remove extra kf hack
    if clusters
        .last()
        .map_or(false, |c| c.simple_blocks.len() == 1)
    {
        clusters.pop();
    }

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

fn transcode_command(
    kind: TrackKind,
    input: &StreamFormatInfo,
    output: &StreamFormatInfo,
    dummy: bool,
    config: &Config,
) -> Result<String> {
    let br = output.bitrate as u64;
    let w = output.width.unwrap_or(0);
    let h = output.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 ")?;
        }

        if dummy {
            write!(o, "-f lavfi -i testsrc2=s={w}x{h},format=yuv420p -to 1 ")?;
        } else {
            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 output.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 output.codec == "A_OPUS" && input.channels.unwrap_or(2) > 2 {
            write!(o, "-ac 2 ")?;
        }
        match output.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)
}