wasmworker 0.2.1

use std::{
    cell::RefCell,
    collections::HashMap,
    rc::Rc,
    sync::atomic::{AtomicU32, Ordering},
};

use super::com::*;
use super::js::*;
use js_sys::Array;
use serde::{Deserialize, Serialize};
use tokio::sync::{oneshot, Semaphore};
use wasm_bindgen::{prelude::Closure, JsCast, JsValue, UnwrapThrowExt};
use web_sys::{
    Blob, BlobPropertyBag, MessageChannel, MessageEvent, MessagePort, Url, Worker, WorkerOptions,
    WorkerType,
};

use crate::{
    channel::Channel,
    channel_task::ChannelTask,
    convert::{from_bytes, to_bytes},
    error::{Full, InitError},
    func::{WebWorkerChannelFn, WebWorkerFn},
};

/// An internal type for the callback.
type Callback = dyn FnMut(MessageEvent);

/// This struct represents a single web worker instance.
/// It can be created using [`WebWorker::new`] or [`WebWorker::with_path`].
/// When an instance of this type is dropped, it also terminates the corresponding web worker.
///
/// Example usage:
/// ```no_run
/// # use serde::{Serialize, Deserialize};
/// # use wasmworker_proc_macro::webworker_fn;
/// # #[derive(Serialize, Deserialize, PartialEq, Debug)]
/// # struct VecType(Vec<u32>);
/// # #[webworker_fn]
/// # pub fn sort_vec(mut v: VecType) -> VecType { v.0.sort(); v }
/// use wasmworker::{webworker, WebWorker};
///
/// # async fn example() {
/// let worker = WebWorker::new(None).await.expect("Couldn't create worker");
/// let res = worker.run(webworker!(sort_vec), &VecType(vec![5, 2, 8])).await;
/// assert_eq!(res.0, vec![2, 5, 8]);
/// # }
/// # fn main() {}
/// ```
pub struct WebWorker {
    /// The underlying web worker.
    worker: Worker,
    /// An optional limit on the number of tasks queued at the same time.
    task_limit: Option<Semaphore>,
    /// The current task id, which is used to reidentify responses.
    current_task: AtomicU32,
    /// A map between task ids and the channel they need to be sent out with.
    /// [`Response`]s will arrive on our callback and we redistribute them to their origin.
    open_tasks: Rc<RefCell<HashMap<u32, oneshot::Sender<Response>>>>,
    /// The callback handle for the worker.
    _callback: Closure<Callback>,
}

impl WebWorker {
    /// This function takes the [`WORKER_JS`] and creates the corresponding
    /// worker blob after inserting the given path.
    /// If no `wasm_path` is provided, the [`main_js()`] path is used.
    fn worker_blob(
        wasm_path: Option<&str>,
        wasm_bg_path: Option<&str>,
        has_precompiled_module: bool,
    ) -> String {
        let blob_options = BlobPropertyBag::new();
        blob_options.set_type("application/javascript");

        let mut wasm_path_owned = None;
        let wasm_path = wasm_path.unwrap_or_else(|| {
            // Calculate path to wasm import.
            wasm_path_owned = Some(main_js().as_string().unwrap_throw());
            wasm_path_owned.as_ref().unwrap_throw()
        });

        let wasm_bg_path = match wasm_bg_path {
            Some(path) => format!("{{module_or_path: '{path}'}}"),
            None => "undefined".to_string(),
        };

        let worker_js = if has_precompiled_module {
            super::js::WORKER_JS_WITH_PRECOMPILED
        } else {
            super::js::WORKER_JS
        };

        let code = Array::new();
        code.push(&JsValue::from_str(
            &worker_js
                .replace("{{wasm}}", wasm_path)
                .replace("{{wasm_bg}}", &wasm_bg_path),
        ));

        Url::create_object_url_with_blob(
            &Blob::new_with_blob_sequence_and_options(&code.into(), &blob_options)
                .expect_throw("Couldn't create blob"),
        )
        .expect_throw("Couldn't create object URL")
    }

    /// Create a new [`WebWorker`] with an optional limit on the number of tasks queued.
    /// This can fail with an [`InitError`], for example, if the automatically inferred
    /// path for the wasm-bindgen glue is wrong.
    pub async fn new(task_limit: Option<usize>) -> Result<WebWorker, InitError> {
        Self::with_path(None, None, task_limit).await
    }

    /// Create a new [`WebWorker`] with an optional limit on the number of tasks queued.
    /// This function also receives an optional `main_js` path, which should point to the
    /// glue file generated by wasm-bindgen.
    ///
    /// In the standard setup, the path is automatically inferred.
    /// In more complex setups, it might be necessary to manually set this path.
    /// If a wrong path is given, a [`InitError`] will be returned.
    pub async fn with_path(
        main_js: Option<&str>,
        main_bg_js: Option<&str>,
        task_limit: Option<usize>,
    ) -> Result<WebWorker, InitError> {
        Self::with_path_and_module(main_js, main_bg_js, task_limit, None).await
    }

    /// Create a new [`WebWorker`] with an optional limit on the number of tasks queued
    /// and an optional pre-compiled WASM module.
    pub async fn with_path_and_module(
        main_js: Option<&str>,
        main_bg_js: Option<&str>,
        task_limit: Option<usize>,
        wasm_module: Option<js_sys::WebAssembly::Module>,
    ) -> Result<WebWorker, InitError> {
        // Create worker
        let worker_options = WorkerOptions::new();
        worker_options.set_type(WorkerType::Module);
        let script_url = WebWorker::worker_blob(main_js, main_bg_js, wasm_module.is_some());

        let worker = Worker::new_with_options(&script_url, &worker_options)
            .map_err(InitError::WebWorkerCreation)?;

        // Send pre-compiled WASM module if provided
        if let Some(module) = wasm_module {
            let init_msg = js_sys::Object::new();
            js_sys::Reflect::set(
                &init_msg,
                &JsValue::from_str("type"),
                &JsValue::from_str("wasm_module"),
            )
            .expect_throw("Could not set type");
            js_sys::Reflect::set(&init_msg, &JsValue::from_str("module"), &module)
                .expect_throw("Could not set module");

            worker
                .post_message(&init_msg)
                .expect_throw("Could not send WASM module to worker");
        }

        // Wait until worker is initialized.
        let (tx, rx) = oneshot::channel();
        let handler = Closure::once(move |event: MessageEvent| {
            let data = event.data();
            let post_init: PostInit = serde_wasm_bindgen::from_value(data)
                .expect_throw("Error deserializing post init data");
            let _ = tx.send(post_init);
        });
        worker.set_onmessage(Some(handler.as_ref().unchecked_ref()));
        let post_init = rx.await.expect_throw("WebWorker init sender dropped");

        // Handle errors in webworker init
        if !post_init.success {
            return Err(InitError::WebWorkerModuleLoading(
                post_init
                    .message
                    .expect_throw("Post init should have error message"),
            ));
        }

        let tasks = Rc::new(RefCell::new(HashMap::new()));

        let callback_handle = Self::callback(Rc::clone(&tasks));
        worker.set_onmessage(Some(callback_handle.as_ref().unchecked_ref()));

        Ok(WebWorker {
            worker,
            task_limit: task_limit.map(|limit| Semaphore::new(limit)),
            current_task: AtomicU32::new(0),
            open_tasks: tasks,
            _callback: callback_handle,
        })
    }

    /// Function to be called when a result is ready.
    fn callback(
        tasks: Rc<RefCell<HashMap<u32, oneshot::Sender<Response>>>>,
    ) -> Closure<dyn FnMut(MessageEvent)> {
        Closure::new(move |event: MessageEvent| {
            let data = event.data();
            let response: Response =
                serde_wasm_bindgen::from_value(data).expect_throw("Could not deserialize response");
            let mut tasks_wg = tasks.borrow_mut();

            // Send response on channel.
            if let Some(channel) = tasks_wg.remove(&response.id) {
                // Ignore if receiver is already closed.
                let _ = channel.send(response);
            }
        })
    }

    /// This is the most general function to outsource a task on a [`WebWorker`].
    /// It will automatically handle serialization of the argument, scheduling of the task on the worker,
    /// and deserialization of the return value.
    ///
    /// The `func`: [`WebWorkerFn`] argument should normally be instantiated using the [`crate::webworker!`] macro.
    /// This ensures type safety and that the function is correctly exposed to the worker.
    ///
    /// If a task limit has been set, this function will yield until previous tasks have been finished.
    /// This is achieved by a semaphore before task submission to the worker.
    ///
    /// Example:
    /// ```ignore
    /// worker.run(webworker!(sort_vec), &my_vec).await
    /// ```
    pub async fn run<T, R>(&self, func: WebWorkerFn<T, R>, arg: &T) -> R
    where
        T: Serialize + for<'de> Deserialize<'de>,
        R: Serialize + for<'de> Deserialize<'de>,
    {
        self.run_internal(func, arg).await
    }

    /// Run an async function with bidirectional channel support on this [`WebWorker`].
    ///
    /// Returns a [`ChannelTask`] that provides both the communication channel and the
    /// task result. The `MessageChannel` is created internally — callers interact only
    /// through the returned `ChannelTask`.
    ///
    /// The `func`: [`WebWorkerChannelFn`] argument should normally be instantiated using the
    /// [`crate::webworker_channel!`] macro. This ensures type safety and that the function
    /// is correctly exposed to the worker.
    ///
    /// If a task limit has been set, this function will yield until previous tasks have been finished.
    ///
    /// Example:
    /// ```ignore
    /// let task = worker
    ///     .run_channel(webworker_channel!(process_with_progress), &data)
    ///     .await;
    ///
    /// let progress: Progress = task.recv().await.expect("progress");
    /// task.send(&Continue { should_continue: true });
    /// let result: ProcessResult = task.result().await;
    /// ```
    pub async fn run_channel<T, R>(&self, func: WebWorkerChannelFn<T, R>, arg: &T) -> ChannelTask<R>
    where
        T: Serialize + for<'de> Deserialize<'de>,
        R: Serialize + for<'de> Deserialize<'de>,
    {
        self.run_channel_internal(func, arg).await
    }

    /// This function differs from [`WebWorker::run`] by returning early if the given task limit is reached.
    /// In this case a [`Full`] error is returned.
    ///
    /// The `func`: [`WebWorkerFn`] argument should normally be instantiated using the [`crate::webworker!`] macro.
    /// This ensures type safety and that the function is correctly exposed to the worker.
    ///
    /// If no task limit has been set, this function can never return an error.
    ///
    /// Example:
    /// ```ignore
    /// worker.try_run(webworker!(sort_vec), &my_vec).await
    /// ```
    pub async fn try_run<T, R>(&self, func: WebWorkerFn<T, R>, arg: &T) -> Result<R, Full>
    where
        T: Serialize + for<'de> Deserialize<'de>,
        R: Serialize + for<'de> Deserialize<'de>,
    {
        self.try_run_internal(func, arg).await
    }

    /// This function can outsource a task on a [`WebWorker`] which has `Box<[u8]>` both as input and output.
    /// (De)serialization of values needs to be handled by the caller.
    /// For more convenient access, make sure the `serde` feature is enabled and use [`WebWorker::run`].
    ///
    /// The `func`: [`WebWorkerFn`] argument should normally be instantiated using the [`crate::webworker!`] macro.
    /// This ensures type safety and that the function is correctly exposed to the worker.
    ///
    /// If a task limit has been set, this function will yield until previous tasks have been finished.
    /// This is achieved by a semaphore before task submission to the worker.
    ///
    /// Example:
    /// ```ignore
    /// worker.run_bytes(webworker!(sort), &my_box).await
    /// ```
    pub async fn run_bytes(
        &self,
        func: WebWorkerFn<Box<[u8]>, Box<[u8]>>,
        arg: &Box<[u8]>,
    ) -> Box<[u8]> {
        self.run_internal(func, arg).await
    }

    /// This function differs from [`WebWorker::run_bytes`] by returning early if the given task limit is reached.
    /// In this case a [`Full`] error is returned.
    /// (De)serialization of values needs to be handled by the caller.
    /// For more convenient access, make sure the `serde` feature is enabled and use [`WebWorker::try_run`].
    ///
    /// The `func`: [`WebWorkerFn`] argument should normally be instantiated using the [`crate::webworker!`] macro.
    /// This ensures type safety and that the function is correctly exposed to the worker.
    ///
    /// If no task limit has been set, this function can never return an error.
    ///
    /// Example:
    /// ```ignore
    /// worker.try_run_bytes(webworker!(sort), &my_box).await
    /// ```
    pub async fn try_run_bytes(
        &self,
        func: WebWorkerFn<Box<[u8]>, Box<[u8]>>,
        arg: &Box<[u8]>,
    ) -> Result<Box<[u8]>, Full> {
        self.try_run_internal(func, arg).await
    }

    /// Internal function to schedule a simple task to the worker.
    /// This variant returns early if a semaphore permit cannot be obtained immediately.
    pub(crate) async fn try_run_internal<T, R>(
        &self,
        func: WebWorkerFn<T, R>,
        arg: &T,
    ) -> Result<R, Full>
    where
        T: Serialize + for<'de> Deserialize<'de>,
        R: Serialize + for<'de> Deserialize<'de>,
    {
        // Acquire permit if necessary.
        let _permit = if let Some(ref s) = self.task_limit {
            Some(match s.try_acquire() {
                Ok(permit) => permit,
                Err(_) => return Err(Full),
            })
        } else {
            None
        };

        // Convert arg and result.
        Ok(self.force_run(func.name, arg, false, None).await)
    }

    /// Internal function to schedule a simple task to the worker.
    pub(crate) async fn run_internal<T, R>(&self, func: WebWorkerFn<T, R>, arg: &T) -> R
    where
        T: Serialize + for<'de> Deserialize<'de>,
        R: Serialize + for<'de> Deserialize<'de>,
    {
        // Acquire permit if necessary.
        let _permit = if let Some(ref s) = self.task_limit {
            Some(s.acquire().await.unwrap())
        } else {
            None
        };

        // Convert arg and result.
        self.force_run(func.name, arg, false, None).await
    }

    /// Internal function to schedule a channel task to the worker.
    /// Creates a `MessageChannel` internally, sends one port to the worker,
    /// and returns a `ChannelTask` wrapping the other port and the result future.
    pub(crate) async fn run_channel_internal<T, R>(
        &self,
        func: WebWorkerChannelFn<T, R>,
        arg: &T,
    ) -> ChannelTask<R>
    where
        T: Serialize + for<'de> Deserialize<'de>,
        R: Serialize + for<'de> Deserialize<'de>,
    {
        // Acquire permit if necessary.
        let _permit = if let Some(ref s) = self.task_limit {
            Some(s.acquire().await.unwrap())
        } else {
            None
        };

        // Create the MessageChannel internally.
        let msg_channel = MessageChannel::new().expect_throw("Could not create MessageChannel");
        let channel = Channel::from(msg_channel.port1());
        let worker_port = msg_channel.port2();

        // Send the request and get a receiver for the result bytes.
        let result_rx = self.send_channel_request(func.name, arg, worker_port);

        ChannelTask::new(channel, result_rx)
    }

    /// This function handles the communication with the worker
    /// after the task limit has been checked.
    /// It also handles (de)serialization.
    async fn force_run<T, R>(
        &self,
        func_name: &'static str,
        arg: &T,
        is_channel: bool,
        port: Option<MessagePort>,
    ) -> R
    where
        T: Serialize + for<'de> Deserialize<'de>,
        R: Serialize + for<'de> Deserialize<'de>,
    {
        let id = self.current_task.fetch_add(1, Ordering::Relaxed);
        let request = Request {
            id,
            func_name,
            is_channel,
            arg: to_bytes(arg),
        };

        let res = self.send_request(id, request, port).await;
        from_bytes(&res)
    }

    /// Sends a request to the worker and waits for the response.
    /// This is extracted from `force_run` to reduce monomorphisation cost.
    async fn send_request(&self, id: u32, request: Request, port: Option<MessagePort>) -> Vec<u8> {
        // Create channel and add task.
        let (sender, receiver) = oneshot::channel();
        self.open_tasks.borrow_mut().insert(id, sender);

        // send the task to the webworker, either with a port or without one
        if let Some(port) = port {
            let transfer = Array::new();
            transfer.push(&port);

            self.worker
                .post_message_with_transfer(
                    &serde_wasm_bindgen::to_value(&request)
                        .expect_throw("Could not serialize request"),
                    &transfer,
                )
                .expect_throw("WebWorker gone");
        } else {
            self.worker
                .post_message(
                    &serde_wasm_bindgen::to_value(&request)
                        .expect_throw("Could not serialize request"),
                )
                .expect_throw("WebWorker gone");
        }

        // Handle result.
        receiver
            .await
            .expect_throw("WebWorker gone")
            .response
            .expect_throw("Could not find function")
    }

    /// Sends a channel request to the worker and returns a receiver for the result bytes.
    /// Unlike `send_request`, this does not await the result — it returns immediately
    /// so the caller can interact with the channel before consuming the result.
    fn send_channel_request<T>(
        &self,
        func_name: &'static str,
        arg: &T,
        port: MessagePort,
    ) -> oneshot::Receiver<Vec<u8>>
    where
        T: Serialize + for<'de> Deserialize<'de>,
    {
        let id = self.current_task.fetch_add(1, Ordering::Relaxed);
        let request = Request {
            id,
            func_name,
            is_channel: true,
            arg: to_bytes(arg),
        };

        let (sender, receiver) = oneshot::channel();
        self.open_tasks.borrow_mut().insert(id, sender);

        let transfer = Array::new();
        transfer.push(&port);

        self.worker
            .post_message_with_transfer(
                &serde_wasm_bindgen::to_value(&request).expect_throw("Could not serialize request"),
                &transfer,
            )
            .expect_throw("WebWorker gone");

        // Map the receiver to extract just the response bytes.
        let (byte_sender, byte_receiver) = oneshot::channel();
        wasm_bindgen_futures::spawn_local(async move {
            if let Ok(response) = receiver.await {
                let _ = byte_sender.send(response.response.expect("Could not find function"));
            }
        });

        byte_receiver
    }

    /// Return the current capacity for new tasks.
    pub fn capacity(&self) -> Option<usize> {
        self.task_limit.as_ref().map(|s| s.available_permits())
    }

    /// Return the number of tasks currently queued to this worker.
    pub fn current_load(&self) -> usize {
        self.open_tasks.borrow().len()
    }
}

impl Drop for WebWorker {
    fn drop(&mut self) {
        self.worker.terminate();
    }
}