/*
    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::{cues::generate_cues, metadata::read_metadata, SMediaInfo, CONF};
use anyhow::Result;
use jellyremuxer::matroska::{self, Segment, TrackEntry, TrackType};
use jellystream_types::{
    StreamContainer, StreamFormatInfo, StreamInfo, StreamTrackInfo, TrackKind,
};
use std::{
    io::{Cursor, Read},
    path::PathBuf,
    sync::Arc,
};

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

// TODO cache mem
pub(crate) 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 = read_metadata(&path)?;
        let cue_stat = generate_cues(&path)?;
        if let Some(t) = &metadata.tracks {
            let duration = media_duration(&metadata.info);
            for t in &t.entries {
                let byterate = cue_stat
                    .stats
                    .get(&t.track_number)
                    .map(|e| e.total_size)
                    .unwrap_or_default() as f64
                    / duration;
                tracks.push(StreamTrackInfo {
                    name: None,
                    kind: match t.track_type {
                        matroska::TrackType::Video => TrackKind::Video,
                        matroska::TrackType::Audio => TrackKind::Audio,
                        matroska::TrackType::Subtitle => TrackKind::Subtitle,
                        _ => todo!(),
                    },
                    formats: stream_formats(t, byterate * 8.),
                });
                track_to_file.push((i, t.track_number));
            }
        }
        metadata_arr.push(metadata);
        paths.push(path.to_owned());
    }

    let duration = media_duration(&metadata_arr[0].info); // TODO different durations?!
    Ok((
        InternalStreamInfo {
            metadata: metadata_arr,
            paths,
            track_to_file,
        },
        StreamInfo {
            name: info.title.clone(),
            duration,
            tracks,
        },
    ))
}

fn stream_formats(t: &TrackEntry, 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 {
        TrackType::Video => {
            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,
                        });
                    }
                }
            }
        }
        TrackType::Audio => {
            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),
                });
            }
        }
        TrackType::Subtitle => {}
        _ => {}
    }

    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) fn write_stream_info(info: Arc<SMediaInfo>) -> Result<Box<dyn Read + Send + Sync>> {
    let (_, info) = stream_info(info)?;
    Ok(Box::new(Cursor::new(serde_json::to_vec(&info)?)))
}

fn media_duration(info: &matroska::Info) -> f64 {
    (info.duration.unwrap_or_default() * info.timestamp_scale as f64) / 1_000_000_000.
}