#![allow(clippy::result_large_err)]
use super::{BinaryPackage, BinaryPackageCommand};
use crate::{
    RewindState, SpawnError, WasiError, WasiRuntimeError,
    os::task::{
        TaskJoinHandle,
        thread::{RewindResultType, WasiThreadRunGuard},
    },
    runtime::{
        ModuleInput, TaintReason,
        module_cache::HashedModuleData,
        task_manager::{
            TaskWasm, TaskWasmRecycle, TaskWasmRecycleProperties, TaskWasmRunProperties,
        },
    },
    state::context_switching::ContextSwitchingEnvironment,
    syscalls::rewind_ext,
};
use crate::{Runtime, WasiEnv, WasiFunctionEnv};
use std::{borrow::Cow, sync::Arc};
use tracing::*;
use virtual_mio::block_on;
use wasmer::{Function, Memory32, Memory64, Module, RuntimeError, Store, Value};
use wasmer_wasix_types::wasi::Errno;

#[tracing::instrument(level = "trace", skip_all, fields(%name, package_id=%binary.id))]
pub async fn spawn_exec(
    binary: BinaryPackage,
    name: &str,
    env: WasiEnv,
    runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
) -> Result<TaskJoinHandle, SpawnError> {
    spawn_union_fs(&env, &binary).await?;

    let cmd = package_command_by_name(&binary, name)?;
    let input = ModuleInput::Command(Cow::Borrowed(cmd));
    let module = runtime.resolve_module(input, None, None).await?;

    // Free the space used by the binary, since we don't need it
    // any longer
    drop(binary);

    spawn_exec_module(module, env, runtime)
}

#[tracing::instrument(level = "trace", skip_all, fields(%name))]
pub async fn spawn_exec_wasm(
    wasm: HashedModuleData,
    name: &str,
    env: WasiEnv,
    runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
) -> Result<TaskJoinHandle, SpawnError> {
    let module = spawn_load_module(name, wasm, runtime).await?;

    spawn_exec_module(module, env, runtime)
}

pub fn package_command_by_name<'a>(
    pkg: &'a BinaryPackage,
    name: &str,
) -> Result<&'a BinaryPackageCommand, SpawnError> {
    // If an explicit command is provided, use it.
    // Otherwise, use the entrypoint.
    // If no entrypoint exists, and the package has a single
    // command, then use it. This is done for backwards
    // compatibility.
    let cmd = if let Some(cmd) = pkg.get_command(name) {
        cmd
    } else if let Some(cmd) = pkg.get_entrypoint_command() {
        cmd
    } else {
        match pkg.commands.as_slice() {
            // Package only has a single command, so use it.
            [first] => first,
            // Package either has no command, or has multiple commands, which
            // would make the choice ambiguous, so fail.
            _ => {
                return Err(SpawnError::MissingEntrypoint {
                    package_id: pkg.id.clone(),
                });
            }
        }
    };

    Ok(cmd)
}

pub async fn spawn_load_module(
    name: &str,
    wasm: HashedModuleData,
    runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
) -> Result<Module, SpawnError> {
    match runtime.load_hashed_module(wasm, None).await {
        Ok(module) => Ok(module),
        Err(err) => {
            tracing::error!(
                command = name,
                error = &err as &dyn std::error::Error,
                "Failed to compile the module",
            );
            Err(err)
        }
    }
}

pub async fn spawn_union_fs(env: &WasiEnv, binary: &BinaryPackage) -> Result<(), SpawnError> {
    // If the file system has not already been union'ed then do so
    env.state
        .fs
        .conditional_union(binary)
        .await
        .map_err(|err| {
            tracing::warn!("failed to union file system - {err}");
            SpawnError::FileSystemError(crate::ExtendedFsError::with_msg(
                err,
                "could not union filesystems",
            ))
        })?;
    tracing::debug!("{:?}", env.state.fs);
    Ok(())
}

pub fn spawn_exec_module(
    module: Module,
    env: WasiEnv,
    runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
) -> Result<TaskJoinHandle, SpawnError> {
    // Create a new task manager
    let tasks = runtime.task_manager();

    // Create the signaler
    let pid = env.pid();

    let join_handle = env.thread.join_handle();
    {
        // Create a thread that will run this process
        let tasks_outer = tasks.clone();

        tasks_outer
            .task_wasm(
                TaskWasm::new(Box::new(run_exec), env, module, true, true).with_pre_run(Box::new(
                    |ctx, store| {
                        Box::pin(async move {
                            ctx.data(store).state.fs.close_cloexec_fds().await;
                        })
                    },
                )),
            )
            .map_err(|err| {
                error!("wasi[{}]::failed to launch module - {}", pid, err);
                SpawnError::Other(Box::new(err))
            })?
    };

    Ok(join_handle)
}

/// # SAFETY
/// This must be executed from the same thread that owns the instance as
/// otherwise it will cause a panic
unsafe fn run_recycle(
    callback: Option<Box<TaskWasmRecycle>>,
    ctx: WasiFunctionEnv,
    mut store: Store,
) {
    if let Some(callback) = callback {
        let env = ctx.data_mut(&mut store);
        let memory = unsafe { env.memory() }.clone();

        let props = TaskWasmRecycleProperties {
            env: env.clone(),
            memory,
            store,
        };
        callback(props);
    }
}

pub fn run_exec(props: TaskWasmRunProperties) {
    let ctx = props.ctx;
    let mut store = props.store;

    // Create the WasiFunctionEnv
    let thread = WasiThreadRunGuard::new(ctx.data(&store).thread.clone());
    let recycle = props.recycle;

    // Perform the initialization
    // If this module exports an _initialize function, run that first.
    if let Ok(initialize) = ctx
        .data(&store)
        .inner()
        .main_module_instance_handles()
        .instance
        .exports
        .get_function("_initialize")
        .cloned()
    {
        // This does not need a context switching environment as the documentation
        // states that that is only available after the first call to main
        let result = initialize.call(&mut store, &[]);

        if let Err(err) = result {
            thread.thread.set_status_finished(Err(err.into()));
            ctx.data(&store)
                .blocking_on_exit(Some(Errno::Noexec.into()));
            unsafe { run_recycle(recycle, ctx, store) };
            return;
        }
    }

    // Bootstrap the process
    // Unsafe: The bootstrap must be executed in the same thread that runs the
    //         actual WASM code
    let rewind_state = match unsafe { ctx.bootstrap(&mut store) } {
        Ok(r) => r,
        Err(err) => {
            tracing::warn!("failed to bootstrap - {}", err);
            thread.thread.set_status_finished(Err(err));
            ctx.data(&store)
                .blocking_on_exit(Some(Errno::Noexec.into()));
            unsafe { run_recycle(recycle, ctx, store) };
            return;
        }
    };

    // If there is a start function
    debug!("wasi[{}]::called main()", ctx.data(&store).pid());
    // TODO: rewrite to use crate::run_wasi_func

    // Call the module
    call_module(ctx, store, thread, rewind_state, recycle);
}

fn get_start(ctx: &WasiFunctionEnv, store: &Store) -> Option<Function> {
    ctx.data(store)
        .inner()
        .main_module_instance_handles()
        .instance
        .exports
        .get_function("_start")
        .cloned()
        .ok()
}

/// Calls the module
fn call_module(
    ctx: WasiFunctionEnv,
    mut store: Store,
    handle: WasiThreadRunGuard,
    rewind_state: Option<(RewindState, RewindResultType)>,
    recycle: Option<Box<TaskWasmRecycle>>,
) {
    let env = ctx.data(&store);
    let pid = env.pid();
    let tasks = env.tasks().clone();
    handle.thread.set_status_running();
    let runtime = env.runtime.clone();

    // If we need to rewind then do so
    if let Some((rewind_state, rewind_result)) = rewind_state {
        let mut ctx = ctx.env.clone().into_mut(&mut store);
        if rewind_state.is_64bit {
            let res = rewind_ext::<Memory64>(
                &mut ctx,
                Some(rewind_state.memory_stack),
                rewind_state.rewind_stack,
                rewind_state.store_data,
                rewind_result,
            );
            if res != Errno::Success {
                ctx.data().blocking_on_exit(Some(res.into()));
                unsafe { run_recycle(recycle, WasiFunctionEnv { env: ctx.as_ref() }, store) };
                return;
            }
        } else {
            let res = rewind_ext::<Memory32>(
                &mut ctx,
                Some(rewind_state.memory_stack),
                rewind_state.rewind_stack,
                rewind_state.store_data,
                rewind_result,
            );
            if res != Errno::Success {
                ctx.data().blocking_on_exit(Some(res.into()));
                unsafe { run_recycle(recycle, WasiFunctionEnv { env: ctx.as_ref() }, store) };
                return;
            }
        };
    }

    // Invoke the start function
    // Call the module
    let Some(start) = get_start(&ctx, &store) else {
        debug!("wasi[{}]::exec-failed: missing _start function", pid);
        ctx.data(&store)
            .blocking_on_exit(Some(Errno::Noexec.into()));
        unsafe { run_recycle(recycle, ctx, store) };
        return;
    };

    let (mut store, mut call_ret) =
        ContextSwitchingEnvironment::run_main_context(&ctx, store, start.clone(), vec![]);

    let mut store = loop {
        // Technically, it's an error for a vfork to return from main, but anyway...
        store = match resume_vfork(&ctx, store, &start, &call_ret) {
            // A vfork was resumed, there may be another, so loop back
            (store, Ok(Some(ret))) => {
                call_ret = ret;
                store
            }

            // An error was encountered when restoring from the vfork, report it
            (store, Err(e)) => {
                call_ret = Err(RuntimeError::user(Box::new(WasiError::Exit(e.into()))));
                break store;
            }

            // No vfork, keep the call_ret value
            (store, Ok(None)) => break store,
        };
    };

    let ret = if let Err(err) = call_ret {
        match err.downcast::<WasiError>() {
            Ok(WasiError::Exit(code)) if code.is_success() => Ok(Errno::Success),
            Ok(WasiError::ThreadExit) => Ok(Errno::Success),
            Ok(WasiError::Exit(code)) => {
                runtime.on_taint(TaintReason::NonZeroExitCode(code));
                Err(WasiError::Exit(code).into())
            }
            Ok(WasiError::DeepSleep(deep)) => {
                // Create the callback that will be invoked when the thread respawns after a deep sleep
                let rewind = deep.rewind;
                let respawn = {
                    move |ctx, store, rewind_result| {
                        // Call the thread
                        call_module(
                            ctx,
                            store,
                            handle,
                            Some((rewind, RewindResultType::RewindWithResult(rewind_result))),
                            recycle,
                        );
                    }
                };

                // Spawns the WASM process after a trigger
                if let Err(err) = unsafe {
                    tasks.resume_wasm_after_poller(Box::new(respawn), ctx, store, deep.trigger)
                } {
                    debug!("failed to go into deep sleep - {}", err);
                }
                return;
            }
            Ok(WasiError::UnknownWasiVersion) => {
                debug!("failed as wasi version is unknown");
                runtime.on_taint(TaintReason::UnknownWasiVersion);
                Ok(Errno::Noexec)
            }
            Ok(WasiError::DlSymbolResolutionFailed(symbol)) => {
                debug!("failed as a needed DL symbol could not be resolved");
                runtime.on_taint(TaintReason::DlSymbolResolutionFailed(symbol.clone()));
                Err(WasiError::DlSymbolResolutionFailed(symbol).into())
            }
            Err(err) => {
                runtime.on_taint(TaintReason::RuntimeError(err.clone()));
                Err(WasiRuntimeError::from(err))
            }
        }
    } else {
        Ok(Errno::Success)
    };

    let code = if let Err(err) = &ret {
        match err.as_exit_code() {
            Some(s) => s,
            None => {
                let err_display = err.display(&mut store);
                error!("{err_display}");
                eprintln!("{err_display}");
                Errno::Noexec.into()
            }
        }
    } else {
        Errno::Success.into()
    };

    // Cleanup the environment
    ctx.data(&store).blocking_on_exit(Some(code));
    unsafe { run_recycle(recycle, ctx, store) };

    debug!("wasi[{pid}]::main() has exited with {code}");
    handle.thread.set_status_finished(ret.map(|a| a.into()));
}

#[allow(clippy::type_complexity)]
fn resume_vfork(
    ctx: &WasiFunctionEnv,
    mut store: Store,
    start: &Function,
    call_ret: &Result<Box<[Value]>, RuntimeError>,
) -> (
    Store,
    Result<Option<Result<Box<[Value]>, RuntimeError>>, Errno>,
) {
    let (err, code) = match call_ret {
        Ok(_) => (None, wasmer_wasix_types::wasi::ExitCode::from(0u16)),
        Err(err) => match err.downcast_ref::<WasiError>() {
            // If the child process is just deep sleeping, we don't restore the vfork
            Some(WasiError::DeepSleep(..)) => return (store, Ok(None)),

            Some(WasiError::Exit(code)) => (None, *code),
            Some(WasiError::ThreadExit) => (None, wasmer_wasix_types::wasi::ExitCode::from(0u16)),
            Some(WasiError::UnknownWasiVersion) => (None, Errno::Noexec.into()),
            Some(WasiError::DlSymbolResolutionFailed(_)) => (None, Errno::Nolink.into()),
            None => (
                Some(WasiRuntimeError::from(err.clone())),
                Errno::Unknown.into(),
            ),
        },
    };

    if let Some(mut vfork) = ctx.data_mut(&mut store).vfork.take() {
        if let Some(err) = err {
            error!(%err, "Error from child process");
            eprintln!("{err}");
        }

        block_on(
            unsafe { ctx.data(&store).get_memory_and_wasi_state(&store, 0) }
                .1
                .fs
                .close_all(),
        );

        tracing::debug!(
            pid = %ctx.data_mut(&mut store).process.pid(),
            vfork_pid = %vfork.env.process.pid(),
            "Resuming from vfork after child process was terminated"
        );

        // Restore the WasiEnv to the point when we vforked
        vfork.env.swap_inner(ctx.data_mut(&mut store));
        std::mem::swap(vfork.env.as_mut(), ctx.data_mut(&mut store));
        let mut child_env = *vfork.env;
        child_env.owned_handles.push(vfork.handle);

        // Terminate the child process
        child_env.process.terminate(code);

        // If the vfork contained a context-switching environment, exit now
        if ctx.data(&store).context_switching_environment.is_some() {
            // We cannot recover from this situation when using context switching
            tracing::error!(
                "Terminated a vfork in another way than exit or exec which is undefined behaviour. In this case the parent process will be terminated."
            );
            return (store, Err(code.into()));
        }
        let Some(asyncify_info) = vfork.asyncify else {
            // We can only recover from this situation when using asyncify-based vforking; since asyncify is not in use here, we cannot recover and must terminate the parent process
            tracing::error!(
                "Terminated a vfork in another way than exit or exec which is undefined behaviour. In this case the parent process will be terminated."
            );
            return (store, Err(code.into()));
        };
        // TODO: We can also only safely recover if we are not using nested calling
        // TODO: Just delete this branch

        // Jump back to the vfork point and continue execution
        let child_pid = child_env.process.pid();
        let rewind_stack = asyncify_info.rewind_stack.freeze();
        let store_data = asyncify_info.store_data;

        let ctx_cloned = ctx.env.clone().into_mut(&mut store);
        // Now rewind the previous stack and carry on from where we did the vfork
        let rewind_result = if asyncify_info.is_64bit {
            crate::syscalls::rewind::<Memory64, _>(
                ctx_cloned,
                None,
                rewind_stack,
                store_data,
                crate::syscalls::ForkResult {
                    pid: child_pid.raw() as wasmer_wasix_types::wasi::Pid,
                    ret: Errno::Success,
                },
            )
        } else {
            crate::syscalls::rewind::<Memory32, _>(
                ctx_cloned,
                None,
                rewind_stack,
                store_data,
                crate::syscalls::ForkResult {
                    pid: child_pid.raw() as wasmer_wasix_types::wasi::Pid,
                    ret: Errno::Success,
                },
            )
        };

        match rewind_result {
            Errno::Success => {
                // We should only get here, if the engine does not support context switching
                // If the engine supports it, we should exit in the check a few lines above
                let (store, result) = ContextSwitchingEnvironment::run_main_context(
                    ctx,
                    store,
                    start.clone(),
                    vec![],
                );
                (store, Ok(Some(result)))
            }
            err => {
                warn!("fork failed - could not rewind the stack - errno={}", err);
                (store, Err(err))
            }
        }
    } else {
        (store, Ok(None))
    }
}
