// Copyright 2025 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Routers allow converting IPC subchannels to crossbeam channels.
//! The [RouterProxy] provides various methods to register
//! `SubReceiver<T>`s. The router will then either call the appropriate callback or route the
//! message to a crossbeam `Sender<T>` or `Receiver<T>`. You should use the global `ROUTER` to
//! access the `RouterProxy` methods (via `ROUTER`'s `Deref` for `RouterProxy`.

use lazy_static::lazy_static;
use std::collections::HashMap;
use std::sync::Mutex;
use std::thread;

use crate::mux::{
    self, MultiplexError, OpaqueSubReceiver, RawMessage, SubReceiver, SubReceiverSet,
    SubSelectionResult, SubSender,
};
use crossbeam_channel::{self, Receiver, Sender};
use serde::{Deserialize, Serialize};

lazy_static! {
    /// Global object wrapping a `RouterProxy`.
    /// Add routes ([add_typed_route](RouterProxy::add_typed_route)), or convert `SubReceiver<T>`
    /// to crossbeam receivers (e.g. [route_subreceiver_to_new_crossbeam_receiver](RouterProxy::route_subreceiver_to_new_crossbeam_receiver)).
    pub static ref ROUTER: RouterProxy = RouterProxy::new();
}

/// A `RouterProxy` provides methods for establishing and talking to the router.
pub struct RouterProxy {
    comm: Mutex<RouterProxyComm>,
}

#[allow(clippy::new_without_default)]
impl RouterProxy {
    /// Spawn a router thread which waits for events on its registered `SubReceiver<T>`s.
    /// The `RouterProxy`'s methods communicate with the running router thread to
    /// register new `SubReceiver<T>`'s.
    pub fn new() -> RouterProxy {
        // Router acts like a receiver, running in its own thread with both
        // receiver ends.
        // Router proxy takes both sending ends.
        let (msg_sender, msg_receiver) = crossbeam_channel::unbounded();
        let chan = mux::Channel::new().unwrap();
        let (wakeup_sender, wakeup_receiver) = chan.sub_channel();
        thread::spawn(move || Router::new(msg_receiver, wakeup_receiver).run());
        RouterProxy {
            comm: Mutex::new(RouterProxyComm {
                msg_sender,
                wakeup_sender,
                shutdown: false,
            }),
        }
    }

    /// Add a new (receiver, callback) pair to the router, and send a wakeup message
    /// to the router.
    ///
    /// Consider using [add_typed_route](Self::add_typed_route) instead, which prevents
    /// mismatches between the receiver and callback types.
    fn add_route(&self, subreceiver: OpaqueSubReceiver, callback: RouterHandler) {
        let comm = self.comm.lock().unwrap();

        if comm.shutdown {
            return;
        }

        comm.msg_sender
            .send(RouterMsg::AddRoute(subreceiver, callback))
            .unwrap();
        comm.wakeup_sender.send(()).unwrap();
    }

    /// Add a new `(receiver, callback)` pair to the router, and send a wakeup message
    /// to the router. This method is strongly typed and guarantees
    /// that `subreceiver` and `callback` use the same message type.
    pub fn add_typed_route<T>(
        &self,
        subreceiver: SubReceiver<T>,
        mut callback: TypedRouterHandler<T>,
    ) where
        T: Serialize + for<'de> Deserialize<'de> + 'static,
    {
        // Before passing the message on to the callback, turn it into the appropriate type
        let modified_callback = move |msg: RawMessage| {
            let typed_message = msg.to::<T>();
            callback(typed_message)
        };

        self.add_route(subreceiver.to_opaque(), Box::new(modified_callback));
    }

    /// Send a shutdown message to the router containing an ACK sender,
    /// send a wakeup message to the router, and block on the ACK.
    /// Calling it is idempotent,
    /// which can be useful when running a multi-process system in single-process mode.
    pub fn shutdown(&self) {
        let mut comm = self.comm.lock().unwrap();

        if comm.shutdown {
            return;
        }
        comm.shutdown = true;

        let (ack_sender, ack_receiver) = crossbeam_channel::unbounded();
        comm.wakeup_sender
            .send(())
            .map(|_| {
                comm.msg_sender
                    .send(RouterMsg::Shutdown(ack_sender))
                    .unwrap();
                ack_receiver.recv().unwrap();
            })
            .unwrap();
    }

    /// A convenience function to route an `SubReceiver<T>` to an existing `Sender<T>`.
    pub fn route_subreceiver_to_crossbeam_sender<T>(
        &self,
        subreceiver: SubReceiver<T>,
        crossbeam_sender: Sender<T>,
    ) where
        T: for<'de> Deserialize<'de> + Serialize + Send + 'static,
    {
        self.add_typed_route(
            subreceiver,
            Box::new(move |message| drop(crossbeam_sender.send(message.unwrap()))),
        )
    }

    /// A convenience function to route an `SubReceiver<T>` to a `Receiver<T>`: the most common
    /// use of a `Router`.
    pub fn route_subreceiver_to_new_crossbeam_receiver<T>(
        &self,
        subreceiver: SubReceiver<T>,
    ) -> Receiver<T>
    where
        T: for<'de> Deserialize<'de> + Serialize + Send + 'static,
    {
        let (crossbeam_sender, crossbeam_receiver) = crossbeam_channel::unbounded();
        self.route_subreceiver_to_crossbeam_sender(subreceiver, crossbeam_sender);
        crossbeam_receiver
    }
}

struct RouterProxyComm {
    msg_sender: Sender<RouterMsg>,
    wakeup_sender: SubSender<()>,
    shutdown: bool,
}

/// Router runs in its own thread listening for events. Adds events to its SubReceiverSet
/// and listens for events using select().
struct Router {
    /// Get messages from RouterProxy.
    msg_receiver: Receiver<RouterMsg>,
    /// The ID/index of the special channel we use to identify messages from msg_receiver.
    msg_wakeup_id: u64,
    /// Set of all receivers which have been registered for us to select on.
    subreceiver_set: SubReceiverSet,
    /// Maps ids to their handler functions.
    handlers: HashMap<u64, RouterHandler>,
}

impl Router {
    fn new(msg_receiver: Receiver<RouterMsg>, wakeup_receiver: SubReceiver<()>) -> Router {
        let mut subreceiver_set = SubReceiverSet::new().unwrap();
        let msg_wakeup_id = subreceiver_set.add(wakeup_receiver).unwrap();
        Router {
            msg_receiver,
            msg_wakeup_id,
            subreceiver_set,
            handlers: HashMap::new(),
        }
    }

    /// Continuously loop waiting for wakeup signals from router proxy.
    /// Iterate over events either:
    /// 1) If a message comes in from our special `wakeup_receiver` (identified through
    ///    msg_wakeup_id. Read message from `msg_receiver` and add a new receiver
    ///    to our receiver set.
    /// 2) Call appropriate handler based on message id.
    /// 3) Remove handler once channel closes.
    fn run(&mut self) {
        loop {
            // Wait for events to come from our select() new channels are added to
            // our ReceiverSet below.
            let results = match self.subreceiver_set.select() {
                Ok(results) => results,
                Err(_) => break,
            };

            // Iterate over numerous events that were ready at this time.
            for result in results.into_iter() {
                match result {
                    // Message came from the RouterProxy. Listen on our `msg_receiver`
                    // channel.
                    SubSelectionResult::MessageReceived(id, _) if id == self.msg_wakeup_id => {
                        match self.msg_receiver.recv().unwrap() {
                            RouterMsg::AddRoute(receiver, handler) => {
                                let new_receiver_id =
                                    self.subreceiver_set.add_opaque(receiver).unwrap();
                                self.handlers.insert(new_receiver_id, handler);
                            },
                            RouterMsg::Shutdown(sender) => {
                                sender
                                    .send(())
                                    .expect("Failed to send comfirmation of shutdown.");
                                break;
                            },
                        }
                    },
                    // Event from one of our registered receivers, call callback.
                    SubSelectionResult::MessageReceived(id, message) => {
                        self.handlers.get_mut(&id).unwrap()(message)
                    },
                    SubSelectionResult::ChannelClosed(id) => {
                        let _ = self.handlers.remove(&id).unwrap();
                    },
                }
            }
        }
    }
}

enum RouterMsg {
    /// Register the receiver OpaqueSubReceiver for listening for events on.
    /// When a message comes from this receiver, call RouterHandler.
    AddRoute(OpaqueSubReceiver, RouterHandler),
    /// Shutdown the router, providing a sender to send an acknowledgement.
    Shutdown(Sender<()>),
}

/// Function to call when a new event is received from the corresponding receiver.
pub type RouterHandler = Box<dyn FnMut(RawMessage) + Send>;

/// Like [RouterHandler] but includes the type that will be passed to the callback
pub type TypedRouterHandler<T> = Box<dyn FnMut(Result<T, MultiplexError>) + Send>;
