use crate::*;
use brec_consts::*;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};

impl Write for Block {
    fn generate(&self) -> Result<TokenStream, E> {
        let block_name = self.name();
        let mut buf_fillers = Vec::new();
        for field in self.fields.iter().filter(|f| !f.injected) {
            let size = field.size();
            match field.ty {
                BlockTy::U8
                | BlockTy::U16
                | BlockTy::U32
                | BlockTy::U64
                | BlockTy::U128
                | BlockTy::I8
                | BlockTy::I16
                | BlockTy::I32
                | BlockTy::I64
                | BlockTy::I128
                | BlockTy::F32
                | BlockTy::F64
                | BlockTy::Bool
                | BlockTy::LinkedToU8(..) => {
                    let as_bytes = field.to_bytes(true)?;
                    buf_fillers.push(quote! {
                        buffer[offset..offset + #size].copy_from_slice(#as_bytes);
                        offset += #size;
                    });
                }
                BlockTy::Blob(..) => {
                    let name = format_ident!("{}", field.name);
                    buf_fillers.push(quote! {
                        unsafe {
                            let dst = buffer.as_mut_ptr().add(offset);
                            let src = self.#name.as_ptr();
                            std::ptr::copy_nonoverlapping(src, dst, #size);
                        }
                        offset += #size;
                    });
                }
            }
        }
        let const_sig = self.const_sig_name();
        let size = self.size(if cfg!(feature = "resilient") {
            BLOCK_SIZE_FIELD_LEN
        } else {
            0
        });
        let write_len = if cfg!(feature = "resilient") {
            let len = self
                .fields
                .iter()
                .filter(|f| !f.injected)
                .map(|f| f.size())
                .sum::<usize>() as u32;
            quote! {
                buffer[offset..offset + #BLOCK_SIZE_FIELD_LEN]
                    .copy_from_slice(&#len.to_le_bytes());
                offset += #BLOCK_SIZE_FIELD_LEN;
            }
        } else {
            quote! {}
        };
        Ok(quote! {

            impl brec::WriteTo for #block_name {

                fn write<T: std::io::Write>(&self, writer: &mut T) -> std::io::Result<usize> {
                    use brec::prelude::*;
                    let mut buffer = [0u8; #size];
                    let mut offset = 0;
                    let crc = self.crc();
                    buffer[offset..offset + #BLOCK_SIG_LEN].copy_from_slice(&#const_sig);
                    offset += #BLOCK_SIG_LEN;
                    #write_len
                    #(#buf_fillers)*
                    unsafe {
                        let dst = buffer.as_mut_ptr().add(offset);
                        let src = crc.as_ptr();
                        std::ptr::copy_nonoverlapping(src, dst, #BLOCK_CRC_LEN);
                    }
                    writer.write(&buffer)
                }

                fn write_all<T: std::io::Write>(&self, writer: &mut T) -> std::io::Result<()> {
                    use brec::prelude::*;
                    let mut buffer = [0u8; #size];
                    let mut offset = 0;
                    let crc = self.crc();
                    buffer[offset..offset + #BLOCK_SIG_LEN].copy_from_slice(&#const_sig);
                    offset += #BLOCK_SIG_LEN;
                    #write_len
                    #(#buf_fillers)*
                    unsafe {
                        let dst = buffer.as_mut_ptr().add(offset);
                        let src = crc.as_ptr();
                        std::ptr::copy_nonoverlapping(src, dst, #BLOCK_CRC_LEN);
                    }
                    writer.write_all(&buffer)
                }

            }

        })
    }
}

impl WriteVectored for Block {
    fn generate(&self) -> Result<TokenStream, E> {
        let block_name = self.name();
        let mut groups = Vec::new();
        let mut slices = Vec::new();
        let mut fields = Vec::new();
        for field in self.fields.iter().filter(|f| !f.injected) {
            match field.ty {
                BlockTy::U8
                | BlockTy::U16
                | BlockTy::U32
                | BlockTy::U64
                | BlockTy::U128
                | BlockTy::I8
                | BlockTy::I16
                | BlockTy::I32
                | BlockTy::I64
                | BlockTy::I128
                | BlockTy::F32
                | BlockTy::F64
                | BlockTy::Bool
                | BlockTy::LinkedToU8(..) => {
                    if !slices.is_empty() {
                        groups.push(slices);
                        slices = Vec::new();
                    }
                    fields.push(field);
                }
                BlockTy::Blob(..) => {
                    if !fields.is_empty() {
                        groups.push(fields);
                        fields = Vec::new();
                    }
                    slices.push(field);
                }
            }
        }
        if !fields.is_empty() {
            groups.push(fields);
        }
        if !slices.is_empty() {
            groups.push(slices);
        }
        let mut fragments = Vec::new();
        for group in groups.into_iter() {
            let len: usize = group.iter().map(|f| f.size()).sum();
            let last = group.len() - 1;
            for (n, field) in group.into_iter().enumerate() {
                let size = field.size();
                match field.ty {
                    BlockTy::U8
                    | BlockTy::U16
                    | BlockTy::U32
                    | BlockTy::U64
                    | BlockTy::U128
                    | BlockTy::I8
                    | BlockTy::I16
                    | BlockTy::I32
                    | BlockTy::I64
                    | BlockTy::I128
                    | BlockTy::F32
                    | BlockTy::F64
                    | BlockTy::Bool
                    | BlockTy::LinkedToU8(..) => {
                        if n == 0 {
                            fragments.push(quote! {
                                let mut buffer = [0u8; #len];
                                let mut offset = 0;
                            });
                        }
                        let as_bytes = field.to_bytes(true)?;
                        fragments.push(quote! {
                            buffer[offset..offset + #size].copy_from_slice(#as_bytes);
                        });
                        if n != last {
                            fragments.push(quote! {
                                offset += #size;
                            });
                        } else {
                            fragments.push(quote! {
                                slices.add_buffered(buffer.to_vec());
                            });
                        }
                    }
                    BlockTy::Blob(..) => {
                        let name = format_ident!("{}", field.name);
                        fragments.push(quote! {
                            slices.add_slice(&self.#name);
                        });
                    }
                }
            }
        }

        let const_sig = self.const_sig_name();

        let write_len = if cfg!(feature = "resilient") {
            let len = self
                .fields
                .iter()
                .filter(|f| !f.injected)
                .map(|f| f.size())
                .sum::<usize>() as u32;
            quote! {
                slices.add_buffered(#len.to_le_bytes().to_vec());
            }
        } else {
            quote! {}
        };
        Ok(quote! {

            impl brec::WriteVectoredTo for #block_name {

                fn slices(&self) -> std::io::Result<brec::IoSlices<'_>> {
                    use brec::prelude::*;
                    let mut slices = brec::IoSlices::default();
                    slices.add_buffered(#const_sig.to_vec());
                    #write_len;
                    #(#fragments)*
                    slices.add_buffered(self.crc().to_vec());
                    Ok(slices)
                }

            }

        })
    }
}
