/*
    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) 2025 metamuffin <metamuffin.org>
*/
use crate::{SMediaInfo, CONF};
use anyhow::Result;
use jellycommon::stream::{
    StreamContainer, StreamFormatInfo, StreamInfo, StreamSegmentInfo, StreamTrackInfo, TrackKind,
};
use jellyremuxer::{
    metadata::{matroska_metadata, MatroskaMetadata, MatroskaTrackEntry},
    seek_index::get_track_sizes,
};
use std::{collections::BTreeMap, path::PathBuf, sync::Arc};
use tokio::{
    io::{AsyncWriteExt, DuplexStream},
    spawn,
    task::spawn_blocking,
};

async fn async_matroska_metadata(path: PathBuf) -> Result<Arc<MatroskaMetadata>> {
    spawn_blocking(move || matroska_metadata(&path)).await?
}

async fn async_get_track_sizes(path: PathBuf) -> Result<BTreeMap<u64, usize>> {
    spawn_blocking(move || get_track_sizes(&path)).await?
}

pub(crate) struct InternalStreamInfo {
    pub paths: Vec<PathBuf>,
    pub metadata: Vec<Arc<MatroskaMetadata>>,
    pub track_to_file: Vec<(usize, u64)>,
}

// TODO cache mem
pub(crate) async fn stream_info(info: Arc<SMediaInfo>) -> Result<(InternalStreamInfo, StreamInfo)> {
    let mut tracks = Vec::new();
    let mut track_to_file = Vec::new();
    let mut metadata_arr = Vec::new();
    let mut paths = Vec::new();
    for (i, path) in info.files.iter().enumerate() {
        let metadata = async_matroska_metadata(path.clone()).await?;
        let sizes = async_get_track_sizes(path.clone()).await?;
        if let Some(t) = &metadata.tracks {
            let duration = media_duration(&metadata);
            for t in &t.entries {
                let bitrate =
                    sizes.get(&t.track_number).copied().unwrap_or_default() as f64 / duration * 8.;
                tracks.push(StreamTrackInfo {
                    name: None,
                    kind: match t.track_type {
                        1 => TrackKind::Video,
                        2 => TrackKind::Audio,
                        17 => TrackKind::Subtitle,
                        _ => todo!(),
                    },
                    formats: stream_formats(t, bitrate),
                });
                track_to_file.push((i, t.track_number));
            }
        }
        metadata_arr.push(metadata);
        paths.push(path.to_owned());
    }

    let segment = StreamSegmentInfo {
        name: None,
        duration: media_duration(&metadata_arr[0]),
        tracks,
    };
    Ok((
        InternalStreamInfo {
            metadata: metadata_arr,
            paths,
            track_to_file,
        },
        StreamInfo {
            name: info.info.title.clone(),
            segments: vec![segment],
        },
    ))
}

fn stream_formats(t: &MatroskaTrackEntry, remux_bitrate: f64) -> Vec<StreamFormatInfo> {
    let mut formats = Vec::new();
    formats.push(StreamFormatInfo {
        codec: t.codec_id.to_string(),
        remux: true,
        bitrate: remux_bitrate,
        containers: containers_by_codec(&t.codec_id),
        bit_depth: t.audio.as_ref().and_then(|a| a.bit_depth.map(|e| e as u8)),
        samplerate: t.audio.as_ref().map(|a| a.sampling_frequency),
        channels: t.audio.as_ref().map(|a| a.channels as usize),
        width: t.video.as_ref().map(|v| v.pixel_width),
        height: t.video.as_ref().map(|v| v.pixel_height),
    });

    match t.track_type {
        1 => {
            let sw = t.video.as_ref().unwrap().pixel_width;
            let sh = t.video.as_ref().unwrap().pixel_height;
            for (w, br) in [
                (3840, 6000e3),
                (1920, 5000e3),
                (1920, 2000e3),
                (1280, 1500e3),
                (640, 800e3),
                (320, 200e3),
            ] {
                if w > sw {
                    continue;
                }
                // most codecs use chroma subsampling that requires even dims
                let h = ((w * sh) / sw) & !1; // clear last bit to ensure even height.
                for (cid, enable) in [
                    ("V_AV1", CONF.offer_av1),
                    ("V_VP8", CONF.offer_vp8),
                    ("V_VP9", CONF.offer_vp9),
                    ("V_MPEG4/ISO/AVC", CONF.offer_avc),
                    ("V_MPEGH/ISO/HEVC", CONF.offer_hevc),
                ] {
                    if enable {
                        formats.push(StreamFormatInfo {
                            codec: cid.to_string(),
                            bitrate: (remux_bitrate * 3.).min(br),
                            remux: false,
                            containers: containers_by_codec(cid),
                            width: Some(w),
                            height: Some(h),
                            samplerate: None,
                            channels: None,
                            bit_depth: None,
                        });
                    }
                }
            }
        }
        2 => {
            for br in [256e3, 128e3, 64e3] {
                formats.push(StreamFormatInfo {
                    codec: "A_OPUS".to_string(),
                    bitrate: br,
                    remux: false,
                    containers: containers_by_codec("A_OPUS"),
                    width: None,
                    height: None,
                    samplerate: Some(48e3),
                    channels: Some(2),
                    bit_depth: Some(32),
                });
            }
        }
        17 => {}
        _ => {}
    }

    formats
}

fn containers_by_codec(codec: &str) -> Vec<StreamContainer> {
    use StreamContainer::*;
    match codec {
        "V_VP8" | "V_VP9" | "V_AV1" | "A_OPUS" | "A_VORBIS" => vec![Matroska, WebM],
        "V_MPEG4/ISO/AVC" | "A_AAC" => vec![Matroska, MPEG4],
        "S_TEXT/UTF8" | "S_TEXT/WEBVTT" => vec![Matroska, WebVTT, WebM, JVTT],
        _ => vec![Matroska],
    }
}

pub(crate) async fn write_stream_info(info: Arc<SMediaInfo>, mut b: DuplexStream) -> Result<()> {
    let (_, info) = stream_info(info).await?;
    spawn(async move { b.write_all(&serde_json::to_vec(&info)?).await });
    Ok(())
}

fn media_duration(m: &MatroskaMetadata) -> f64 {
    let info = m.info.as_ref().unwrap();
    (info.duration.unwrap_or_default() * info.timestamp_scale as f64) / 1_000_000_000.
}