use crate::*;

#[derive(Debug)]
pub enum MemoryKind {
    Owned(Rc<Memory>),
    Import(ModuleRef),
    ImportHost(HostRef),
}

#[derive(Debug)]
pub enum TableKind {
    Owned((Rc<[TableElement]>, TableType)),
    Import(ModuleRef),
    ImportHost(HostRef),
}

#[derive(Clone, Copy, Default, Debug)]
pub enum TableElement {
    #[default]
    Uninitialized,
    /// A function in a host module
    Host { module: HostModuleRef, index: u16 },
    /// A WebAssembly function
    Func { module: ModuleRef, index: u16 },
}

#[derive(Debug)]
pub struct Module {
    pub name: String,
    pub types: Box<[FuncType]>,
    pub functions: Box<[Func]>,
    pub table: Option<TableKind>,
    pub memory: Option<MemoryKind>,
    pub globals: Box<[Global]>,
    pub imports: Box<[Import]>,
    pub exports: Box<[Export]>,
    pub start: Option<Ref>,
}

pub trait CustomSectionHandler {
    /// Called when we reach a custom section of the Wasm binary.
    /// This reader must read _exactly_ `size` bytes out from the reader.
    /// If the reader does not follow this rule, a validation error will be triggered.
    fn custom_section(
        &mut self,
        name: &str,
        size: usize,
        reader: &mut Reader,
    ) -> Result<(), ValidationError>;
}

struct DefaultCustomSectionHandler;
impl CustomSectionHandler for DefaultCustomSectionHandler {
    fn custom_section(
        &mut self,
        name: &str,
        size: usize,
        reader: &mut Reader,
    ) -> Result<(), ValidationError> {
        let _ = name;
        reader.skip(size)
    }
}

impl Module {
    /// Decode, validate, compile, and instantiate a module from `stream` in a
    /// single pass.
    ///
    /// # Errors invalidate the engine
    ///
    /// Unlike the Wasm specification, instantiation is not atomic. After an
    /// `Err` (or a failed [`Engine::push_module`]) the embedder must drop the
    /// entire [`Engine`] and its [`CodeBuilder`] and reconstruct them before
    /// loading or invoking anything else. [`Engine::reset`] only resets
    /// execution state and is not a recovery path for a failed load.
    ///
    /// - **No rollback of imported state.** Data and element segments are
    ///   written into their target memory or table as they are read, rather
    ///   than buffered until the whole module validates. When the target is an
    ///   imported memory or table, a rejected module can leave writes in
    ///   another module's memory or table, including table entries that
    ///   reference the module index the rejected module would have received.
    /// - **No rollback of allocations.** Since deallocation cannot precede
    ///   allocation, freeing the partial state of a rejected module poisons the
    ///   page allocator. Every later allocation fails with
    ///   [`AllocError::HeapPoisoned`] until all pages are released. Code pages
    ///   consumed by the rejected module are not returned to the
    ///   `max_code_pages` budget.
    pub fn new<const MAX_CONTROL_FRAMES: usize, const MAX_STACK_DEPTH: usize>(
        name: &str,
        stream: &mut dyn WasmStream,
        store: &mut Store,
        code_builder: &mut CodeBuilder,
        allocator: Rc<dyn WasmMemoryAllocator>,
    ) -> Result<Module, ParseError> {
        let mut wasm = Reader::new(stream);

        Module::read::<MAX_CONTROL_FRAMES, MAX_STACK_DEPTH>(
            name,
            &mut wasm,
            store,
            code_builder,
            &mut DefaultCustomSectionHandler,
            allocator,
        )
        .map_err(|err| ParseError {
            offset: wasm.offset() as u32,
            err,
        })
    }

    #[allow(clippy::too_many_arguments)]
    fn read<const MAX_CONTROL_FRAMES: usize, const MAX_STACK_DEPTH: usize>(
        name: &str,
        wasm: &mut Reader,
        store: &mut Store,
        code_builder: &mut CodeBuilder,
        custom_handler: &mut dyn CustomSectionHandler,
        allocator: Rc<dyn WasmMemoryAllocator>,
    ) -> Result<Module, SectionDecodeError> {
        let magic = wasm.strip_bytes::<4>()?;
        if magic != [0x00, 0x61, 0x73, 0x6D] {
            return Err(ValidationError::MalformedMagic.into());
        }

        let version = wasm.strip_bytes::<4>()?;

        if version != [0x01, 0x00, 0x00, 0x00] {
            return Err(ValidationError::MalformedVersion.into());
        }

        // Make sure that the module name is not a duplicate in the store
        // Allow multiple modules with empty names since they can't be imported
        if !name.is_empty() {
            if store.modules().iter().any(|m| m.name == name) {
                return Err(ValidationError::DuplicateModuleName.into());
            }

            if store.host_modules().iter().any(|m| m.name == name) {
                return Err(ValidationError::DuplicateModuleName.into());
            }
        }

        let mut module = Module {
            name: name.try_into()?,
            types: Vec::zero().into(),
            functions: Vec::zero().into(),
            table: None,
            memory: None,
            globals: Vec::zero().into(),
            imports: Vec::zero().into(),
            exports: Vec::zero().into(),
            start: None,
        };

        let mut last_section: SectionKind = SectionKind::Custom;
        let mut seen_code_section = false;

        loop {
            let section_ty = match SectionKind::read(wasm) {
                Ok(section) => section,
                Err(ValidationError::Eof) => break,
                Err(e) => return Err(e.into()),
            };

            // Validate the section ordering
            // Custom sections can be interspersed as needed
            if section_ty != SectionKind::Custom {
                if last_section > section_ty && last_section != SectionKind::Custom {
                    return Err(
                        ValidationError::InvalidSectionOrdering(last_section, section_ty).into(),
                    );
                } else if last_section == section_ty {
                    return Err(ValidationError::DuplicateSection(section_ty).into());
                }

                last_section = section_ty;
            }

            if section_ty == SectionKind::Code {
                seen_code_section = true;
            }

            let section_size = wasm.read_u32()? as usize;
            let section_start = wasm.offset();

            wasm.with_limit(section_size, |wasm| {
                module
                    .read_section::<MAX_CONTROL_FRAMES, MAX_STACK_DEPTH>(
                        wasm,
                        store,
                        section_size,
                        section_ty,
                        custom_handler,
                        code_builder,
                        allocator.clone(),
                    )
                    .map_err(|e| e.with_section(section_ty))
            })?;

            // Validate we actually read the entire section
            let section_end = wasm.offset();
            let section_length = section_end - section_start;
            if section_length != section_size {
                return Err(ValidationError::MalformedSectionSize.with_section(section_ty));
            }
        }

        // If we have a function section we should also have a code section and vis versa
        if !seen_code_section && !module.functions.is_empty() {
            return Err(ValidationError::InvalidCodeSectionFunctionCount.into());
        }

        Ok(module)
    }

    #[allow(clippy::too_many_arguments)]
    fn read_section<const MAX_CONTROL_FRAMES: usize, const MAX_STACK_DEPTH: usize>(
        &mut self,
        wasm: &mut Reader,
        store: &mut Store,
        section_size: usize,
        section_ty: SectionKind,
        custom_handler: &mut dyn CustomSectionHandler,
        code_builder: &mut CodeBuilder,
        allocator: Rc<dyn WasmMemoryAllocator>,
    ) -> Result<(), ValidationError> {
        use SectionKind::*;
        match section_ty {
            Custom => {
                CustomSection::read(wasm, section_size, custom_handler)?;
            }
            Type => {
                self.types = TypeSection::read(wasm)?.into();
            }
            Import => {
                self.imports = ImportSection::read(wasm, store, self)?.into();

                // We need to resolve some imports into the module (i.e. memory and table)
                for import in &self.imports {
                    match import {
                        crate::Import::Memory { module, .. } => {
                            if self.memory.is_some() {
                                return Err(ValidationError::MultipleMemories);
                            }

                            // The import should have already validated with linkage
                            // We can make the assertion here
                            let Some(MemoryKind::Owned(_)) =
                                &store.modules()[module.0 as usize].memory
                            else {
                                unreachable!()
                            };

                            self.memory = Some(MemoryKind::Import(*module));
                        }
                        crate::Import::HostMemory { module, index } => {
                            if self.memory.is_some() {
                                return Err(ValidationError::MultipleMemories);
                            }

                            // Make sure this module actually defines a memory
                            let _ = &store.host_modules()[module.0 as usize].memory
                                [*index as usize]
                                .value;

                            self.memory = Some(MemoryKind::ImportHost(HostRef {
                                module: *module,
                                index: *index,
                            }))
                        }
                        crate::Import::Table { module, .. } => {
                            if self.table.is_some() {
                                return Err(ValidationError::MultipleTables);
                            }

                            let Some(TableKind::Owned(_)) =
                                &mut store.modules_mut()[module.0 as usize].table
                            else {
                                unreachable!()
                            };

                            self.table = Some(TableKind::Import(*module));
                        }
                        crate::Import::HostTable { module, index } => {
                            if self.table.is_some() {
                                return Err(ValidationError::MultipleTables);
                            }

                            // Make sure this module actually defines a table
                            let _ = &store.host_modules()[module.0 as usize].table[*index as usize]
                                .value;

                            self.table = Some(TableKind::ImportHost(HostRef {
                                module: *module,
                                index: *index,
                            }));
                        }
                        _ => {}
                    }
                }
            }
            Function => {
                self.functions = FunctionSection::read(wasm, self)?.into();
            }
            Table => {
                if self.table.is_some() {
                    return Err(ValidationError::MultipleTables);
                }

                self.table = TableSection::read(wasm)?;
            }
            Memory => {
                self.memory = MemorySection::read(wasm, self, allocator)?;
            }
            Global => {
                self.globals = GlobalSection::read(wasm, store, self)?.into();
            }
            Export => {
                self.exports = ExportSection::read(wasm, self)?.into();
            }
            Start => {
                let idx = FuncIdx::read(wasm)?;
                let r = self
                    .get_func_ref(idx)
                    .ok_or(ValidationError::FunctionIdxOutOfRange)?;

                // Start functions must be [] -> []
                match r {
                    Ref::Module(index) => {
                        let f = &self.functions[index as usize];
                        let ty = &self.types[f.ty.0 as usize];
                        if !ty.params.is_empty() || !ty.returns.is_empty() {
                            return Err(ValidationError::InvalidStartFunctionSignature);
                        }
                    }
                    Ref::Host { .. } => {
                        // We explicitely do not support start mapped to host functions
                        // It decreases the state space. Just write a wrapper function!
                        return Err(ValidationError::InvalidHostStartFunction);
                    }
                    Ref::Extern { module, index } => {
                        let m = &store.modules()[module.0 as usize];
                        let f = &m.functions[index as usize];
                        let ty = &m.types[f.ty.0 as usize];
                        if !ty.params.is_empty() || !ty.returns.is_empty() {
                            return Err(ValidationError::InvalidStartFunctionSignature);
                        }
                    }
                }

                self.start = Some(r);
            }
            Element => {
                ElementSection::read(wasm, store, self)?;
            }
            Code => {
                CodeSection::read::<MAX_CONTROL_FRAMES, MAX_STACK_DEPTH>(
                    wasm,
                    code_builder,
                    store,
                    self,
                )?;
            }
            Data => {
                DataSection::read(wasm, store, self)?;
            }
            _ => unreachable!(),
        }

        Ok(())
    }

    pub fn get_func_ref(&self, x: FuncIdx) -> Option<Ref> {
        // Check if this function index is a host function or an internal function
        let mut n = 0;
        for f in self.imports.iter() {
            match f {
                Import::Func { module, index } => {
                    if x.0 == n {
                        // We are at the proper index, this is our function
                        // This import has already been resolved to an embedded function
                        return Some(Ref::Extern {
                            module: *module,
                            index: *index,
                        });
                    }

                    n += 1;
                }
                Import::HostFunc { module, index } => {
                    if x.0 == n {
                        return Some(Ref::Host {
                            module: *module,
                            index: *index,
                        });
                    }

                    n += 1;
                }
                _ => {}
            }
        }

        // 'n' is the number of imported functions which offset the local function index
        let i = (x.0 - n) as usize;

        if i >= self.functions.len() {
            None
        } else {
            Some(Ref::Module(i as u16))
        }
    }

    pub fn get_global_ref(&self, x: GlobalIdx) -> Option<Ref> {
        // Check if this global index is a host function or an internal global
        let mut n = 0;
        for f in self.imports.iter() {
            match f {
                Import::Global { module, index } => {
                    if x.0 == n {
                        // We are at the proper index, this is our global
                        // This import has already been resolved to an embedded global
                        return Some(Ref::Extern {
                            module: *module,
                            index: *index,
                        });
                    }

                    n += 1;
                }
                Import::HostGlobal { module, index } => {
                    if x.0 == n {
                        return Some(Ref::Host {
                            module: *module,
                            index: *index,
                        });
                    }

                    n += 1;
                }
                _ => {}
            }
        }

        // 'n' is the number of imported globals which offset the internal global index
        let i = (x.0 - n) as usize;

        if i >= self.globals.len() {
            None
        } else {
            Some(Ref::Module(i as u16))
        }
    }
}

/// All Wasm sections ordered by the order expected in the file
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub enum SectionKind {
    Custom,
    Type,
    Import,
    Function,
    Table,
    Memory,
    Global,
    Export,
    Start,
    Element,
    Code,
    Data,

    N,
}

impl SectionKind {
    pub fn convert(value: u8) -> Result<SectionKind, ValidationError> {
        use SectionKind::*;
        let ty = match value {
            0 => Custom,
            1 => Type,
            2 => Import,
            3 => Function,
            4 => Table,
            5 => Memory,
            6 => Global,
            7 => Export,
            8 => Start,
            9 => Element,
            10 => Code,
            11 => Data,
            other => return Err(ValidationError::MalformedSectionId(other)),
        };

        Ok(ty)
    }

    pub fn read(wasm: &mut Reader) -> Result<Self, ValidationError> {
        SectionKind::convert(wasm.read_u8()?)
    }
}

pub struct CustomSection;

impl CustomSection {
    pub const MAX_NAME_LENGTH: usize = 32;

    pub fn read(
        wasm: &mut Reader,
        size: usize,
        handler: &mut dyn CustomSectionHandler,
    ) -> Result<(), ValidationError> {
        let start = wasm.offset();

        // We only support up to MAX_NAME_LENGTH-byte names
        let name: StaticVec<u8, { CustomSection::MAX_NAME_LENGTH }> =
            wasm.read_vec_stack(|w| w.read_u8())?;
        let name_str = core::str::from_utf8(&name).map_err(|_| ValidationError::MalformedUtf8)?;

        let name_length = wasm.offset() - start;
        if name_length > size {
            return Err(ValidationError::MalformedSectionSize);
        }

        let payload_len = size - name_length;
        wasm.with_limit(payload_len, |wasm| {
            handler.custom_section(name_str, payload_len, wasm)
        })
    }
}

pub struct TypeSection;

impl TypeSection {
    pub fn read(wasm: &mut Reader) -> Result<Vec<FuncType>, ValidationError> {
        wasm.read_vec(FuncType::read)
    }
}

macro_rules! read_impl_u32 {
    ($type_name:ident) => {
        #[derive(Debug, Clone, Copy, PartialEq, Eq)]
        pub struct $type_name(pub u32);
        impl $type_name {
            pub fn read(wasm: &mut Reader) -> Result<Self, ValidationError> {
                Ok($type_name(wasm.read_u32()?))
            }
        }

        impl From<u32> for $type_name {
            fn from(v: u32) -> Self {
                $type_name(v)
            }
        }
    };
}

read_impl_u32!(TypeIdx);
read_impl_u32!(FuncIdx);
read_impl_u32!(TableIdx);
read_impl_u32!(MemIdx);
read_impl_u32!(GlobalIdx);
read_impl_u32!(LocalIdx);
read_impl_u32!(LabelIdx);

pub enum ImportDesc {
    Func(TypeIdx),
    Table(TableType),
    Mem(MemType),
    Global(GlobalType),
}

impl ImportDesc {
    pub fn read(wasm: &mut Reader) -> Result<Self, ValidationError> {
        match wasm.read_u8()? {
            0x00 => Ok(ImportDesc::Func(TypeIdx::read(wasm)?)),
            0x01 => Ok(ImportDesc::Table(TableType::read(wasm)?)),
            0x02 => Ok(ImportDesc::Mem(MemType::read(wasm)?)),
            0x03 => Ok(ImportDesc::Global(GlobalType::read(wasm)?)),
            c => Err(ValidationError::MalformedImportExportDesc(c)),
        }
    }
}

pub struct ImportSection;

impl ImportSection {
    pub fn read(
        wasm: &mut Reader,
        store: &Store,
        module: &Module,
    ) -> Result<Vec<Import>, ValidationError> {
        wasm.read_vec(|w| Import::read(w, module, store))
    }
}

impl Func {
    pub fn read_func_section(wasm: &mut Reader, module: &Module) -> Result<Func, ValidationError> {
        let ty_idx = TypeIdx::read(wasm)?;

        let ty = module
            .types
            .get(ty_idx.0 as usize)
            .ok_or(ValidationError::TypeIdxOutOfRange)?;

        let parameter_size = ty.params.iter().fold(0, |sum, a_ty| sum + a_ty.size()) / 4;

        let return_ty = if ty.returns.len() > 1 {
            return Err(ValidationError::FunctionReturnsTooLarge);
        } else if ty.returns.is_empty() {
            None
        } else {
            Some(ty.returns[0])
        };

        if parameter_size > 0xFF {
            return Err(ValidationError::FunctionParametersTooLarge);
        }

        Ok(Func {
            ty: ty_idx,
            stack_usage: 0,
            local_size: 0,
            parameter_size: parameter_size as u8,
            return_ty,
            locals: Vec::zero().into(),
            expr: Expr::zero(),
        })
    }
}

struct FunctionSection;
impl FunctionSection {
    pub fn read(wasm: &mut Reader, module: &Module) -> Result<Vec<Func>, ValidationError> {
        let functions = wasm.read_vec(|w| Func::read_func_section(w, module))?;

        // Module functions are referenced by a 16-bit index (`Ref::Module`)
        if functions.len() > u16::MAX as usize {
            return Err(ValidationError::IdxTooLarge);
        }

        Ok(functions)
    }
}

pub struct TableSection;

impl TableSection {
    pub fn read(wasm: &mut Reader) -> Result<Option<TableKind>, ValidationError> {
        let n = wasm.read_u32()?;
        if n == 0 {
            Ok(None)
        } else if n == 1 {
            let table_type = TableType::read(wasm)?;
            let table = Rc::new_slice_with_default(table_type.limits.min as usize)?;

            Ok(Some(TableKind::Owned((table, table_type))))
        } else {
            Err(ValidationError::MultipleTables)
        }
    }
}

pub struct MemorySection;
impl MemorySection {
    pub fn read(
        wasm: &mut Reader,
        module: &Module,
        allocator: Rc<dyn WasmMemoryAllocator>,
    ) -> Result<Option<MemoryKind>, ValidationError> {
        let len = wasm.read_u32()?;
        if len > 1 {
            Err(ValidationError::MultipleMemories)
        } else if len == 0 {
            Ok(None)
        } else if module.memory.is_some() {
            Err(ValidationError::MultipleMemories)
        } else {
            // We are allocating memory for this module
            let ty = MemType::read(wasm)?;
            let memory = Memory::new(ty, allocator)
                .map_err(|_| ValidationError::GuestMemoryAllocationFailure)?;
            Ok(Some(MemoryKind::Owned(Rc::new(memory)?)))
        }
    }
}

#[derive(Debug, Clone)]
pub struct Global {
    pub type_: GlobalType,
    pub value: RawValue,
}

impl Global {
    pub fn value(&self) -> Value {
        self.value.to_value(self.type_.ty)
    }

    pub fn read(
        wasm: &mut Reader,
        store: &Store,
        module: &Module,
    ) -> Result<Self, ValidationError> {
        let type_ = GlobalType::read(wasm)?;
        let init = Expr::read_constant(wasm, store, module)?;
        let init = match init {
            Value::I32(i) => {
                if type_.ty != ValType::I32 {
                    return Err(ValidationError::GlobalTypeMismatch);
                }

                RawValue::from_i32(i)
            }
            Value::I64(i) => {
                if type_.ty != ValType::I64 {
                    return Err(ValidationError::GlobalTypeMismatch);
                }

                RawValue::from_i64(i)
            }
            Value::F32(z) => {
                if type_.ty != ValType::F32 {
                    return Err(ValidationError::GlobalTypeMismatch);
                }

                RawValue::from_f32(z)
            }
            Value::F64(z) => {
                if type_.ty != ValType::F64 {
                    return Err(ValidationError::GlobalTypeMismatch);
                }

                RawValue::from_f64(z)
            }
        };

        Ok(Global { type_, value: init })
    }
}

pub struct GlobalSection;
impl GlobalSection {
    pub fn read(
        wasm: &mut Reader,
        store: &Store,
        module: &Module,
    ) -> Result<Vec<Global>, ValidationError> {
        let globals = wasm.read_vec(|wasm| Global::read(wasm, store, module))?;

        // Module globals are referenced by a 16-bit index (`Ref::Module`)
        if globals.len() > u16::MAX as usize {
            return Err(ValidationError::IdxTooLarge);
        }

        Ok(globals)
    }
}

#[derive(Debug, Clone)]
pub enum ExportDesc {
    Func(FuncIdx),
    Table(TableIdx),
    Mem(MemIdx),
    Global(GlobalIdx),
}

impl ExportDesc {
    pub fn read(wasm: &mut Reader) -> Result<Self, ValidationError> {
        match wasm.read_u8()? {
            0x00 => Ok(ExportDesc::Func(FuncIdx::read(wasm)?)),
            0x01 => Ok(ExportDesc::Table(TableIdx::read(wasm)?)),
            0x02 => Ok(ExportDesc::Mem(MemIdx::read(wasm)?)),
            0x03 => Ok(ExportDesc::Global(GlobalIdx::read(wasm)?)),
            c => Err(ValidationError::MalformedImportExportDesc(c)),
        }
    }
}

#[derive(Debug, Clone)]
pub struct Export {
    pub name: String,
    pub desc: ExportDesc,
}

impl Export {
    pub fn read(wasm: &mut Reader, module: &Module) -> Result<Self, ValidationError> {
        let name = Name::read(wasm)?;
        let desc = ExportDesc::read(wasm)?;

        // Check if the export refers to a defined symbol
        match desc {
            ExportDesc::Func(i) => {
                let _ = module
                    .get_func_ref(i)
                    .ok_or(ValidationError::FunctionIdxOutOfRange)?;
            }
            ExportDesc::Table(i) => {
                if i.0 != 0 {
                    return Err(ValidationError::InvalidTableIndex);
                } else if module.table.is_none() {
                    return Err(ValidationError::TableNotDefined);
                }
            }
            ExportDesc::Mem(i) => {
                if i.0 != 0 {
                    return Err(ValidationError::InvalidMemIndex);
                } else if module.memory.is_none() {
                    return Err(ValidationError::MemoryNotDefined);
                }
            }
            ExportDesc::Global(i) => {
                let _ = module
                    .get_global_ref(i)
                    .ok_or(ValidationError::GlobalIdxOutOfRange)?;
            }
        }

        Ok(Export { name, desc })
    }
}

pub struct ExportSection;
impl ExportSection {
    pub fn read(wasm: &mut Reader, module: &Module) -> Result<Vec<Export>, ValidationError> {
        let len = wasm.read_u32()?;
        let mut out: Vec<Export> = Vec::new(len as usize)?;
        for _ in 0..len {
            let e = Export::read(wasm, module)?;
            // Check for duplicate export name
            if out.iter().any(|ei| *ei.name == *e.name) {
                return Err(ValidationError::DuplicateExportName);
            }

            out.push(e);
        }

        Ok(out)
    }
}

impl Module {
    pub(crate) fn get_table<'store>(
        &'store mut self,
        store: &'store mut Store,
    ) -> Result<Option<&'store mut [TableElement]>, ValidationError> {
        if let Some(table) = &mut self.table {
            let table = match table {
                TableKind::Owned(table) => &mut table.0,
                TableKind::Import(i) => {
                    let Some(TableKind::Owned(table)) =
                        &mut store.modules_mut()[i.0 as usize].table
                    else {
                        unreachable!()
                    };

                    &mut table.0
                }
                TableKind::ImportHost(i) => {
                    &mut store.host_modules_mut()[i.module.0 as usize].table[i.index as usize]
                        .value
                        .0
                }
            };

            // The table may still be aliased (e.g. an imported table kept alive
            // by the engine's `Rc` clone, or a host-held table `Rc`). Reject
            // such a module instead of panicking on the failed unique borrow.
            match table.get_mut() {
                Some(t) => Ok(Some(t)),
                None => Err(ValidationError::TableRefNotUnique),
            }
        } else {
            Ok(None)
        }
    }
}

pub struct Element;

impl Element {
    pub fn read(
        wasm: &mut Reader,
        store: &mut Store,
        module: &mut Module,
    ) -> Result<(), ValidationError> {
        let table = TableIdx::read(wasm)?;
        if table.0 != 0 {
            return Err(ValidationError::InvalidTableIndex);
        }

        let offset = match Expr::read_constant(wasm, store, module)? {
            Value::I32(i) => i as u32,
            _ => return Err(ValidationError::InvalidElementOffset),
        };

        let init_len = wasm.read_u32()?;
        if let Some(table) = module.get_table(store)? {
            if offset as usize > table.len() {
                return Err(ValidationError::InvalidElementOffset);
            } else if (offset as u64 + init_len as u64) > table.len() as u64 {
                return Err(ValidationError::InvalidElementOutOfBounds);
            }
        } else {
            return Err(ValidationError::TableNotDefined);
        }

        // Write the function indexes into the table
        for i in 0..init_len as usize {
            let idx = FuncIdx::read(wasm)?;

            let r = module
                .get_func_ref(idx)
                .ok_or(ValidationError::FunctionIdxOutOfRange)?;

            let tr = match r {
                Ref::Module(index) => TableElement::Func {
                    module: ModuleRef(store.modules().len() as u8),
                    index,
                },
                Ref::Host { module, index } => TableElement::Host { module, index },
                Ref::Extern { module, index } => TableElement::Func { module, index },
            };

            let index = (offset as usize) + i;
            if let Some(table) = module.get_table(store)? {
                if index >= table.len() {
                    return Err(ValidationError::InvalidElementOutOfBounds);
                }

                table[index] = tr;
            } else {
                return Err(ValidationError::TableNotDefined);
            }
        }

        Ok(())
    }
}

pub struct ElementSection;
impl ElementSection {
    pub fn read(
        wasm: &mut Reader,
        store: &mut Store,
        module: &mut Module,
    ) -> Result<(), ValidationError> {
        let len = wasm.read_u32()?;
        for _ in 0..len {
            Element::read(wasm, store, module)?;
        }

        Ok(())
    }
}

pub struct CodeSection;

impl CodeSection {
    pub fn read<const MAX_CONTROL_FRAMES: usize, const MAX_STACK_DEPTH: usize>(
        wasm: &mut Reader,
        builder: &mut CodeBuilder,
        store: &Store,
        module: &mut Module,
    ) -> Result<(), ValidationError> {
        let n = wasm.read_u32()?;
        if n as usize != module.functions.len() {
            return Err(ValidationError::InvalidCodeSectionFunctionCount);
        }

        for i in 0..n as usize {
            module.read_function_code::<MAX_CONTROL_FRAMES, MAX_STACK_DEPTH>(
                wasm, store, builder, i,
            )?;
        }

        Ok(())
    }
}

#[derive(Clone)]
pub struct Data;

impl Module {
    pub(crate) fn check_memory_defined(&self) -> Result<(), ValidationError> {
        if self.memory.is_none() {
            Err(ValidationError::MemoryNotDefined)
        } else {
            Ok(())
        }
    }
}

impl Data {
    pub fn read(
        wasm: &mut Reader,
        store: &Store,
        module: &Module,
        memory: &Rc<Memory>,
    ) -> Result<(), ValidationError> {
        let mem = MemIdx::read(wasm)?;
        if mem.0 != 0 {
            return Err(ValidationError::InvalidMemIndex);
        }

        // Make sure we actually defined a linear memory for this module
        module.check_memory_defined()?;

        let offset = match Expr::read_constant(wasm, store, module)? {
            Value::I32(i) => i as u32,
            _ => return Err(ValidationError::InvalidMemOffsetType),
        };

        // Read the data onto the linear memory
        // Note! This is a deviation from the Wasm specification, technically we
        // should be reading into a separate buffer and only populating the linear memory
        // _after_ the full module validates. This is sound because any load failure
        // invalidates the engine (see `Module::new`).
        let len = wasm.read_u32()?;
        for i in 0..len {
            memory.store_u8(Memory::effective_address(offset, i)?, wasm.read_u8()?)?;
        }

        if len == 0 {
            // Make sure the data offset is in-range
            memory.store(offset as usize, &[])?;
        }

        Ok(())
    }
}

pub struct DataSection;
impl DataSection {
    pub fn read(wasm: &mut Reader, store: &Store, module: &Module) -> Result<(), ValidationError> {
        let len = wasm.read_u32()?;
        if len == 0 {
            return Ok(());
        }

        if let Some(memory) = &module.memory {
            let memory = match memory {
                MemoryKind::Owned(memory) => memory,
                MemoryKind::Import(i) => {
                    let Some(MemoryKind::Owned(memory)) = &store.modules()[i.0 as usize].memory
                    else {
                        unreachable!()
                    };

                    memory
                }
                MemoryKind::ImportHost(i) => {
                    &store.host_modules()[i.module.0 as usize].memory[i.index as usize].value
                }
            };

            for _ in 0..len {
                Data::read(wasm, store, module, memory)?;
            }

            Ok(())
        } else {
            Err(ValidationError::MemoryNotDefined)
        }
    }
}