/*
    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) 2023 metamuffin <metamuffin.org>
*/
pub mod import;

use anyhow::{anyhow, Context, Result};
use jellycommon::{ItemInfo, SeekIndex, SourceTrack, SourceTrackKind};
use jellymatroska::{
    block::Block,
    read::EbmlReader,
    unflatten::{IterWithPos, Unflat, Unflatten},
    write::EbmlWriter,
    Master, MatroskaTag,
};
use log::{debug, info, trace, warn};
use std::{collections::VecDeque, fs::File, io::Write, path::PathBuf};

#[derive(Debug, Clone)]
pub struct RemuxerContext {}

impl RemuxerContext {
    pub fn new() -> Self {
        Self {}
    }

    pub fn generate_into(
        &self,
        writer: impl Write + 'static,
        _offset: usize,
        path_base: PathBuf,
        iteminfo: ItemInfo,
        selection: Vec<usize>,
        webm: bool,
    ) -> anyhow::Result<()> {
        info!("remuxing {:?} to have tracks {selection:?}", iteminfo.title);
        let mut output = EbmlWriter::new(writer, 0);

        struct ReaderC {
            info: SourceTrack,
            reader: EbmlReader,
            mapped: u64,
            index: SeekIndex,
        }

        let mut inputs = selection
            .iter()
            .enumerate()
            .map(|(index, sel)| {
                let info = iteminfo
                    .tracks
                    .get(sel)
                    .ok_or(anyhow!("track not available"))?
                    .to_owned();
                let source_path = path_base.join(&info.path);
                let mapped = index as u64 + 1;
                info!(
                    "\t- {sel} {source_path:?} ({} => {mapped})",
                    info.track_number
                );
                info!("\t    {}", info);
                let file = File::open(&source_path).context("opening source file")?;
                let mut index =
                    File::open(source_path.with_extension(&format!("si.{}", info.track_number)))
                        .context("opening seek index file")?;
                let index = bincode::decode_from_std_read::<SeekIndex, _, _>(
                    &mut index,
                    bincode::config::standard(),
                )?;
                debug!("\t    seek index: {} blocks loaded", index.blocks.len());
                let reader = EbmlReader::new(file);
                Ok(ReaderC {
                    index,
                    reader,
                    info,
                    mapped,
                })
            })
            .into_iter()
            .collect::<anyhow::Result<Vec<_>>>()?;

        output.write_tag(&MatroskaTag::Ebml(Master::Collected(vec![
            MatroskaTag::EbmlVersion(1),
            MatroskaTag::EbmlReadVersion(1),
            MatroskaTag::EbmlMaxIdLength(4),
            MatroskaTag::EbmlMaxSizeLength(8),
            MatroskaTag::DocType(if webm {
                "webm".to_string()
            } else {
                "matroska".to_string()
            }),
            MatroskaTag::DocTypeVersion(4),
            MatroskaTag::DocTypeReadVersion(2),
        ])))?;

        output.write_tag(&MatroskaTag::Segment(Master::Start))?;

        output.write_tag(&MatroskaTag::Info(Master::Collected(vec![
            MatroskaTag::TimestampScale(1_000_000),
            MatroskaTag::Duration(iteminfo.duration * 1000.0),
            MatroskaTag::Title(iteminfo.title.clone()),
            MatroskaTag::MuxingApp("jellyremux".to_string()),
            MatroskaTag::WritingApp("jellything".to_string()),
        ])))?;
        output.write_tag(&MatroskaTag::Tags(Master::Collected(vec![])))?;
        // output.write_tag(&MatroskaTag::Cues(Master::Collected(vec![])))?;

        let tracks_header = inputs
            .iter()
            .map(|rc| track_to_ebml(rc.mapped, &rc.info))
            .collect();
        output.write_tag(&MatroskaTag::Tracks(Master::Collected(tracks_header)))?;

        struct ReaderD<'a> {
            _info: SourceTrack,
            peek: Option<AbsoluteBlock>,
            stream: SegmentExtractIter<'a>,
            mapped: u64,
        }

        // read until start of the segment
        let mut ks = vec![];
        for i in &mut inputs {
            loop {
                let t = i.reader.next().ok_or(anyhow!("early eof"))??;
                if let MatroskaTag::Segment(Master::Start) = t {
                    break;
                }
            }
            let mut stream = SegmentExtractIter {
                segment: Unflatten::new_with_end(
                    &mut i.reader,
                    MatroskaTag::Segment(Master::Start),
                ),
                extract: i.info.track_number,
                emission_queue: VecDeque::new(),
            };
            ks.push(ReaderD {
                mapped: i.mapped,
                peek: Some(stream.next()?),
                stream,
                _info: i.info.clone(),
            });
        }

        loop {
            let next_index = ks
                .iter()
                .enumerate()
                .fold((0, u64::MAX), |(bi, bpts), (i, r)| {
                    if let Some(peek) = &r.peek {
                        let pts = peek.pts();
                        if pts < bpts {
                            return (i, pts);
                        }
                    }
                    (bi, bpts)
                })
                .0;

            let kn = &mut ks[next_index];
            let mut next = kn.peek.replace(kn.stream.next()?).unwrap(); //.ok_or(anyhow!("eof?"));

            let pts = next.pts();
            next.block.track = kn.mapped;
            next.block.timestamp_off = 0;
            let buf = next.block.dump();
            output.write_tag(&MatroskaTag::Cluster(Master::Collected(vec![
                MatroskaTag::Timestamp(pts),
                MatroskaTag::SimpleBlock(buf),
            ])))?;
        }

        // output.write_tag(&MatroskaTag::Segment(Master::End))?;
        // Ok(())
    }
}

struct AbsoluteBlock {
    pts_base: u64,
    block: Block,
}

struct SegmentExtractIter<'a> {
    segment: Unflatten<'a>,
    extract: u64,
    emission_queue: VecDeque<AbsoluteBlock>,
}

impl AbsoluteBlock {
    pub fn pts(&self) -> u64 {
        self.block.timestamp_off as u64 + self.pts_base
    }
}

impl SegmentExtractIter<'_> {
    pub fn next(&mut self) -> Result<AbsoluteBlock> {
        loop {
            if let Some(b) = self.emission_queue.pop_front() {
                break Ok(b);
            }
            self.read()?;
        }
    }

    pub fn read(&mut self) -> Result<()> {
        let Unflat { children, item } = self.segment.next().ok_or(anyhow!("eof"))??;
        let mut pts_base = 0;
        match item {
            MatroskaTag::SeekHead(_) => {}
            MatroskaTag::Info(_) => {}
            MatroskaTag::Cluster(_) => {
                let mut children = children.unwrap();
                while let Some(Ok(Unflat { children, item })) = children.next() {
                    match item {
                        MatroskaTag::Crc32(_) => (),
                        MatroskaTag::Timestamp(ts) => {
                            info!("ts={ts}");
                            pts_base = ts;
                        }
                        MatroskaTag::BlockGroup(_) => {
                            debug!("group");
                            let mut children = children.unwrap();

                            let mut duration = None;
                            let mut block = None;

                            while let Some(Ok(Unflat { children: _, item })) = children.next() {
                                match item {
                                    MatroskaTag::Block(buf) => block = Some(buf),
                                    MatroskaTag::BlockDuration(v) => duration = Some(v),
                                    _ => debug!("ignored {item:?}"),
                                }
                            }
                            // TODO duration
                            drop(duration);
                            let block = Block::parse(&block.unwrap())?;
                            if block.track == self.extract {
                                trace!("block: track={} tso={}", block.track, block.timestamp_off);
                                self.emission_queue
                                    .push_back(AbsoluteBlock { pts_base, block });
                            }
                        }
                        MatroskaTag::SimpleBlock(buf) => {
                            let block = Block::parse(&buf)?;
                            if block.track == self.extract {
                                trace!("block: track={} tso={}", block.track, block.timestamp_off);
                                self.emission_queue
                                    .push_back(AbsoluteBlock { pts_base, block });
                            }
                        }
                        _ => warn!("(rsc) tag ignored: {item:?}"),
                    }
                }
            }
            MatroskaTag::Tags(_) => {}
            MatroskaTag::Cues(_) => {}
            MatroskaTag::Chapters(_) => {}
            MatroskaTag::Tracks(_) => {}
            MatroskaTag::Void(_) => {}
            _ => debug!("(rs) tag ignored: {item:?}"),
        }
        Ok(())
    }
}

pub fn track_to_ebml(number: u64, track: &SourceTrack) -> MatroskaTag {
    let mut els = Vec::new();
    els.push(MatroskaTag::TrackNumber(number));
    els.push(MatroskaTag::TrackUID(number));
    els.push(MatroskaTag::FlagLacing(0));
    els.push(MatroskaTag::Language(track.language.clone()));
    els.push(MatroskaTag::CodecID(track.codec.clone()));
    if let Some(d) = &track.default_duration {
        els.push(MatroskaTag::DefaultDuration(*d));
    }
    match track.kind {
        SourceTrackKind::Video {
            width,
            height,
            fps: _,
        } => {
            els.push(MatroskaTag::TrackType(1));
            els.push(MatroskaTag::Video(Master::Collected(vec![
                MatroskaTag::PixelWidth(width),
                MatroskaTag::PixelHeight(height),
            ])))
        }
        SourceTrackKind::Audio {
            channels,
            sample_rate,
            bit_depth,
        } => {
            els.push(MatroskaTag::TrackType(2));
            els.push(MatroskaTag::Audio(Master::Collected(vec![
                MatroskaTag::SamplingFrequency(sample_rate),
                MatroskaTag::Channels(channels.try_into().unwrap()),
            ])));
            els.push(MatroskaTag::BitDepth(bit_depth.try_into().unwrap()));
        }
        SourceTrackKind::Subtitles => {
            els.push(MatroskaTag::TrackType(19));
        }
    }
    if let Some(d) = &track.codec_private {
        els.push(MatroskaTag::CodecPrivate(d.clone()));
    }
    MatroskaTag::TrackEntry(Master::Collected(els))
}