/*
    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) 2024 metamuffin <metamuffin.org>
*/
use anyhow::{Context, Result};
use jellybase::cache::cache_memory;
use jellycommon::seek_index::{BlockIndex, SeekIndex};
use jellymatroska::{
    block::Block,
    read::EbmlReader,
    unflatten::{Unflat, Unflatten},
    MatroskaTag,
};
use log::{debug, info, trace, warn};
use std::{collections::BTreeMap, fs::File, io::BufReader, path::Path, sync::Arc};

pub fn get_seek_index(path: &Path) -> anyhow::Result<Arc<BTreeMap<u64, Arc<SeekIndex>>>> {
    cache_memory(&["seekindex", path.to_str().unwrap()], move || {
        info!("generating seek index for {path:?}");
        let input = File::open(&path).context("opening source file")?;
        let mut input = EbmlReader::new(BufReader::new(input));
        let index = import_seek_index(&mut input)?;
        info!("done");
        Ok(index.into_iter().map(|(k, v)| (k, Arc::new(v))).collect())
    })
}

pub fn import_seek_index(input: &mut EbmlReader) -> Result<BTreeMap<u64, SeekIndex>> {
    let mut seek_index = BTreeMap::new();
    while let Some(item) = input.next() {
        let item = match item {
            Ok((_, item)) => item,
            Err(e) => {
                if !matches!(e, jellymatroska::error::Error::Io(_)) {
                    warn!("{e}");
                }
                break;
            }
        };
        match item {
            MatroskaTag::Segment(_) => {
                info!("segment start");
                let mut children = Unflatten::new_with_end(input, item);
                import_seek_index_segment(&mut children, &mut seek_index)?;
                info!("segment end");
            }
            _ => debug!("(r) tag ignored: {item:?}"),
        }
    }
    Ok(seek_index)
}

fn import_seek_index_segment(
    segment: &mut Unflatten,
    seek_index: &mut BTreeMap<u64, SeekIndex>,
) -> Result<()> {
    while let Some(Ok(Unflat { children, item, .. })) = segment.n() {
        match item {
            MatroskaTag::SeekHead(_) => {}
            MatroskaTag::Info(_) => {}
            MatroskaTag::Tags(_) => {}
            MatroskaTag::Cues(_) => {}
            MatroskaTag::Chapters(_) => {}
            MatroskaTag::Tracks(_) => {}
            MatroskaTag::Void(_) => {}
            MatroskaTag::Cluster(_) => {
                let mut children = children.unwrap();
                let mut pts = 0;
                loop {
                    if let Some(Ok(Unflat {
                        children,
                        item,
                        position,
                    })) = children.n()
                    {
                        match item {
                            MatroskaTag::Timestamp(ts) => pts = ts,
                            MatroskaTag::BlockGroup(_) => {
                                trace!("group");
                                let mut children = children.unwrap();
                                while let Some(Ok(Unflat {
                                    children: _, item, ..
                                })) = children.n()
                                {
                                    match item {
                                        MatroskaTag::Block(ref buf) => {
                                            let block = Block::parse(buf)?;
                                            debug!(
                                                "block: track={} tso={}",
                                                block.track, block.timestamp_off
                                            );
                                            seek_index_add(
                                                seek_index,
                                                &block,
                                                position.unwrap(),
                                                pts,
                                            );
                                        }
                                        _ => trace!("{item:?}"),
                                    }
                                }
                            }
                            MatroskaTag::SimpleBlock(buf) => {
                                let block = Block::parse(&buf)?;
                                trace!(
                                    "simple block: track={} tso={}",
                                    block.track,
                                    block.timestamp_off
                                );
                                trace!("{pts} {}", block.timestamp_off);
                                seek_index_add(seek_index, &block, position.unwrap(), pts);
                            }
                            _ => trace!("(rsc) tag ignored: {item:?}"),
                        }
                    } else {
                        break;
                    }
                }
            }
            _ => debug!("(rs) tag ignored: {item:?}"),
        };
    }
    Ok(())
}

fn seek_index_add(
    seek_index: &mut BTreeMap<u64, SeekIndex>,
    block: &Block,
    position: u64,
    pts_base: u64,
) {
    //* I heard this helped debugging once.
    // {
    //     let mut f = File::open("/home/muffin/videos/itte-yorushika.mkv").unwrap();
    //     f.seek(std::io::SeekFrom::Start(position.try_into().unwrap()))
    //         .unwrap();
    //     let mut buf = [0u8];
    //     f.read_exact(&mut buf).unwrap();

    //     eprintln!("{}", buf[0]);
    //     if buf[0] != 0xa0 && buf[0] != 0xa3 {
    //         warn!("invalid position {position}")
    //     }
    // }

    let trs = seek_index
        .entry(block.track)
        .or_insert(SeekIndex::default());

    if block.keyframe {
        trs.keyframes.push(trs.blocks.len());
    }
    trs.blocks.push(BlockIndex {
        pts: (pts_base as i64 + block.timestamp_off as i64) as u64,
        source_off: position,
        size: block.data.len(),
    });
}