use alloc::vec;
use alloc::vec::Vec;
use crate::CURRENT_VERSION;
use crate::types::{CBTFlags, CBTHeader, HeaderError};

#[derive(Debug)]
pub struct CBT {
    pub header: CBTHeader,
    pub frames: Vec<Vec<u8>>,
}

pub struct CBTReader {
    pub frames: Vec<Vec<u8>>,
}

#[derive(Debug)]
pub enum CBTError {
    HeaderError(HeaderError),
    TooShort,
}

impl CBT {
    /// Creates a new CBT from the given bytes.
    /// Bytes should be the entire contents of a CBT file.
    pub fn new_from_bytes(bytes: &[u8]) -> Result<CBT, CBTError> {
        let header = CBTHeader::get_from_bytes(&bytes[0..16]);
        if let Err(e) = header {
            return Err(CBTError::HeaderError(e));
        }
        let header = header.unwrap();
        let data = bytes[16..].to_vec();
        let is_alpha = header.flags.has_alpha_channel;
        if data.len() < (header.width as usize * header.height as usize * header.frame_count as usize * {if is_alpha { 4 } else { 3 }}) {
            return Err(CBTError::TooShort);
        }
        let mut frames = Vec::new();
        for i in 0..header.frame_count as usize {
            let frame_size = header.width as usize * header.height as usize * {if is_alpha { 4 } else { 3 }};
            let frame_start = i * frame_size;
            let frame_end = frame_start + frame_size;
            let frame = &data[frame_start as usize..frame_end as usize];
            frames.push(frame.to_vec());
        }
        Ok(CBT {
            header,
            frames,
        })
    }

    /// Creates a new CBT with no data contained and a frame count of 0.
    /// Preferably use new_with_capacity instead, to save cpu time.
    pub fn new_empty(width: u16, height: u16, with_alpha_channel: bool) -> CBT {
        let frame_count = 0;
        let flags = CBTFlags {
            has_alpha_channel: with_alpha_channel,
        };
        let header = CBTHeader {
            version: CURRENT_VERSION,
            width,
            height,
            frame_count,
            flags,
        };
        let frames = Vec::new();
        CBT {
            header,
            frames,
        }
    }

    /// Creates a new CBT with a pre-allocated frame count.
    /// Preferred over new_empty, to save cpu time.
    pub fn new_with_capacity(width: u16, height: u16, frame_count: u16, with_alpha_channel: bool) -> CBT {
        let flags = CBTFlags {
            has_alpha_channel: with_alpha_channel,
        };
        let header = CBTHeader {
            version: CURRENT_VERSION,
            width,
            height,
            frame_count,
            flags,
        };
        let mut frames = Vec::new();
        let frame_size = header.width * header.height * {
            if header.flags.has_alpha_channel { 4 } else { 3 }
        };
        for _ in 0..frame_count {
            frames.push(vec![0; frame_size as usize]);
        }
        CBT {
            header,
            frames,
        }
    }

    /// Inserts the given data into the frame at the given index.
    pub fn fill_frame(&mut self, frame_index: usize, data: &[u8]) {
        self.frames[frame_index] = data.to_vec();
    }

    /// Converts the CBT struct to a byte array that would be in the CBT file format.
    pub fn to_file_data(&self) -> Vec<u8> {
        let header_bytes = self.header.to_bytes();
        let mut data = Vec::new();
        for frame in &self.frames {
            data.extend_from_slice(frame);
        }
        let final_size = header_bytes.len() + data.len();
        let mut final_bytes = Vec::with_capacity(final_size);
        final_bytes.extend_from_slice(&header_bytes);
        final_bytes.extend_from_slice(&data);
        final_bytes
    }

    /// Returns a CBTReader that can be used to read the frames in a specific format.
    /// This should only be used if you're too lazy to do the correct readings yourself,
    /// as it will consume a lot of memory.
    pub fn get_reader(&self, with_alpha_channel: bool) -> CBTReader {
        // if the file has no alpha channel and we don't want one, or vice versa, then we can just return the same CBT
        if self.header.flags.has_alpha_channel == with_alpha_channel {
            let mut frames = Vec::new();
            for frame in &self.frames {
                frames.push(frame.clone());
            }
            return CBTReader { frames };
        }
        // otherwise, we need to convert the frames
        let mut frames = Vec::new();
        match self.header.flags.has_alpha_channel {
            true => {
                for frame in &self.frames {
                    let mut new_frame = Vec::new();
                    let mut i = 0;
                    while i < frame.len() {
                        new_frame.push(frame[i]);
                        new_frame.push(frame[i + 1]);
                        new_frame.push(frame[i + 2]);
                        if with_alpha_channel {
                            new_frame.push(frame[i + 3]);
                        }
                        i += 4;
                    }
                    frames.push(new_frame);
                }
            },
            false => {
                for frame in &self.frames {
                    let mut new_frame = Vec::new();
                    let mut i = 0;
                    while i < frame.len() {
                        new_frame.push(frame[i]);
                        new_frame.push(frame[i + 1]);
                        new_frame.push(frame[i + 2]);
                        if with_alpha_channel {
                            new_frame.push(255);
                        }
                        i += 4;
                    }
                    frames.push(new_frame);
                }
            },
        }
        CBTReader { frames }
    }
}

impl PartialEq for CBT {
    fn eq(&self, other: &CBT) -> bool {
        self.header == other.header && self.frames == other.frames
    }
}