/*!
Plugins are WebAssembly modules that can be attached to an [`Environment`][0]. They have 2 tasks:
- Modify other modules that are added to the [`Environment`][0]
- Spawn environment bound processes (TODO)

Each module that is being added to an [`Environment`][0] is going to get parsed into a
`ModuleContext` and Plugins will have an opportunity to modify this `ModuleContext` before it's JIT
compiled. Theses modifications can include adding new types, functions and exports or wrapping
existing functions with additional code (TODO).

One common use case for Plugins is the addition of utility functions that can't be expressed in
some languages. For example `lunatic_call_indirect(table_index: u32)` is used inside Rust code
compiled to WebAssembly, but it can't be generated by the stable Rust compiler.

Another use case is inserting additional code into particular functions to collect data during
execution, for example as the `heap_profiler` plugin does.

[0]: crate::environment::Environment
*/

use std::{collections::HashMap, convert::TryInto};

use anyhow::{anyhow, Result};
use wasmparser::{
    Chunk, Export, FunctionBody, Import, Name, NameSectionReader, Parser, Payload, Range,
    SectionReader, TypeDef, Validator,
};
use wasmtime::{Linker, Module, Store};

use crate::{api, environment::PLUGIN_ENV, state::PluginState};

/// WebAssembly module that can be attached to an [`Environment`](crate::environment::Environment).
#[derive(Clone)]
pub struct Plugin {
    module: Module,
}

impl Plugin {
    /// Creates a new [`Plugin`] from raw WebAssembly data.
    pub fn new(module: Vec<u8>) -> Result<Self> {
        let module = Module::new(&PLUGIN_ENV.engine, module)?;
        Ok(Self { module })
    }
}

// [`ModuleContext`] is part of the plugin process' state.
//
// It is shared between all the plugins and allows them to perform modifications to a module that
// is being added to the [`Environment`].
//
// Having such an abstraction allows us to perform the module parsing once and skip
// sections that we are not interested in modifying.
pub(crate) struct ModuleContext<'a> {
    module: &'a [u8],
    types: Vec<ContextType<'a>>,
    functions: Vec<u32>,
    function_names: HashMap<&'a str, u32>,
    code_section: Vec<ContextCode<'a>>,
    imports: Vec<Import<'a>>,
    exports: Vec<ContextExport<'a>>,
    sections: Vec<wasm_encoder::RawSection<'a>>,
}

impl<'a> ModuleContext<'a> {
    // Creates a new [`ModuleContext`] from raw WebAssembly data.
    pub(crate) fn new(module: &'a [u8]) -> Result<Self> {
        let mut context = Self {
            module,
            types: Vec::new(),
            functions: Vec::new(),
            function_names: HashMap::new(),
            code_section: Vec::new(),
            imports: Vec::new(),
            exports: Vec::new(),
            sections: Vec::new(),
        };

        context.parse()?;

        Ok(context)
    }

    // Adds a raw function to the end of the function & code section and returns the function
    // index.
    pub(crate) fn add_function(
        &mut self,
        type_index: u32,
        func_locals: &[u8],
        func_body: Vec<u8>,
    ) -> Result<u64> {
        self.functions.push(type_index as u32);
        // Add code section
        let func_locals = func_locals
            .chunks_exact(5) // (count: u32, type :u8)
            .map(|data| {
                let count = u32::from_le_bytes(data[0..4].try_into()?);
                Ok((count, Self::u8_to_valtype(data[4])?))
            })
            .collect::<Result<Vec<_>>>()?;
        self.code_section
            .push(ContextCode::New(func_locals, func_body));
        // Return index of inserted function
        Ok((self.functions.len() - 1) as u64)
    }

    // Adds a function type to the end off the types section and returns the type index.
    //
    // Types are represented as `u8` values, following the WebAssembly binary format.
    pub(crate) fn add_function_type(
        &mut self,
        param_types: &[u8],
        return_types: &[u8],
    ) -> Result<u64> {
        let param_types = param_types
            .iter()
            .map(|value| Self::u8_to_valtype(*value))
            .collect::<Result<Vec<_>>>()?;
        let return_types = return_types
            .iter()
            .map(|value| Self::u8_to_valtype(*value))
            .collect::<Result<Vec<_>>>()?;
        self.types.push(ContextType::New(param_types, return_types));
        Ok((self.types.len() - 1) as u64)
    }

    // Adds a function as an export.
    pub(crate) fn add_function_export(&mut self, name: String, id: u32) {
        self.exports.push(ContextExport::NewFunction(name, id));
    }

    // Adds a raw section to module without parsing it.
    fn add_section(&mut self, id: wasm_encoder::SectionId, range: Range) {
        self.sections.push(wasm_encoder::RawSection {
            id: id as u8,
            data: &self.module[range.start..range.end],
        });
    }

    // Parses the binary module and generates structured information so that we can easily add
    // types, functions and exports. At the same time it validates the module.
    fn parse(&mut self) -> Result<()> {
        let mut parser = Parser::new(0);
        let mut validator = Validator::new();

        let mut data = self.module;
        loop {
            let payload = match parser.parse(data, true)? {
                Chunk::NeedMoreData(_) => unreachable!(),
                Chunk::Parsed { payload, consumed } => {
                    data = &data[consumed..];
                    payload
                }
            };

            match payload {
                Payload::Version { num, range } => validator.version(num, &range)?,
                Payload::TypeSection(types) => {
                    validator.type_section(&types)?;
                    self.types.reserve(types.get_count() as usize);
                    for ty in types {
                        self.types.push(ContextType::Unchanged(ty?));
                    }
                }
                Payload::ImportSection(imports) => {
                    validator.import_section(&imports)?;
                    self.imports.reserve(imports.get_count() as usize);
                    for import in imports {
                        self.imports.push(import?);
                    }
                }
                Payload::AliasSection(aliases) => {
                    validator.alias_section(&aliases)?;
                    self.add_section(wasm_encoder::SectionId::Alias, aliases.range());
                }
                Payload::InstanceSection(instances) => {
                    validator.instance_section(&instances)?;
                    self.add_section(wasm_encoder::SectionId::Instance, instances.range());
                }
                Payload::FunctionSection(functions) => {
                    validator.function_section(&functions)?;
                    self.functions.reserve(functions.get_count() as usize);
                    for f in functions {
                        self.functions.push(f?);
                    }
                }
                Payload::TableSection(tables) => {
                    validator.table_section(&tables)?;
                    self.add_section(wasm_encoder::SectionId::Table, tables.range());
                }
                Payload::MemorySection(memories) => {
                    validator.memory_section(&memories)?;
                    self.add_section(wasm_encoder::SectionId::Memory, memories.range());
                }
                Payload::TagSection(tags) => {
                    validator.tag_section(&tags)?;
                    self.add_section(wasm_encoder::SectionId::Tag, tags.range());
                }
                Payload::GlobalSection(globals) => {
                    validator.global_section(&globals)?;
                    self.add_section(wasm_encoder::SectionId::Global, globals.range());
                }
                Payload::ExportSection(exports) => {
                    validator.export_section(&exports)?;
                    self.exports.reserve(exports.get_count() as usize);
                    for export in exports {
                        self.exports.push(ContextExport::Unchanged(export?));
                    }
                }
                Payload::StartSection { func, range } => {
                    validator.start_section(func, &range)?;
                    self.add_section(wasm_encoder::SectionId::Start, range);
                }
                Payload::ElementSection(elements) => {
                    validator.element_section(&elements)?;
                    self.add_section(wasm_encoder::SectionId::Element, elements.range());
                }
                Payload::DataCountSection { count, range } => {
                    validator.data_count_section(count, &range)?;
                    self.add_section(wasm_encoder::SectionId::DataCount, range)
                }
                Payload::DataSection(data) => {
                    validator.data_section(&data)?;
                    self.add_section(wasm_encoder::SectionId::Data, data.range());
                }
                Payload::CodeSectionStart { count, range, .. } => {
                    validator.code_section_start(count, &range)?;
                }
                Payload::CodeSectionEntry(body) => {
                    let mut validator = validator.code_section_entry()?;
                    validator.validate(&body)?;
                    self.code_section.push(ContextCode::Unchanged(body));
                }
                Payload::ModuleSectionStart { count, range, size } => {
                    validator.module_section_start(count, &range)?;
                    self.add_section(wasm_encoder::SectionId::Module, range);
                    // Skip individual module parsing
                    parser.skip_section();
                    data = &data[size as usize..];
                }
                // Should never be called, because we skip it in `Payload::ModuleSectionStart`
                Payload::ModuleSectionEntry { .. } => unreachable!(),
                Payload::CustomSection {
                    name, range, data, ..
                } => {
                    // The name of the custom name section is "name".
                    if name == "name" {
                        // Extract the function names
                        let mut name_section_reader = NameSectionReader::new(data, 0)?;
                        while let Ok(subsection) = name_section_reader.read() {
                            if let Name::Function(function_names) = subsection {
                                let mut name_map = function_names.get_map()?;
                                let name_map_count = name_map.get_count() as usize;
                                for _ in 0..name_map_count {
                                    if let Ok(function) = name_map.read() {
                                        self.function_names.insert(function.name, function.index);
                                    } else {
                                        return Err(anyhow!("Couldn't parse all function names of CustomName section."));
                                    }
                                }
                                // Finish as soon as function names are parsed.
                                break;
                            }
                        }
                    }
                    self.add_section(wasm_encoder::SectionId::Custom, range)
                }
                Payload::UnknownSection { id, range, .. } => {
                    validator.unknown_section(id, &range)?;
                }
                Payload::End => break,
            }
        }

        Ok(())
    }

    // Encodes the module back to the WebAssembly binary format.
    //
    // The order of sections is significant in WebAssembly and we need to make sure that we
    // correctly mix together parsed and raw sections. For this we relay on the validator in the
    // `parse` functions having validated the initial order.
    pub(crate) fn encode(&mut self) -> Result<Vec<u8>> {
        let mut module = wasm_encoder::Module::new();
        let mut section_iterator = self.sections.iter();
        let mut current_section = section_iterator.next();

        if !self.types.is_empty() {
            module.section(&self.encode_types()?);
        }
        if !self.imports.is_empty() {
            module.section(&self.encode_imports()?);
        }
        if !self.functions.is_empty() {
            module.section(&self.encode_functions());
        }
        if let Some(section) = current_section {
            if section.id == wasm_encoder::SectionId::Table as u8 {
                module.section(section);
                current_section = section_iterator.next();
            }
        }
        if let Some(section) = current_section {
            if section.id == wasm_encoder::SectionId::Memory as u8 {
                module.section(section);
                current_section = section_iterator.next();
            }
        }
        if let Some(section) = current_section {
            if section.id == wasm_encoder::SectionId::Global as u8 {
                module.section(section);
                current_section = section_iterator.next();
            }
        }
        if !self.exports.is_empty() {
            module.section(&self.encode_exports()?);
        }
        if let Some(section) = current_section {
            if section.id == wasm_encoder::SectionId::Start as u8 {
                module.section(section);
                current_section = section_iterator.next();
            }
        }
        if let Some(section) = current_section {
            if section.id == wasm_encoder::SectionId::Element as u8 {
                module.section(section);
                current_section = section_iterator.next();
            }
        }
        if !self.code_section.is_empty() {
            module.section(&self.encode_code()?);
        }
        if let Some(section) = current_section {
            if section.id == wasm_encoder::SectionId::Data as u8 {
                module.section(section);
                current_section = section_iterator.next();
            }
        }
        if let Some(section) = current_section {
            if section.id == wasm_encoder::SectionId::DataCount as u8 {
                module.section(section);
            }
        }
        // All other sections
        for section in section_iterator {
            module.section(section);
        }

        Ok(module.finish())
    }

    fn encode_types(&self) -> Result<wasm_encoder::TypeSection> {
        let mut type_section = wasm_encoder::TypeSection::new();
        for context_type in &self.types {
            match context_type {
                ContextType::New(params, returns) => {
                    type_section.function(params.iter().copied(), returns.iter().copied());
                }
                ContextType::Unchanged(type_def) => match type_def {
                    TypeDef::Func(function) => {
                        type_section.function(
                            function
                                .params
                                .iter()
                                .map(Self::type_translate)
                                .collect::<Result<Vec<_>>>()?,
                            function
                                .returns
                                .iter()
                                .map(Self::type_translate)
                                .collect::<Result<Vec<_>>>()?,
                        );
                    }
                    TypeDef::Instance(_) => {
                        return Err(anyhow!("Unsupported Wasm type: 'Instance'"))
                    }
                    TypeDef::Module(_) => return Err(anyhow!("Unsupported Wasm type: 'Module'")),
                },
            };
        }
        Ok(type_section)
    }

    fn encode_functions(&self) -> wasm_encoder::FunctionSection {
        let mut function_section = wasm_encoder::FunctionSection::new();
        for type_index in &self.functions {
            function_section.function(*type_index);
        }
        function_section
    }

    fn encode_code(&self) -> Result<wasm_encoder::CodeSection> {
        let mut code_section = wasm_encoder::CodeSection::new();
        for context_code in &self.code_section {
            match context_code {
                ContextCode::New(locals, instructions) => {
                    let mut function = wasm_encoder::Function::new(locals.iter().copied());
                    function.raw(instructions.iter().copied());
                    code_section.function(&function);
                }
                ContextCode::Unchanged(function_body) => {
                    // Extract locals from original function
                    let mut locals_reader = function_body.get_locals_reader()?;
                    let locls_count = locals_reader.get_count() as usize;
                    let mut locals = Vec::with_capacity(locls_count);
                    for _ in 0..locls_count {
                        if let Ok((count, type_)) = locals_reader.read() {
                            locals.push((count, Self::type_translate(&type_)?));
                        } else {
                            return Err(anyhow!("Couldn't parse all locals of Wasm function."));
                        }
                    }
                    let mut function = wasm_encoder::Function::new(locals);
                    // Copy original instruction body from function
                    let instructions_start =
                        function_body.get_operators_reader()?.original_position();
                    let body_end = function_body.range().end;
                    let instructions = self.module[instructions_start..body_end].iter().copied();
                    function.raw(instructions);
                    code_section.function(&function);
                }
            }
        }
        Ok(code_section)
    }

    fn encode_imports(&self) -> Result<wasm_encoder::ImportSection> {
        let mut import_section = wasm_encoder::ImportSection::new();
        for import in &self.imports {
            import_section.import(
                import.module,
                import.field,
                Self::entity_type_translate(&import.ty)?,
            );
        }
        Ok(import_section)
    }

    fn encode_exports(&self) -> Result<wasm_encoder::ExportSection> {
        let mut export_section = wasm_encoder::ExportSection::new();
        for context_export in &self.exports {
            match context_export {
                ContextExport::NewFunction(name, id) => {
                    export_section.export(name, wasm_encoder::Export::Function(*id))
                }
                ContextExport::Unchanged(export) => export_section.export(
                    export.field,
                    match export.kind {
                        wasmparser::ExternalKind::Function => {
                            wasm_encoder::Export::Function(export.index)
                        }
                        wasmparser::ExternalKind::Table => {
                            wasm_encoder::Export::Table(export.index)
                        }
                        wasmparser::ExternalKind::Memory => {
                            wasm_encoder::Export::Memory(export.index)
                        }
                        wasmparser::ExternalKind::Tag => wasm_encoder::Export::Tag(export.index),
                        wasmparser::ExternalKind::Global => {
                            wasm_encoder::Export::Global(export.index)
                        }
                        wasmparser::ExternalKind::Type => {
                            return Err(anyhow!("Unsupported Wasm export: 'Type'"))
                        }
                        wasmparser::ExternalKind::Module => {
                            wasm_encoder::Export::Module(export.index)
                        }
                        wasmparser::ExternalKind::Instance => {
                            wasm_encoder::Export::Instance(export.index)
                        }
                    },
                ),
            };
        }
        Ok(export_section)
    }

    // Translate WebAssembly raw type value to the `wasm_encoder::ValType` enum.
    fn u8_to_valtype(value: u8) -> Result<wasm_encoder::ValType> {
        let result = match value {
            0x7F => wasm_encoder::ValType::I32,
            0x7E => wasm_encoder::ValType::I64,
            0x7D => wasm_encoder::ValType::F32,
            0x7C => wasm_encoder::ValType::F64,
            0x7B => wasm_encoder::ValType::V128,
            0x70 => wasm_encoder::ValType::FuncRef,
            0x6F => wasm_encoder::ValType::ExternRef,
            _ => return Err(anyhow!(format!("Unsupported Wasm type ID: {}", value))),
        };
        Ok(result)
    }

    // Translate `wasmparser` WebAssembly types to `wasm_encoder` equivalents.
    fn type_translate(parser_type: &wasmparser::Type) -> Result<wasm_encoder::ValType> {
        let type_ = match parser_type {
            wasmparser::Type::I32 => wasm_encoder::ValType::I32,
            wasmparser::Type::I64 => wasm_encoder::ValType::I64,
            wasmparser::Type::F32 => wasm_encoder::ValType::F32,
            wasmparser::Type::F64 => wasm_encoder::ValType::F64,
            wasmparser::Type::V128 => wasm_encoder::ValType::V128,
            wasmparser::Type::FuncRef => wasm_encoder::ValType::FuncRef,
            wasmparser::Type::ExternRef => wasm_encoder::ValType::ExternRef,
            wasmparser::Type::ExnRef => return Err(anyhow!("Unsupported Wasm type: 'ExnRef'")),
            wasmparser::Type::Func => return Err(anyhow!("Unsupported Wasm type: 'Func'")),
            wasmparser::Type::EmptyBlockType => {
                return Err(anyhow!("Unsupported Wasm type: 'EmptyBlockType'"))
            }
        };
        Ok(type_)
    }

    // Translate `wasmparser::ImportSectionEntryType` to `wasm_encoder::EntityType`.
    fn entity_type_translate(
        parser_type: &wasmparser::ImportSectionEntryType,
    ) -> Result<wasm_encoder::EntityType> {
        let type_ = match parser_type {
            wasmparser::ImportSectionEntryType::Function(id) => {
                wasm_encoder::EntityType::Function(*id)
            }
            wasmparser::ImportSectionEntryType::Table(table) => {
                wasm_encoder::EntityType::Table(wasm_encoder::TableType {
                    element_type: Self::type_translate(&table.element_type)?,
                    minimum: table.initial,
                    maximum: table.maximum,
                })
            }
            wasmparser::ImportSectionEntryType::Memory(memory) => {
                wasm_encoder::EntityType::Memory(wasm_encoder::MemoryType {
                    minimum: memory.initial,
                    maximum: memory.maximum,
                    memory64: memory.memory64,
                })
            }
            wasmparser::ImportSectionEntryType::Tag(tag) => {
                wasm_encoder::EntityType::Tag(wasm_encoder::TagType {
                    kind: wasm_encoder::TagKind::Exception,
                    func_type_idx: tag.type_index,
                })
            }
            wasmparser::ImportSectionEntryType::Global(global) => {
                wasm_encoder::EntityType::Global(wasm_encoder::GlobalType {
                    val_type: Self::type_translate(&global.content_type)?,
                    mutable: global.mutable,
                })
            }
            wasmparser::ImportSectionEntryType::Module(id) => wasm_encoder::EntityType::Module(*id),
            wasmparser::ImportSectionEntryType::Instance(id) => {
                wasm_encoder::EntityType::Instance(*id)
            }
        };
        Ok(type_)
    }
}

enum ContextType<'a> {
    // Parameter & return types
    New(Vec<wasm_encoder::ValType>, Vec<wasm_encoder::ValType>),
    Unchanged(TypeDef<'a>),
}

enum ContextCode<'a> {
    // Locals & instruction body
    New(Vec<(u32, wasm_encoder::ValType)>, Vec<u8>),
    Unchanged(FunctionBody<'a>),
}

enum ContextExport<'a> {
    NewFunction(String, u32),
    Unchanged(Export<'a>),
}

// Runs all the plugins on the module, allowing them to modify it.
// Modifications a plugin can do:
// * Adding function types
// * Adding functions
// * Adding imports & exports
// * Wrapping existing host functions in additional instructions
pub(crate) fn patch_module(module: &[u8], plugins: &[Plugin]) -> Result<Vec<u8>> {
    // Return early if no plugins are used
    if plugins.is_empty() {
        return Ok(module.into());
    }

    let mut module_context = ModuleContext::new(module)?;
    let state = PluginState::new(&mut module_context);
    let mut store = Store::new(&PLUGIN_ENV.engine, state);
    let mut linker = Linker::new(&PLUGIN_ENV.engine);
    // Add plugin related host functions to the linker
    api::plugin::register(&mut linker)?;

    for plugin in plugins {
        let instance = linker.instantiate(&mut store, &plugin.module)?;
        // Call the initialize method on the plugin if it exists
        if let Some(initialize) = instance.get_func(&mut store, "_initialize") {
            initialize.call(&mut store, &[], &mut [])?;
        }
        if let Some(hook) = instance.get_func(&mut store, "lunatic_create_module_hook") {
            hook.call(&mut store, &[], &mut [])?;
        };
    }
    store.into_data().module_context().encode()
}

#[cfg(test)]
mod tests {
    use std::{fs, path::Path};

    use super::*;
    use pretty_assertions::assert_eq;

    #[test]
    fn module_context_parses_and_encodes_empty_module() {
        let input = "(module)";
        let input_wasm = wat::parse_str(input).unwrap();
        // Normalize input wasm
        let input = wabt::wasm2wat(&input_wasm).unwrap();
        let output_wasm = patch_module(&input_wasm, &[]).unwrap();
        let output = wabt::wasm2wat(&output_wasm).unwrap();
        assert_eq!(input, output);
    }

    #[test]
    fn module_context_parses_and_encodes_simple_module() {
        let input = r#"
        (module
            (import "env" "mem" (memory 1))
            (func $a (import "env" "imported_func") (param i32))
            (func $b (param $lhs i32) (param $rhs i32) (result i32)
                local.get $lhs
                local.get $rhs
                i32.add
            )
            (func $data_types
                (local $int_32 i32)
                (local $int_64 i64)
                (local $float_32 f32)
                (local $float_64 f64)

                (local.set $int_32 (i32.const 16))
                (local.set $int_64 (i64.const 128))
                (local.set $float_32 (f32.const 3.14))
                (local.set $float_64 (f64.const 1.28))
            )
            (func (export "exported_func")
              i32.const 42
              call $a
            )
            (table 2 funcref)
            (data (i32.const 0) "Hi")
            (elem (i32.const 0) $a $b)
          )
        "#;
        let input_wasm = wat::parse_str(input).unwrap();
        // Normalize input wasm
        let input = wabt::wasm2wat(&input_wasm).unwrap();
        let output_wasm = patch_module(&input_wasm, &[]).unwrap();
        let output = wabt::wasm2wat(&output_wasm).unwrap();
        assert_eq!(input, output);
    }

    // #[test]
    fn _simple_add_function_plugin() {
        let input = r#"
        (module)
        "#;
        let expected_output = r#"
        (module
            (type (;0;) (func))
            (func (param i32)
                local.get 0
                call_indirect (type 0) 0
            )
            (export "lunatic_call_indirect" (func 0))
        )
        "#;
        let input_wasm = wat::parse_str(input).unwrap();
        let expected_output_wasm = wat::parse_str(expected_output).unwrap();
        // Normalize expected output
        let expected_output = wabt::wasm2wat(&expected_output_wasm).unwrap();

        let path = Path::new("target/wasm/stdlib.wasm");
        let module = fs::read(path).unwrap();
        let plugin = Plugin::new(module).unwrap();
        let output_wasm = patch_module(&input_wasm, &[plugin]).unwrap();
        let output = wabt::wasm2wat(&output_wasm).unwrap();
        assert_eq!(expected_output, output);
    }
}