pub mod format;
pub mod import;

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

pub struct RemuxerContext {}

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

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

        let input = File::open(source_path)?;
        let mut input = EbmlReader::new(input);
        let mut output = EbmlWriter::new(writer, 0);

        // maps original to remuxed
        let mapping = BTreeMap::from_iter(
            selection
                .iter()
                .enumerate()
                .map(|(i, e)| (*e, i as u64 + 1)),
        );
        info!("track mapping: {mapping:?}");

        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 = MatroskaTag::Tracks(Master::Collected(
            mapping
                .iter()
                .map(|(norig, nrem)| track_to_ebml(*nrem, &iteminfo.tracks.get(norig).unwrap()))
                .collect(),
        ));
        info!("track header: {tracks_header:?}");
        output.write_tag(&tracks_header)?;

        while let Some(item) = input.next() {
            let item = match item {
                Ok(item) => item,
                Err(e) => {
                    warn!("{e}");
                    break;
                }
            };
            match item {
                MatroskaTag::Ebml(_) => {
                    Unflatten::new_with_end(&mut input, item);
                }
                MatroskaTag::Segment(_) => {
                    info!("segment start");
                    let mut children = Unflatten::new_with_end(&mut input, item);
                    filter_segment(&mut children, &mut output, &iteminfo, &mapping)?;
                    info!("segment end");
                }
                _ => debug!("(r) tag ignored: {item:?}"),
            }
        }

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

        Ok(())
    }
}

fn filter_segment(
    children: &mut Unflatten,
    writer: &mut EbmlWriter,
    _iteminfo: &ItemInfo,
    mapping: &BTreeMap<u64, u64>,
) -> Result<()> {
    while let Some(Ok(Unflat { children, item })) = children.next() {
        match item {
            MatroskaTag::SeekHead(_) => {}
            MatroskaTag::Info(_) => {}
            MatroskaTag::Cluster(_) => {
                let mut cluster = vec![];
                let mut children = children.unwrap();
                while let Some(Ok(Unflat { children, item })) = children.next() {
                    match item {
                        MatroskaTag::Crc32(_) => (),
                        MatroskaTag::Timestamp(ts) => {
                            info!("ts={ts}");
                            cluster.push(MatroskaTag::Timestamp(ts));
                        }
                        MatroskaTag::BlockGroup(_) => {
                            debug!("group");
                            let mut children = children.unwrap();
                            let mut group = vec![];
                            while let Some(Ok(Unflat { children: _, item })) = children.next() {
                                match item {
                                    MatroskaTag::Block(buf) => {
                                        let mut block = Block::parse(&buf)?;
                                        if let Some(outnum) = mapping.get(&block.track) {
                                            block.track = *outnum;
                                            trace!(
                                                "block: track={} tso={}",
                                                block.track,
                                                block.timestamp_off
                                            );
                                            group.push(MatroskaTag::Block(block.dump()));
                                        }
                                    }
                                    MatroskaTag::BlockDuration(v) => {
                                        group.push(MatroskaTag::BlockDuration(v));
                                    }
                                    _ => debug!("ignored {item:?}"),
                                }
                            }
                            cluster.push(MatroskaTag::BlockGroup(Master::Collected(group)));
                        }
                        MatroskaTag::SimpleBlock(buf) => {
                            let mut block = Block::parse(&buf)?;
                            if let Some(outnum) = mapping.get(&block.track) {
                                block.track = *outnum;
                                trace!("block: track={} tso={}", block.track, block.timestamp_off);
                                cluster.push(MatroskaTag::SimpleBlock(block.dump()));
                            }
                        }
                        _ => warn!("(rsc) tag ignored: {item:?}"),
                    }
                }
                writer.write_tag(&MatroskaTag::Cluster(Master::Collected(cluster)))?;
            }
            MatroskaTag::Tags(_) => {}
            MatroskaTag::Cues(_) => {}
            MatroskaTag::Chapters(_) => {}
            MatroskaTag::Tracks(_) => {}
            _ => 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::DefaultDuration(41708333)); // TODO pls not hardcoded
            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))
}

// pub struct SendWriter(pub Sender<Vec<u8>>);

// impl Write for SendWriter {
//     fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
//         self.0.blocking_send(buf.to_owned()).unwrap();
//         Ok(buf.len())
//     }

//     fn flush(&mut self) -> std::io::Result<()> {
//         Ok(()) // TODO should we actually do something here?
//     }
// }