use ebml_struct::{
    Size,
    ids::{
        EL_ATTACHMENTS, EL_CHAPTERS, EL_CLUSTER, EL_CRC32, EL_CUES, EL_EBML, EL_INFO, EL_SEEKHEAD,
        EL_SEGMENT, EL_TAGS, EL_TRACKS, EL_VOID,
    },
    matroska::{Cluster, Ebml, Info, Tags, Tracks},
    read::{EbmlReadExt, TagRead},
    write::{EbmlWriteExt, TagWrite},
};
use std::io::{BufReader, BufWriter, ErrorKind, Read, Result, stdin, stdout};

pub trait BitstreamFilter: Sized {
    const INPUT_CODEC_ID: &str;
    const OUTPUT_CODEC_ID: &str;
    fn new(width: u32, height: u32) -> Self;
    fn process_block(&mut self, a: Vec<u8>) -> Vec<u8>;
}

pub fn bitstream_filter_main<F: BitstreamFilter>() -> Result<()> {
    let mut inp = BufReader::new(stdin());
    let mut out = BufWriter::new(stdout());

    let (id, inner) = inp.read_tag()?;
    assert_eq!(id, EL_EBML);
    inner.consume()?;

    Ebml {
        ebml_version: 1,
        ebml_read_version: 1,
        ebml_max_id_length: 4,
        ebml_max_size_length: 8,
        doc_type: "matroska".to_string(),
        doc_type_version: 4,
        doc_type_read_version: 2,
        doc_type_extensions: vec![],
    }
    .write(&mut out)?;

    let (id, mut segment) = inp.read_tag()?;
    assert_eq!(id, EL_SEGMENT);

    out.write_tag_id(EL_SEGMENT)?;
    out.write_size(Size::Unknown)?;

    let (mut info, mut tracks) = read_info_and_tracks(&mut segment)?;
    info.muxing_app = "ebml-struct".to_string();
    info.writing_app = "video-codec-experiments".to_string();

    let mut encode_track = None;
    for t in &mut tracks.entries {
        if t.track_type == 1 && encode_track.is_none() {
            assert_eq!(t.codec_id, F::INPUT_CODEC_ID);
            t.codec_id = F::OUTPUT_CODEC_ID.to_string();
            t.codec_private = None;
            t.codec_download_urls =
                vec!["https://codeberg.org/metamuffin/video-codec-experiments".to_string()];

            encode_track = Some(t.clone());
        }
    }
    let filter_track = encode_track.expect("no video track found");

    tracks.write(&mut out)?;
    info.write(&mut out)?;

    let mut filter = F::new(
        filter_track.video.as_ref().unwrap().pixel_width as u32,
        filter_track.video.as_ref().unwrap().pixel_height as u32,
    );

    loop {
        let (id, mut element) = match segment.read_tag() {
            Ok(x) => x,
            Err(e) if e.kind() == ErrorKind::UnexpectedEof => break,
            Err(e) => return Err(e),
        };
        match id {
            EL_CUES | EL_ATTACHMENTS | EL_SEEKHEAD | EL_CHAPTERS | EL_CRC32 | EL_VOID => {
                element.consume()?;
            }
            EL_TAGS => {
                Tags::read(&mut element)?.write(&mut out)?;
            }
            EL_CLUSTER => {
                let mut cluster = Cluster::read(&mut element)?;

                for b in &mut cluster.simple_blocks {
                    if b.track == filter_track.track_number {
                        assert_eq!(b.lacing, None);
                        assert_eq!(b.discardable, false);
                        assert_eq!(b.keyframe, true);
                        b.data = filter.process_block(b.data.clone());
                    }
                }

                cluster.write(&mut out)?;
            }
            _ => {
                eprintln!("unhandled element {id:x}");
                element.consume()?;
            }
        }
    }

    Ok(())
}

fn read_info_and_tracks(mut segment: impl Read) -> Result<(Info, Tracks)> {
    let mut info = None;
    let mut tracks = None;
    loop {
        let (id, mut element) = segment.read_tag()?;
        match id {
            EL_CUES | EL_TAGS | EL_ATTACHMENTS | EL_SEEKHEAD | EL_CHAPTERS | EL_CRC32 | EL_VOID => {
                element.consume()?;
            }
            EL_INFO => {
                info = Some(Info::read(&mut element)?);
            }
            EL_TRACKS => {
                tracks = Some(Tracks::read(&mut element)?);
            }
            EL_CLUSTER => {
                eprintln!("not ready for clusters yet");
                element.consume()?;
            }
            _ => {
                eprintln!("unhandled element {id:x}");
                element.consume()?;
            }
        };
        if tracks.is_some() && info.is_some() {
            break Ok((info.unwrap(), tracks.unwrap()));
        }
    }
}