#[cfg(not(features = "std"))]
use crate::std::collections::BTreeSet as Set;
use crate::std::{mem, vec::Vec};

use crate::symbols::{expand_symbols, push_code_symbols, resolve_function, Symbol};
use log::trace;
use parity_wasm::elements;

#[derive(Debug)]
pub enum Error {
    /// Since optimizer starts with export entries, export
    ///   section is supposed to exist.
    NoExportSection,
}

pub fn optimize(
    module: &mut elements::Module, // Module to optimize
    used_exports: Vec<&str>,       // List of only exports that will be usable after optimization
) -> Result<(), Error> {
    // WebAssembly exports optimizer
    // Motivation: emscripten compiler backend compiles in many unused exports
    //   which in turn compile in unused imports and leaves unused functions

    // try to parse name section
    let module_temp = mem::take(module);
    let module_temp = module_temp
        .parse_names()
        .unwrap_or_else(|(_err, module)| module);
    *module = module_temp;

    // Algo starts from the top, listing all items that should stay
    let mut stay: Set<_> = module
        .export_section()
        .ok_or(Error::NoExportSection)?
        .entries()
        .iter()
        .enumerate()
        .filter_map(|(index, entry)| {
            if used_exports.iter().any(|e| *e == entry.field()) {
                Some(Symbol::Export(index))
            } else {
                None
            }
        })
        .collect();

    // If there is start function in module, it should start
    if let Some(ss) = module.start_section() {
        stay.insert(resolve_function(module, ss));
    }

    // All symbols used in data/element segments are also should be preserved
    let mut init_symbols = Vec::new();
    if let Some(data_section) = module.data_section() {
        for segment in data_section.entries() {
            push_code_symbols(
                module,
                segment
                    .offset()
                    .as_ref()
                    .expect("parity-wasm is compiled without bulk-memory operations")
                    .code(),
                &mut init_symbols,
            );
        }
    }
    if let Some(elements_section) = module.elements_section() {
        for segment in elements_section.entries() {
            push_code_symbols(
                module,
                segment
                    .offset()
                    .as_ref()
                    .expect("parity-wasm is compiled without bulk-memory operations")
                    .code(),
                &mut init_symbols,
            );
            for func_index in segment.members() {
                stay.insert(resolve_function(module, *func_index));
            }
        }
    }

    stay.extend(init_symbols.drain(..));

    // Call function which will traverse the list recursively, filling stay with all symbols
    // that are already used by those which already there
    expand_symbols(module, &mut stay);

    for symbol in stay.iter() {
        trace!("symbol to stay: {:?}", symbol);
    }

    // Keep track of referreable symbols to rewire calls/globals
    let mut eliminated_funcs = Vec::new();
    let mut eliminated_globals = Vec::new();
    let mut eliminated_types = Vec::new();

    // First, iterate through types
    let mut index = 0;
    let mut old_index = 0;

    loop {
        if type_section(module)
            .map(|section| section.types_mut().len())
            .unwrap_or(0)
            == index
        {
            break;
        }

        if stay.contains(&Symbol::Type(old_index)) {
            index += 1;
        } else {
            type_section(module)
                    .expect("If type section does not exists, the loop will break at the beginning of first iteration")
                    .types_mut().remove(index);
            eliminated_types.push(old_index);
            trace!("Eliminated type({})", old_index);
        }
        old_index += 1;
    }

    // Second, iterate through imports
    let mut top_funcs = 0;
    let mut top_globals = 0;
    index = 0;
    old_index = 0;

    if let Some(imports) = import_section(module) {
        loop {
            let mut remove = false;
            match imports.entries()[index].external() {
                elements::External::Function(_) => {
                    if stay.contains(&Symbol::Import(old_index)) {
                        index += 1;
                    } else {
                        remove = true;
                        eliminated_funcs.push(top_funcs);
                        trace!(
                            "Eliminated import({}) func({}, {})",
                            old_index,
                            top_funcs,
                            imports.entries()[index].field()
                        );
                    }
                    top_funcs += 1;
                }
                elements::External::Global(_) => {
                    if stay.contains(&Symbol::Import(old_index)) {
                        index += 1;
                    } else {
                        remove = true;
                        eliminated_globals.push(top_globals);
                        trace!(
                            "Eliminated import({}) global({}, {})",
                            old_index,
                            top_globals,
                            imports.entries()[index].field()
                        );
                    }
                    top_globals += 1;
                }
                _ => {
                    index += 1;
                }
            }
            if remove {
                imports.entries_mut().remove(index);
            }

            old_index += 1;

            if index == imports.entries().len() {
                break;
            }
        }
    }

    // Third, iterate through globals
    if let Some(globals) = global_section(module) {
        index = 0;
        old_index = 0;

        loop {
            if globals.entries_mut().len() == index {
                break;
            }
            if stay.contains(&Symbol::Global(old_index)) {
                index += 1;
            } else {
                globals.entries_mut().remove(index);
                eliminated_globals.push(top_globals + old_index);
                trace!("Eliminated global({})", top_globals + old_index);
            }
            old_index += 1;
        }
    }

    // Forth, delete orphaned functions
    if function_section(module).is_some() && code_section(module).is_some() {
        index = 0;
        old_index = 0;

        loop {
            if function_section(module)
                .expect("Functons section to exist")
                .entries_mut()
                .len()
                == index
            {
                break;
            }
            if stay.contains(&Symbol::Function(old_index)) {
                index += 1;
            } else {
                function_section(module)
                    .expect("Functons section to exist")
                    .entries_mut()
                    .remove(index);
                code_section(module)
                    .expect("Code section to exist")
                    .bodies_mut()
                    .remove(index);

                eliminated_funcs.push(top_funcs + old_index);
                trace!("Eliminated function({})", top_funcs + old_index);
            }
            old_index += 1;
        }
    }

    // Fifth, eliminate unused exports
    {
        let exports = export_section(module).ok_or(Error::NoExportSection)?;

        index = 0;
        old_index = 0;

        loop {
            if exports.entries_mut().len() == index {
                break;
            }
            if stay.contains(&Symbol::Export(old_index)) {
                index += 1;
            } else {
                trace!(
                    "Eliminated export({}, {})",
                    old_index,
                    exports.entries_mut()[index].field()
                );
                exports.entries_mut().remove(index);
            }
            old_index += 1;
        }
    }

    if !eliminated_globals.is_empty()
        || !eliminated_funcs.is_empty()
        || !eliminated_types.is_empty()
    {
        // Finaly, rewire all calls, globals references and types to the new indices
        //   (only if there is anything to do)
        // When sorting primitives sorting unstable is faster without any difference in result.
        eliminated_globals.sort_unstable();
        eliminated_funcs.sort_unstable();
        eliminated_types.sort_unstable();

        for section in module.sections_mut() {
            match section {
                elements::Section::Start(func_index) if !eliminated_funcs.is_empty() => {
                    let totalle = eliminated_funcs
                        .iter()
                        .take_while(|i| (**i as u32) < *func_index)
                        .count();
                    *func_index -= totalle as u32;
                }
                elements::Section::Function(function_section) if !eliminated_types.is_empty() => {
                    for func_signature in function_section.entries_mut() {
                        let totalle = eliminated_types
                            .iter()
                            .take_while(|i| (**i as u32) < func_signature.type_ref())
                            .count();
                        *func_signature.type_ref_mut() -= totalle as u32;
                    }
                }
                elements::Section::Import(import_section) if !eliminated_types.is_empty() => {
                    for import_entry in import_section.entries_mut() {
                        if let elements::External::Function(type_ref) = import_entry.external_mut()
                        {
                            let totalle = eliminated_types
                                .iter()
                                .take_while(|i| (**i as u32) < *type_ref)
                                .count();
                            *type_ref -= totalle as u32;
                        }
                    }
                }
                elements::Section::Code(code_section)
                    if !eliminated_globals.is_empty() || !eliminated_funcs.is_empty() =>
                {
                    for func_body in code_section.bodies_mut() {
                        if !eliminated_funcs.is_empty() {
                            update_call_index(func_body.code_mut(), &eliminated_funcs);
                        }
                        if !eliminated_globals.is_empty() {
                            update_global_index(
                                func_body.code_mut().elements_mut(),
                                &eliminated_globals,
                            )
                        }
                        if !eliminated_types.is_empty() {
                            update_type_index(func_body.code_mut(), &eliminated_types)
                        }
                    }
                }
                elements::Section::Export(export_section) => {
                    for export in export_section.entries_mut() {
                        match export.internal_mut() {
                            elements::Internal::Function(func_index) => {
                                let totalle = eliminated_funcs
                                    .iter()
                                    .take_while(|i| (**i as u32) < *func_index)
                                    .count();
                                *func_index -= totalle as u32;
                            }
                            elements::Internal::Global(global_index) => {
                                let totalle = eliminated_globals
                                    .iter()
                                    .take_while(|i| (**i as u32) < *global_index)
                                    .count();
                                *global_index -= totalle as u32;
                            }
                            _ => {}
                        }
                    }
                }
                elements::Section::Global(global_section) => {
                    for global_entry in global_section.entries_mut() {
                        update_global_index(
                            global_entry.init_expr_mut().code_mut(),
                            &eliminated_globals,
                        )
                    }
                }
                elements::Section::Data(data_section) => {
                    for segment in data_section.entries_mut() {
                        update_global_index(
                            segment
                                .offset_mut()
                                .as_mut()
                                .expect("parity-wasm is compiled without bulk-memory operations")
                                .code_mut(),
                            &eliminated_globals,
                        )
                    }
                }
                elements::Section::Element(elements_section) => {
                    for segment in elements_section.entries_mut() {
                        update_global_index(
                            segment
                                .offset_mut()
                                .as_mut()
                                .expect("parity-wasm is compiled without bulk-memory operations")
                                .code_mut(),
                            &eliminated_globals,
                        );
                        // update all indirect call addresses initial values
                        for func_index in segment.members_mut() {
                            let totalle = eliminated_funcs
                                .iter()
                                .take_while(|i| (**i as u32) < *func_index)
                                .count();
                            *func_index -= totalle as u32;
                        }
                    }
                }
                elements::Section::Name(name_section) => {
                    if let Some(func_name) = name_section.functions_mut() {
                        let mut func_name_map = mem::take(func_name.names_mut());
                        for index in &eliminated_funcs {
                            func_name_map.remove(*index as u32);
                        }
                        let updated_map = func_name_map
                            .into_iter()
                            .map(|(index, value)| {
                                let totalle = eliminated_funcs
                                    .iter()
                                    .take_while(|i| (**i as u32) < index)
                                    .count() as u32;
                                (index - totalle, value)
                            })
                            .collect();
                        *func_name.names_mut() = updated_map;
                    }

                    if let Some(local_name) = name_section.locals_mut() {
                        let mut local_names_map = mem::take(local_name.local_names_mut());
                        for index in &eliminated_funcs {
                            local_names_map.remove(*index as u32);
                        }
                        let updated_map = local_names_map
                            .into_iter()
                            .map(|(index, value)| {
                                let totalle = eliminated_funcs
                                    .iter()
                                    .take_while(|i| (**i as u32) < index)
                                    .count() as u32;
                                (index - totalle, value)
                            })
                            .collect();
                        *local_name.local_names_mut() = updated_map;
                    }
                }
                _ => {}
            }
        }
    }

    // Also drop all custom sections
    module
        .sections_mut()
        .retain(|section| !matches!(section, elements::Section::Custom(_)));

    Ok(())
}

pub fn update_call_index(instructions: &mut elements::Instructions, eliminated_indices: &[usize]) {
    use parity_wasm::elements::Instruction::*;
    for instruction in instructions.elements_mut().iter_mut() {
        if let Call(call_index) = instruction {
            let totalle = eliminated_indices
                .iter()
                .take_while(|i| (**i as u32) < *call_index)
                .count();
            trace!(
                "rewired call {} -> call {}",
                *call_index,
                *call_index - totalle as u32
            );
            *call_index -= totalle as u32;
        }
    }
}

/// Updates global references considering the _ordered_ list of eliminated indices
pub fn update_global_index(
    instructions: &mut [elements::Instruction],
    eliminated_indices: &[usize],
) {
    use parity_wasm::elements::Instruction::*;
    for instruction in instructions.iter_mut() {
        match instruction {
            GetGlobal(index) | SetGlobal(index) => {
                let totalle = eliminated_indices
                    .iter()
                    .take_while(|i| (**i as u32) < *index)
                    .count();
                trace!(
                    "rewired global {} -> global {}",
                    *index,
                    *index - totalle as u32
                );
                *index -= totalle as u32;
            }
            _ => {}
        }
    }
}

/// Updates global references considering the _ordered_ list of eliminated indices
pub fn update_type_index(instructions: &mut elements::Instructions, eliminated_indices: &[usize]) {
    use parity_wasm::elements::Instruction::*;
    for instruction in instructions.elements_mut().iter_mut() {
        if let CallIndirect(call_index, _) = instruction {
            let totalle = eliminated_indices
                .iter()
                .take_while(|i| (**i as u32) < *call_index)
                .count();
            trace!(
                "rewired call_indrect {} -> call_indirect {}",
                *call_index,
                *call_index - totalle as u32
            );
            *call_index -= totalle as u32;
        }
    }
}

pub fn import_section(module: &mut elements::Module) -> Option<&mut elements::ImportSection> {
    for section in module.sections_mut() {
        if let elements::Section::Import(sect) = section {
            return Some(sect);
        }
    }
    None
}

pub fn global_section(module: &mut elements::Module) -> Option<&mut elements::GlobalSection> {
    for section in module.sections_mut() {
        if let elements::Section::Global(sect) = section {
            return Some(sect);
        }
    }
    None
}

pub fn function_section(module: &mut elements::Module) -> Option<&mut elements::FunctionSection> {
    for section in module.sections_mut() {
        if let elements::Section::Function(sect) = section {
            return Some(sect);
        }
    }
    None
}

pub fn code_section(module: &mut elements::Module) -> Option<&mut elements::CodeSection> {
    for section in module.sections_mut() {
        if let elements::Section::Code(sect) = section {
            return Some(sect);
        }
    }
    None
}

pub fn export_section(module: &mut elements::Module) -> Option<&mut elements::ExportSection> {
    for section in module.sections_mut() {
        if let elements::Section::Export(sect) = section {
            return Some(sect);
        }
    }
    None
}

pub fn type_section(module: &mut elements::Module) -> Option<&mut elements::TypeSection> {
    for section in module.sections_mut() {
        if let elements::Section::Type(sect) = section {
            return Some(sect);
        }
    }
    None
}

#[cfg(test)]
mod tests {

    use super::*;
    use parity_wasm::{builder, elements};

    /// @spec 0
    /// Optimizer presumes that export section exists and contains
    /// all symbols passed as a second parameter. Since empty module
    /// obviously contains no export section, optimizer should return
    /// error on it.
    #[test]
    fn empty() {
        let mut module = builder::module().build();
        let result = optimize(&mut module, vec!["_call"]);

        assert!(result.is_err());
    }

    /// @spec 1
    /// Imagine the unoptimized module has two own functions, `_call` and `_random`
    /// and exports both of them in the export section. During optimization, the `_random`
    /// function should vanish completely, given we pass `_call` as the only function to stay
    /// in the module.
    #[test]
    fn minimal() {
        let mut module = builder::module()
            .function()
            .signature()
            .param()
            .i32()
            .build()
            .build()
            .function()
            .signature()
            .param()
            .i32()
            .param()
            .i32()
            .build()
            .build()
            .export()
            .field("_call")
            .internal()
            .func(0)
            .build()
            .export()
            .field("_random")
            .internal()
            .func(1)
            .build()
            .build();
        assert_eq!(
            module
                .export_section()
                .expect("export section to be generated")
                .entries()
                .len(),
            2
        );

        optimize(&mut module, vec!["_call"]).expect("optimizer to succeed");

        assert_eq!(
            1,
            module
                .export_section()
                .expect("export section to be generated")
                .entries()
                .len(),
            "There should only 1 (one) export entry in the optimized module"
        );

        assert_eq!(
            1,
            module
                .function_section()
                .expect("functions section to be generated")
                .entries()
                .len(),
            "There should 2 (two) functions in the optimized module"
        );
    }

    /// @spec 2
    /// Imagine there is one exported function in unoptimized module, `_call`, that we specify as the one
    /// to stay during the optimization. The code of this function uses global during the execution.
    /// This sayed global should survive the optimization.
    #[test]
    fn globals() {
        let mut module = builder::module()
            .global()
            .value_type()
            .i32()
            .build()
            .function()
            .signature()
            .param()
            .i32()
            .build()
            .body()
            .with_instructions(elements::Instructions::new(vec![
                elements::Instruction::GetGlobal(0),
                elements::Instruction::End,
            ]))
            .build()
            .build()
            .export()
            .field("_call")
            .internal()
            .func(0)
            .build()
            .build();

        optimize(&mut module, vec!["_call"]).expect("optimizer to succeed");

        assert_eq!(
            1,
            module.global_section().expect("global section to be generated").entries().len(),
            "There should 1 (one) global entry in the optimized module, since _call function uses it"
        );
    }

    /// @spec 2
    /// Imagine there is one exported function in unoptimized module, `_call`, that we specify as the one
    /// to stay during the optimization. The code of this function uses one global during the execution,
    /// but we have a bunch of other unused globals in the code. Last globals should not survive the optimization,
    /// while the former should.
    #[test]
    fn globals_2() {
        let mut module = builder::module()
            .global()
            .value_type()
            .i32()
            .build()
            .global()
            .value_type()
            .i64()
            .build()
            .global()
            .value_type()
            .f32()
            .build()
            .function()
            .signature()
            .param()
            .i32()
            .build()
            .body()
            .with_instructions(elements::Instructions::new(vec![
                elements::Instruction::GetGlobal(1),
                elements::Instruction::End,
            ]))
            .build()
            .build()
            .export()
            .field("_call")
            .internal()
            .func(0)
            .build()
            .build();

        optimize(&mut module, vec!["_call"]).expect("optimizer to succeed");

        assert_eq!(
            1,
            module.global_section().expect("global section to be generated").entries().len(),
            "There should 1 (one) global entry in the optimized module, since _call function uses only one"
        );
    }

    /// @spec 3
    /// Imagine the unoptimized module has two own functions, `_call` and `_random`
    /// and exports both of them in the export section. Function `_call` also calls `_random`
    /// in its function body. The optimization should kick `_random` function from the export section
    /// but preserve it's body.
    #[test]
    fn call_ref() {
        let mut module = builder::module()
            .function()
            .signature()
            .param()
            .i32()
            .build()
            .body()
            .with_instructions(elements::Instructions::new(vec![
                elements::Instruction::Call(1),
                elements::Instruction::End,
            ]))
            .build()
            .build()
            .function()
            .signature()
            .param()
            .i32()
            .param()
            .i32()
            .build()
            .build()
            .export()
            .field("_call")
            .internal()
            .func(0)
            .build()
            .export()
            .field("_random")
            .internal()
            .func(1)
            .build()
            .build();
        assert_eq!(
            module
                .export_section()
                .expect("export section to be generated")
                .entries()
                .len(),
            2
        );

        optimize(&mut module, vec!["_call"]).expect("optimizer to succeed");

        assert_eq!(
            1,
            module
                .export_section()
                .expect("export section to be generated")
                .entries()
                .len(),
            "There should only 1 (one) export entry in the optimized module"
        );

        assert_eq!(
            2,
            module
                .function_section()
                .expect("functions section to be generated")
                .entries()
                .len(),
            "There should 2 (two) functions in the optimized module"
        );
    }

    /// @spec 4
    /// Imagine the unoptimized module has an indirect call to function of type 1
    /// The type should persist so that indirect call would work
    #[test]
    fn call_indirect() {
        let mut module = builder::module()
            .function()
            .signature()
            .param()
            .i32()
            .param()
            .i32()
            .build()
            .build()
            .function()
            .signature()
            .param()
            .i32()
            .param()
            .i32()
            .param()
            .i32()
            .build()
            .build()
            .function()
            .signature()
            .param()
            .i32()
            .build()
            .body()
            .with_instructions(elements::Instructions::new(vec![
                elements::Instruction::CallIndirect(1, 0),
                elements::Instruction::End,
            ]))
            .build()
            .build()
            .export()
            .field("_call")
            .internal()
            .func(2)
            .build()
            .build();

        optimize(&mut module, vec!["_call"]).expect("optimizer to succeed");

        assert_eq!(
            2,
            module
                .type_section()
                .expect("type section to be generated")
                .types()
                .len(),
            "There should 2 (two) types left in the module, 1 for indirect call and one for _call"
        );

        let indirect_opcode = &module
            .code_section()
            .expect("code section to be generated")
            .bodies()[0]
            .code()
            .elements()[0];
        match *indirect_opcode {
            elements::Instruction::CallIndirect(0, 0) => {}
            _ => {
                panic!(
                    "Expected call_indirect to use index 0 after optimization, since previois 0th was eliminated, but got {:?}",
                    indirect_opcode
                );
            }
        }
    }
}