/*
    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>
*/
use crate::{matroska::MatroskaTag, size::EbmlSize, Master};
use anyhow::{bail, Result};
use std::io::Write;

pub struct EbmlWriter {
    inner: Box<dyn Write>,
    position: usize,
}

impl EbmlWriter {
    pub fn new<T: Write + 'static>(inner: T, position: usize) -> Self {
        Self {
            inner: Box::new(inner),
            position,
        }
    }

    pub fn write(&mut self, data: &[u8]) -> Result<()> {
        self.inner.write_all(data)?;
        self.position += data.len();
        Ok(())
    }

    pub fn write_tag(&mut self, tag: &MatroskaTag) -> Result<()> {
        let mut buf = vec![];
        tag.write_full(&mut buf)?;
        self.write(&buf)?;
        Ok(())
    }

    pub fn write_vint(&mut self, i: u64) -> Result<()> {
        if i > (1 << 56) - 1 {
            bail!("vint does not fit");
        }
        let mut len = 1;
        while len <= 8 {
            if i < (1 << ((7 * len) - 1)) {
                break;
            }
            len += 1;
        }
        let mut bytes = i.to_be_bytes();
        let trunc = &mut bytes[(8 - len)..];
        trunc[0] |= 1 << (8 - len);
        self.write(&trunc)
    }
}

impl MatroskaTag {
    pub fn write_full(&self, w: &mut Vec<u8>) -> Result<()> {
        let mut buf = vec![];
        buf.extend(self.id().to_be_bytes().iter().skip_while(|&v| *v == 0u8));
        // note: it is relevant here, to pass the buffer with the id, such that closing tags, can clear it
        self.write(&mut buf)?;
        w.extend_from_slice(&buf);
        Ok(())
    }
}

pub trait WriteValue {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()>;
}

impl WriteValue for i64 {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
        Ok(match 64 - self.leading_zeros() {
            x if x <= 8 => {
                w.push(0x81);
                w.extend_from_slice(&(*self as i8).to_be_bytes());
            }
            x if x <= 16 => {
                w.push(0x82);
                w.extend_from_slice(&(*self as i16).to_be_bytes());
            }
            x if x <= 32 => {
                w.push(0x84);
                w.extend_from_slice(&(*self as i32).to_be_bytes());
            }
            _ => {
                w.push(0x88);
                w.extend_from_slice(&self.to_be_bytes());
            }
        })
    }
}
impl WriteValue for u64 {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
        Ok(match 64 - self.leading_zeros() {
            x if x <= 8 => {
                w.push(0x81);
                w.extend_from_slice(&(*self as u8).to_be_bytes());
            }
            x if x <= 16 => {
                w.push(0x82);
                w.extend_from_slice(&(*self as u16).to_be_bytes());
            }
            x if x <= 32 => {
                w.push(0x84);
                w.extend_from_slice(&(*self as u32).to_be_bytes());
            }
            _ => {
                w.push(0x88);
                w.extend_from_slice(&self.to_be_bytes());
            }
        })
    }
}
impl WriteValue for f64 {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<(), anyhow::Error> {
        w.push(0x88);
        w.extend_from_slice(&self.to_be_bytes());
        Ok(())
    }
}
impl WriteValue for Vec<u8> {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<(), anyhow::Error> {
        write_vint(w, self.len() as u64)?;
        w.extend_from_slice(&self);
        Ok(())
    }
}
impl WriteValue for String {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<(), anyhow::Error> {
        let sl = self.as_bytes();
        write_vint(w, sl.len() as u64)?;
        w.extend_from_slice(sl);
        Ok(())
    }
}
impl WriteValue for EbmlSize {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
        match self {
            EbmlSize::Exact(s) => write_vint(w, *s as u64)?,
            EbmlSize::Unknown => w.extend_from_slice(&(u64::MAX >> 7).to_be_bytes()),
        }
        Ok(())
    }
}

impl WriteValue for Master {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
        match self {
            Master::Start => EbmlSize::Unknown.write_to(w),
            Master::End => Ok(w.clear()),
            Master::Collected(c) => {
                let mut ib = vec![];
                for c in c {
                    c.write_full(&mut ib)?;
                }
                EbmlSize::Exact(ib.len()).write_to(w)?;
                w.extend_from_slice(&ib);
                Ok(())
            }
        }
    }
}

pub fn write_vint(w: &mut Vec<u8>, i: u64) -> Result<()> {
    if i > (1 << 56) - 1 {
        bail!("vint does not fit");
    }
    let len = (64 - i.leading_zeros() as usize) / 7 + 1;
    let mut bytes = i.to_be_bytes();
    let trunc = &mut bytes[(8 - len)..];
    trunc[0] |= 1 << (8 - len);
    w.extend_from_slice(&trunc);
    Ok(())
}