/*
    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::{BlockIndex, 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 {
    #[allow(clippy::new_without_default)]
    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,
            temp_index: usize,
        }

        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,
                    temp_index: 0,
                })
            })
            .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),
            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 ClusterLayout {
            position: usize,
            timestamp: u64,
            blocks: Vec<(usize, BlockIndex)>,
        }

        let segment_layout = {
            let mut cluster_pts = 0;
            let mut clusters = vec![];
            let mut cluster = vec![];
            let mut gp = 0usize; // cluster position (in the segment)
            let mut p = 0usize; // block position (in the cluster)
            loop {
                let mut best_block = BlockIndex {
                    pts: u64::MAX,
                    size: 0,
                    source_off: 0,
                };
                let mut best_index = 0;
                for (i, r) in inputs.iter().enumerate() {
                    if let Some(v) = r.index.blocks.get(r.temp_index) {
                        if v.pts < best_block.pts {
                            best_block = v.to_owned();
                            best_index = i;
                        }
                    };
                }
                inputs[best_index].temp_index += 1;
                if best_block.pts > cluster_pts + 2_000 {
                    let cluster_content_size = 1 // timestamp tag
                     + 1 // timestamp tag size
                     + EbmlWriter::vint_length(cluster_pts) // timestamp tag value
                     + p;
                    let cluster_header_size = 4 // tag length
                      + EbmlWriter::vint_length(cluster_content_size as u64)// size varint
                      + cluster_content_size;
                    clusters.push(ClusterLayout {
                        position: gp,
                        timestamp: cluster_pts,
                        blocks: std::mem::take(&mut cluster),
                    });

                    cluster_pts = best_block.pts;
                    gp += p + cluster_header_size;
                    p = 0;
                }
                if best_block.pts == u64::MAX {
                    break;
                }
                p += 1; // simpleblock tag
                p += EbmlWriter::vint_length(1 + 2 + 1 + best_block.size as u64); // simpleblock size vint
                p += 1 + 2 + 1; // block {tracknum, pts_off, flags}
                                // TODO does not work, if more than 127 tracks are present
                p += best_block.size; // block payload
                cluster.push((best_index, best_block))
            }
            info!("segment layout computed ({} clusters)", clusters.len());
            clusters
        };

        output.write_tag(&MatroskaTag::Cues(Master::Collected(
            segment_layout
                .iter()
                .map(|cluster| {
                    MatroskaTag::CuePoint(Master::Collected(
                        [
                            MatroskaTag::CueTime(cluster.timestamp),
                            MatroskaTag::CueTrackPositions(Master::Collected(
                                [
                                    MatroskaTag::CueTrack(0),
                                    MatroskaTag::CueClusterPosition(cluster.position as u64),
                                ]
                                .to_vec(),
                            )),
                            MatroskaTag::CueTrackPositions(Master::Collected(
                                [
                                    MatroskaTag::CueTrack(1),
                                    MatroskaTag::CueClusterPosition(cluster.position as u64),
                                ]
                                .to_vec(),
                            )),
                        ]
                        .to_vec(),
                    ))
                })
                .collect(),
        )))?;

        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(),
            });
        }

        let segment_start_position = output.position();
        for (cluster_index, cluster) in segment_layout.into_iter().enumerate() {
            info!(
                "writing cluster {cluster_index} (pts_base={}) with {} blocks",
                cluster.timestamp,
                cluster.blocks.len()
            );
            debug!(
                "calculation was {} bytes off",
                cluster.position as i64 - (output.position() - segment_start_position) as i64
            );
            let mut cluster_blocks = vec![MatroskaTag::Timestamp(cluster.timestamp)];
            for (block_index, iblock) in cluster.blocks {
                let kn = &mut ks[block_index];
                let mut block = kn
                    .peek
                    .replace(kn.stream.next()?)
                    .expect("source file too short");

                assert_eq!(iblock.size, block.inner.data.len(), "seek index is wrong");
                assert_eq!(iblock.pts, block.pts(), "seek index is wrong");

                block.inner.track = kn.mapped;
                block.inner.timestamp_off = (iblock.pts - cluster.timestamp).try_into().unwrap();
                debug!("n={} tso={}", block.inner.track, block.inner.timestamp_off);
                let buf = block.inner.dump();
                cluster_blocks.push(MatroskaTag::SimpleBlock(buf))
            }
            output.write_tag(&MatroskaTag::Cluster(Master::Collected(cluster_blocks)))?;
        }
        output.write_tag(&MatroskaTag::Segment(Master::End))?;
        Ok(())
    }
}

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

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

impl AbsoluteBlock {
    pub fn pts(&self) -> u64 {
        self.inner.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.n().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.n() {
                    match item {
                        MatroskaTag::Crc32(_) => (),
                        MatroskaTag::Timestamp(ts) => {
                            trace!("read pts={ts}");
                            pts_base = ts;
                        }
                        MatroskaTag::BlockGroup(_) => {
                            trace!("group");
                            let mut children = children.unwrap();

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

                            while let Some(Ok(Unflat { children: _, item })) = children.n() {
                                match item {
                                    MatroskaTag::Block(buf) => block = Some(buf),
                                    // MatroskaTag::BlockDuration(v) => duration = Some(v),
                                    _ => debug!("ignored {item:?}"),
                                }
                            }
                            // TODO 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,
                                    inner: 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,
                                    inner: 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![
        MatroskaTag::TrackNumber(number),
        MatroskaTag::TrackUID(number),
        MatroskaTag::FlagLacing(0),
        MatroskaTag::Language(track.language.clone()),
        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))
}