#![cfg(test)]

use std::collections::HashMap;
use std::fs::File;

use wabt::script::{self, Action, Command, CommandKind, ScriptParser, Value};
use wasmi::memory_units::Pages;
use wasmi::{
    Error as InterpreterError, Externals, FuncInstance, FuncRef, GlobalDescriptor, GlobalInstance,
    GlobalRef, ImportResolver, ImportsBuilder, MemoryDescriptor, MemoryInstance, MemoryRef, Module,
    ModuleImportResolver, ModuleInstance, ModuleRef, RuntimeArgs, RuntimeValue, Signature,
    TableDescriptor, TableInstance, TableRef, Trap,
};

fn spec_to_runtime_value(val: Value<u32, u64>) -> RuntimeValue {
    match val {
        Value::I32(v) => RuntimeValue::I32(v),
        Value::I64(v) => RuntimeValue::I64(v),
        Value::F32(v) => RuntimeValue::F32(v.into()),
        Value::F64(v) => RuntimeValue::F64(v.into()),
    }
}

#[derive(Debug)]
enum Error {
    Load(String),
    Start(Trap),
    Script(script::Error),
    Interpreter(InterpreterError),
}

impl From<InterpreterError> for Error {
    fn from(e: InterpreterError) -> Error {
        Error::Interpreter(e)
    }
}

impl From<script::Error> for Error {
    fn from(e: script::Error) -> Error {
        Error::Script(e)
    }
}

struct SpecModule {
    table: TableRef,
    memory: MemoryRef,
    global_i32: GlobalRef,
    global_f32: GlobalRef,
    global_f64: GlobalRef,
}

impl SpecModule {
    fn new() -> Self {
        SpecModule {
            table: TableInstance::alloc(10, Some(20)).unwrap(),
            memory: MemoryInstance::alloc(Pages(1), Some(Pages(2))).unwrap(),
            global_i32: GlobalInstance::alloc(RuntimeValue::I32(666), false),
            global_f32: GlobalInstance::alloc(RuntimeValue::F32(666.0.into()), false),
            global_f64: GlobalInstance::alloc(RuntimeValue::F64(666.0.into()), false),
        }
    }
}

const PRINT_FUNC_INDEX: usize = 0;

impl Externals for SpecModule {
    fn invoke_index(
        &mut self,
        index: usize,
        args: RuntimeArgs,
    ) -> Result<Option<RuntimeValue>, Trap> {
        match index {
            PRINT_FUNC_INDEX => {
                println!("print: {:?}", args);
                Ok(None)
            }
            _ => panic!("SpecModule doesn't provide function at index {}", index),
        }
    }
}

impl ModuleImportResolver for SpecModule {
    fn resolve_func(
        &self,
        field_name: &str,
        func_type: &Signature,
    ) -> Result<FuncRef, InterpreterError> {
        let index = match field_name {
            "print" => PRINT_FUNC_INDEX,
            "print_i32" => PRINT_FUNC_INDEX,
            "print_i32_f32" => PRINT_FUNC_INDEX,
            "print_f64_f64" => PRINT_FUNC_INDEX,
            "print_f32" => PRINT_FUNC_INDEX,
            "print_f64" => PRINT_FUNC_INDEX,
            _ => {
                return Err(InterpreterError::Instantiation(format!(
                    "Unknown host func import {}",
                    field_name
                )));
            }
        };

        if func_type.return_type().is_some() {
            return Err(InterpreterError::Instantiation(
                "Function `print_` have unit return type".into(),
            ));
        }

        let func = FuncInstance::alloc_host(func_type.clone(), index);
        return Ok(func);
    }

    fn resolve_global(
        &self,
        field_name: &str,
        _global_type: &GlobalDescriptor,
    ) -> Result<GlobalRef, InterpreterError> {
        match field_name {
            "global_i32" => Ok(self.global_i32.clone()),
            "global_f32" => Ok(self.global_f32.clone()),
            "global_f64" => Ok(self.global_f64.clone()),
            _ => Err(InterpreterError::Instantiation(format!(
                "Unknown host global import {}",
                field_name
            ))),
        }
    }

    fn resolve_memory(
        &self,
        field_name: &str,
        _memory_type: &MemoryDescriptor,
    ) -> Result<MemoryRef, InterpreterError> {
        if field_name == "memory" {
            return Ok(self.memory.clone());
        }

        Err(InterpreterError::Instantiation(format!(
            "Unknown host memory import {}",
            field_name
        )))
    }

    fn resolve_table(
        &self,
        field_name: &str,
        _table_type: &TableDescriptor,
    ) -> Result<TableRef, InterpreterError> {
        if field_name == "table" {
            return Ok(self.table.clone());
        }

        Err(InterpreterError::Instantiation(format!(
            "Unknown host table import {}",
            field_name
        )))
    }
}

struct SpecDriver {
    spec_module: SpecModule,
    instances: HashMap<String, ModuleRef>,
    last_module: Option<ModuleRef>,
}

impl SpecDriver {
    fn new() -> SpecDriver {
        SpecDriver {
            spec_module: SpecModule::new(),
            instances: HashMap::new(),
            last_module: None,
        }
    }

    fn spec_module(&mut self) -> &mut SpecModule {
        &mut self.spec_module
    }

    fn add_module(&mut self, name: Option<String>, module: ModuleRef) {
        self.last_module = Some(module.clone());
        if let Some(name) = name {
            self.instances.insert(name, module);
        }
    }

    fn module(&self, name: &str) -> Result<ModuleRef, InterpreterError> {
        self.instances.get(name).cloned().ok_or_else(|| {
            InterpreterError::Instantiation(format!("Module not registered {}", name))
        })
    }

    fn module_or_last(&self, name: Option<&str>) -> Result<ModuleRef, InterpreterError> {
        match name {
            Some(name) => self.module(name),
            None => self
                .last_module
                .clone()
                .ok_or_else(|| InterpreterError::Instantiation("No modules registered".into())),
        }
    }
}

impl ImportResolver for SpecDriver {
    fn resolve_func(
        &self,
        module_name: &str,
        field_name: &str,
        func_type: &Signature,
    ) -> Result<FuncRef, InterpreterError> {
        if module_name == "spectest" {
            self.spec_module.resolve_func(field_name, func_type)
        } else {
            self.module(module_name)?
                .resolve_func(field_name, func_type)
        }
    }

    fn resolve_global(
        &self,
        module_name: &str,
        field_name: &str,
        global_type: &GlobalDescriptor,
    ) -> Result<GlobalRef, InterpreterError> {
        if module_name == "spectest" {
            self.spec_module.resolve_global(field_name, global_type)
        } else {
            self.module(module_name)?
                .resolve_global(field_name, global_type)
        }
    }

    fn resolve_memory(
        &self,
        module_name: &str,
        field_name: &str,
        memory_type: &MemoryDescriptor,
    ) -> Result<MemoryRef, InterpreterError> {
        if module_name == "spectest" {
            self.spec_module.resolve_memory(field_name, memory_type)
        } else {
            self.module(module_name)?
                .resolve_memory(field_name, memory_type)
        }
    }

    fn resolve_table(
        &self,
        module_name: &str,
        field_name: &str,
        table_type: &TableDescriptor,
    ) -> Result<TableRef, InterpreterError> {
        if module_name == "spectest" {
            self.spec_module.resolve_table(field_name, table_type)
        } else {
            self.module(module_name)?
                .resolve_table(field_name, table_type)
        }
    }
}

fn try_load_module(wasm: &[u8]) -> Result<Module, Error> {
    Module::from_buffer(wasm).map_err(|e| Error::Load(e.to_string()))
}

fn try_load(wasm: &[u8], spec_driver: &mut SpecDriver) -> Result<(), Error> {
    let module = try_load_module(wasm)?;
    let instance = ModuleInstance::new(&module, &ImportsBuilder::default())?;
    instance
        .run_start(spec_driver.spec_module())
        .map_err(|trap| Error::Start(trap))?;
    Ok(())
}

fn load_module(
    wasm: &[u8],
    name: &Option<String>,
    spec_driver: &mut SpecDriver,
) -> Result<ModuleRef, Error> {
    let module = try_load_module(wasm)?;
    let instance = ModuleInstance::new(&module, spec_driver)
        .map_err(|e| Error::Load(e.to_string()))?
        .run_start(spec_driver.spec_module())
        .map_err(|trap| Error::Start(trap))?;

    let module_name = name.clone();
    spec_driver.add_module(module_name, instance.clone());

    Ok(instance)
}

fn run_action(
    program: &mut SpecDriver,
    action: &Action<u32, u64>,
) -> Result<Option<RuntimeValue>, InterpreterError> {
    match *action {
        Action::Invoke {
            ref module,
            ref field,
            ref args,
        } => {
            let module = program
                .module_or_last(module.as_ref().map(|x| x.as_ref()))
                .expect(&format!(
                    "Expected program to have loaded module {:?}",
                    module
                ));
            let vec_args = args
                .iter()
                .cloned()
                .map(spec_to_runtime_value)
                .collect::<Vec<_>>();
            module.invoke_export(field, &vec_args, program.spec_module())
        }
        Action::Get {
            ref module,
            ref field,
            ..
        } => {
            let module = program
                .module_or_last(module.as_ref().map(|x| x.as_ref()))
                .expect(&format!(
                    "Expected program to have loaded module {:?}",
                    module
                ));
            let global = module
                .export_by_name(&field)
                .ok_or_else(|| {
                    InterpreterError::Global(format!("Expected to have export with name {}", field))
                })?
                .as_global()
                .cloned()
                .ok_or_else(|| {
                    InterpreterError::Global(format!("Expected export {} to be a global", field))
                })?;
            Ok(Some(global.get()))
        }
    }
}

pub fn spec(name: &str) {
    println!("running test: {}", name);
    try_spec(name).expect("Failed to run spec");
}

fn try_spec(name: &str) -> Result<(), Error> {
    let mut spec_driver = SpecDriver::new();
    let spec_script_path = format!("tests/spec/testsuite/{}.wast", name);

    use std::io::Read;
    let mut spec_source = Vec::new();
    let mut spec_file = File::open(&spec_script_path).expect("Can't open file");
    spec_file
        .read_to_end(&mut spec_source)
        .expect("Can't read file");

    let mut parser = ScriptParser::from_source_and_name(&spec_source, &format!("{}.wast", name))
        .expect("Can't read spec script");
    let mut errors = vec![];

    while let Some(Command { kind, line }) = parser.next()? {
        macro_rules! assert_eq {
            ($a:expr, $b:expr) => {{
                let (a, b) = ($a, $b);

                if a != b {
                    errors.push(format!(
                        r#"ERROR (line {}):
  expected: {:?}
  got:      {:?}
"#,
                        line, b, a,
                    ));
                }
            }};
        }

        println!("Running spec cmd {}: {:?}", line, kind);

        match kind {
            CommandKind::Module { name, module, .. } => {
                load_module(&module.into_vec(), &name, &mut spec_driver)
                    .expect("Failed to load module");
            }
            CommandKind::AssertReturn { action, expected } => {
                let result = run_action(&mut spec_driver, &action);
                match result {
                    Ok(result) => {
                        let spec_expected = expected
                            .iter()
                            .cloned()
                            .map(spec_to_runtime_value)
                            .collect::<Vec<_>>();
                        let actual_result = result.into_iter().collect::<Vec<RuntimeValue>>();
                        for (actual_result, spec_expected) in
                            actual_result.iter().zip(spec_expected.iter())
                        {
                            assert_eq!(actual_result.value_type(), spec_expected.value_type());
                            // f32::NAN != f32::NAN
                            match spec_expected {
                                &RuntimeValue::F32(val) if val.is_nan() => match actual_result {
                                    &RuntimeValue::F32(val) => assert!(val.is_nan()),
                                    _ => unreachable!(), // checked above that types are same
                                },
                                &RuntimeValue::F64(val) if val.is_nan() => match actual_result {
                                    &RuntimeValue::F64(val) => assert!(val.is_nan()),
                                    _ => unreachable!(), // checked above that types are same
                                },
                                spec_expected @ _ => assert_eq!(actual_result, spec_expected),
                            }
                        }
                    }
                    Err(e) => {
                        panic!("Expected action to return value, got error: {:?}", e);
                    }
                }
            }
            CommandKind::AssertReturnCanonicalNan { action }
            | CommandKind::AssertReturnArithmeticNan { action } => {
                let result = run_action(&mut spec_driver, &action);
                match result {
                    Ok(result) => {
                        for actual_result in result.into_iter().collect::<Vec<RuntimeValue>>() {
                            match actual_result {
                                RuntimeValue::F32(val) => {
                                    if !val.is_nan() {
                                        panic!("Expected nan value, got {:?}", val)
                                    }
                                }
                                RuntimeValue::F64(val) => {
                                    if !val.is_nan() {
                                        panic!("Expected nan value, got {:?}", val)
                                    }
                                }
                                val @ _ => {
                                    panic!("Expected action to return float value, got {:?}", val)
                                }
                            }
                        }
                    }
                    Err(e) => {
                        panic!("Expected action to return value, got error: {:?}", e);
                    }
                }
            }
            CommandKind::AssertExhaustion { action, .. } => {
                let result = run_action(&mut spec_driver, &action);
                match result {
                    Ok(result) => panic!("Expected exhaustion, got result: {:?}", result),
                    Err(_e) => {}
                }
            }
            CommandKind::AssertTrap { action, .. } => {
                let result = run_action(&mut spec_driver, &action);
                match result {
                    Ok(result) => {
                        panic!(
                            "Expected action to result in a trap, got result: {:?}",
                            result
                        );
                    }
                    Err(_e) => {}
                }
            }
            CommandKind::AssertInvalid { module, .. }
            | CommandKind::AssertMalformed { module, .. }
            | CommandKind::AssertUnlinkable { module, .. } => {
                let module_load = try_load(&module.into_vec(), &mut spec_driver);
                match module_load {
                    Ok(_) => panic!("Expected invalid module definition, got some module!"),
                    Err(_e) => {}
                }
            }
            CommandKind::AssertUninstantiable { module, .. } => {
                match try_load(&module.into_vec(), &mut spec_driver) {
                    Ok(_) => panic!("Expected error running start function at line {}", line),
                    Err(_e) => {}
                }
            }
            CommandKind::Register { name, as_name, .. } => {
                let module = match spec_driver.module_or_last(name.as_ref().map(|x| x.as_ref())) {
                    Ok(module) => module,
                    Err(e) => panic!("No such module, at line {} - ({:?})", e, line),
                };
                spec_driver.add_module(Some(as_name.clone()), module);
            }
            CommandKind::PerformAction(action) => match run_action(&mut spec_driver, &action) {
                Ok(_) => {}
                Err(e) => panic!("Failed to invoke action at line {}: {:?}", line, e),
            },
        }
    }

    if !errors.is_empty() {
        use std::fmt::Write;
        let mut out = "\n".to_owned();
        for err in errors {
            write!(out, "{}", err).expect("Error formatting errors");
        }
        panic!(out);
    }

    Ok(())
}