/*
    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 log::debug;
use std::io::{Seek, Write};

pub struct EbmlWriter<W> {
    inner: W,
    position: usize,
}

impl<W: Write> EbmlWriter<W> {
    pub fn new(inner: W, position: usize) -> Self {
        Self { inner, position }
    }

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

    pub fn position(&self) -> usize {
        self.position
    }

    pub fn write_padding(&mut self, position: usize) -> Result<()> {
        debug!("padding up to {position}");
        let mut size = position - self.position;
        match size {
            0 => return Ok(()),
            1 => bail!("this is sadly not possible"),
            _ => (),
        }
        size -= 1; // subtract tag size
        size -= 4; // subtract vint size

        // match size {
        //     _ if size < (1 << 7) => size -= 1,
        //     _ if size < (1 << 14) => size -= 2,
        //     _ if size < (1 << 21) => size -= 3,
        //     _ if size < (1 << 28) => size -= 4,
        //     _ if size < (1 << 35) => size -= 5,
        //     _ => bail!("padding to large"),
        // }

        self.write(&[0xec])?;
        self.write_vint_len(size.try_into().unwrap(), 4)?;
        self.write(&vec![0; size])?;
        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");
        }
        self.write_vint_len(i, vint_length(i))
    }

    pub fn write_vint_len(&mut self, i: u64, len: usize) -> Result<()> {
        let mut bytes = i.to_be_bytes();
        let trunc = &mut bytes[(8 - len)..];
        trunc[0] |= 1 << (8 - len);
        self.write(trunc)
    }
}

impl<W: Seek> Seek for EbmlWriter<W> {
    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
        self.inner.seek(pos)?;
        match pos {
            std::io::SeekFrom::Start(s) => self.position = s as usize,
            std::io::SeekFrom::End(_) => unimplemented!(),
            std::io::SeekFrom::Current(s) => self.position += s as usize,
        }
        Ok(self.position as u64)
    }
}

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

/// this routine works only, if the varint is as small as it can possibly be.
/// thats not always what we do though - see below
pub fn vint_length(v: u64) -> usize {
    let mut len = 1;
    while len <= 8 {
        if v < (1 << ((7 * len) - 1)) {
            break;
        }
        len += 1;
    }
    len
}
pub fn bad_vint_length(v: u64) -> usize {
    match 64 - v.leading_zeros() {
        x if x <= 8 => 1,
        x if x <= 16 => 2,
        x if x <= 32 => 4,
        _ => 8,
    }
}

pub trait WriteValue {
    /// writes the contents of a tag, including the size but excluding the id.
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()>;
}

impl WriteValue for i64 {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
        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());
            }
        };
        Ok(())
    }
}
impl WriteValue for u64 {
    fn write_to(&self, w: &mut Vec<u8>) -> Result<()> {
        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());
            }
        };
        Ok(())
    }
}
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 => {
                w.clear();
                Ok(())
            }
            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(())
}