// Structure of module

mod runtime;

mod messages;

mod connector;

mod native;

mod port;

mod scheduler;

mod inbox;

#[cfg(test)] mod tests;

// Imports

use std::collections::VecDeque;

use std::sync::{Arc, Condvar, Mutex, RwLock};

use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};

use std::thread::{self, JoinHandle};

use crate::collections::RawVec;

use crate::ProtocolDescription;

use inbox::Message;

use connector::{ConnectorPDL, ConnectorPublic, ConnectorScheduling};

use scheduler::{Scheduler, ControlMessageHandler};

use native::{Connector, ConnectorApplication, ApplicationInterface};

use crate::runtime2::port::{Port, PortState};

use crate::runtime2::scheduler::{ComponentCtxFancy, SchedulerCtx};

/// A kind of token that, once obtained, allows mutable access to a connector.

/// We're trying to use move semantics as much as possible: the owner of this

/// key is the only one that may execute the connector's code.

pub(crate) struct ConnectorKey {

    pub index: u32, // of connector

}

impl ConnectorKey {

    /// Downcasts the `ConnectorKey` type, which can be used to obtain mutable

    /// access, to a "regular ID" which can be used to obtain immutable access.

    #[inline]

    pub fn downcast(&self) -> ConnectorId {

        return ConnectorId(self.index);

    }

    /// Turns the `ConnectorId` into a `ConnectorKey`, marked as unsafe as it

    /// bypasses the type-enforced `ConnectorKey`/`ConnectorId` system

    #[inline]

    pub unsafe fn from_id(id: ConnectorId) -> ConnectorKey {

        return ConnectorKey{ index: id.0 };

    }

}

/// A kind of token that allows shared access to a connector. Multiple threads

/// may hold this

#[derive(Debug, Copy, Clone, PartialEq, Eq)]

pub struct ConnectorId(pub u32);

impl ConnectorId {

    // TODO: Like the other `new_invalid`, maybe remove

    #[inline]

    pub fn new_invalid() -> ConnectorId {

        return ConnectorId(u32::MAX);

    }

    #[inline]

    pub(crate) fn is_valid(&self) -> bool {

        return self.0 != u32::MAX;

    }

}

// TODO: Change this, I hate this. But I also don't want to put `public` and

//  `router` of `ScheduledConnector` back into `Connector`. The reason I don't

//  want `Box<dyn Connector>` everywhere is because of the v-table overhead. But

//  to truly design this properly I need some benchmarks.

pub(crate) enum ConnectorVariant {

    UserDefined(ConnectorPDL),

    Native(Box<dyn Connector>),

}

impl Connector for ConnectorVariant {

    fn run(&mut self, scheduler_ctx: SchedulerCtx, comp_ctx: &mut ComponentCtxFancy) -> ConnectorScheduling {

        match self {

            ConnectorVariant::UserDefined(c) => c.run(scheduler_ctx, comp_ctx),

            ConnectorVariant::Native(c) => c.run(scheduler_ctx, comp_ctx),

        }

    }

}

pub(crate) struct ScheduledConnector {

    pub connector: ConnectorVariant, // access by connector

    pub ctx_fancy: ComponentCtxFancy,

    pub public: ConnectorPublic, // accessible by all schedulers and connectors

    pub router: ControlMessageHandler,

    pub shutting_down: bool,

}

// -----------------------------------------------------------------------------

// Runtime

// -----------------------------------------------------------------------------

/// Externally facing runtime.

pub struct Runtime {

    inner: Arc<RuntimeInner>,

}

impl Runtime {

    pub fn new(num_threads: u32, protocol_description: ProtocolDescription) -> Runtime {

        // Setup global state

        assert!(num_threads > 0, "need a thread to run connectors");

        let runtime_inner = Arc::new(RuntimeInner{

            protocol_description,

            port_counter: AtomicU32::new(0),

            connectors: RwLock::new(ConnectorStore::with_capacity(32)),

            connector_queue: Mutex::new(VecDeque::with_capacity(32)),

            schedulers: Mutex::new(Vec::new()),

            scheduler_notifier: Condvar::new(),

            active_connectors: AtomicU32::new(0),

            active_interfaces: AtomicU32::new(1), // this `Runtime` instance

            should_exit: AtomicBool::new(false),

        });

        // Launch threads

        {

            let mut schedulers = Vec::with_capacity(num_threads as usize);

            for thread_index in 0..num_threads {

                let cloned_runtime_inner = runtime_inner.clone();

                let thread = thread::Builder::new()

                    .name(format!("thread-{}", thread_index))

                    .spawn(move || {

                        let mut scheduler = Scheduler::new(cloned_runtime_inner, thread_index);

                        scheduler.run();

                    })

                    .unwrap();

                schedulers.push(thread);

            }

            let mut lock = runtime_inner.schedulers.lock().unwrap();

            *lock = schedulers;

        }

        // Return runtime

        return Runtime{ inner: runtime_inner };

    }

    /// Returns a new interface through which channels and connectors can be

    /// created.

    pub fn create_interface(&self) -> ApplicationInterface {

        self.inner.increment_active_interfaces();

        let (connector, mut interface) = ConnectorApplication::new(self.inner.clone());

        let connector_key = self.inner.create_interface_component(connector);

        interface.set_connector_id(connector_key.downcast());

        // Note that we're not scheduling. That is done by the interface in case

        // it is actually needed.

        return interface;

    }

}

impl Drop for Runtime {

    fn drop(&mut self) {

        self.inner.decrement_active_interfaces();

        let mut lock = self.inner.schedulers.lock().unwrap();

        for handle in lock.drain(..) {

            handle.join().unwrap();

        }

    }

}

// -----------------------------------------------------------------------------

// RuntimeInner

// -----------------------------------------------------------------------------

pub(crate) struct RuntimeInner {

    // Protocol

    pub(crate) protocol_description: ProtocolDescription,

    // Regular counter for port IDs

    port_counter: AtomicU32,

    // Storage of connectors and the work queue

    connectors: RwLock<ConnectorStore>,

    connector_queue: Mutex<VecDeque<ConnectorKey>>,

    schedulers: Mutex<Vec<JoinHandle<()>>>,

    // Conditions to determine whether the runtime can exit

    scheduler_notifier: Condvar,  // coupled to mutex on `connector_queue`.

    // TODO: Figure out if we can simply merge the counters?

    active_connectors: AtomicU32, // active connectors (if sleeping, then still considered active)

    active_interfaces: AtomicU32, // active API interfaces that can add connectors/channels

    should_exit: AtomicBool,

}

impl RuntimeInner {

    // --- Managing the components queued for execution

    /// Wait until there is a connector to run. If there is one, then `Some`

    /// will be returned. If there is no more work, then `None` will be

    /// returned.

    pub(crate) fn wait_for_work(&self) -> Option<ConnectorKey> {

        let mut lock = self.connector_queue.lock().unwrap();

        while lock.is_empty() && !self.should_exit.load(Ordering::Acquire) {

            lock = self.scheduler_notifier.wait(lock).unwrap();

        }

        return lock.pop_front();

    }

    pub(crate) fn push_work(&self, key: ConnectorKey) {

        let mut lock = self.connector_queue.lock().unwrap();

        lock.push_back(key);

        self.scheduler_notifier.notify_one();

    }

    // --- Creating/using ports

    /// Creates a new port pair. Note that these are stored globally like the

    /// connectors are. Ports stored by components belong to those components.

    pub(crate) fn create_channel(&self, creating_connector: ConnectorId) -> (Port, Port) {

        use port::{PortIdLocal, PortKind};

        let getter_id = self.port_counter.fetch_add(2, Ordering::SeqCst);

        let putter_id = PortIdLocal::new(getter_id + 1);

        let getter_id = PortIdLocal::new(getter_id);

        let getter_port = Port{

            self_id: getter_id,

            peer_id: putter_id,

            kind: PortKind::Getter,

            state: PortState::Open,

            peer_connector: creating_connector,

        };

        let putter_port = Port{

            self_id: putter_id,

            peer_id: getter_id,

            kind: PortKind::Putter,

            state: PortState::Open,

            peer_connector: creating_connector,

        };

        return (getter_port, putter_port);

    }

    /// Sends a message to a particular connector. If the connector happened to

    /// be sleeping then it will be scheduled for execution.

    pub(crate) fn send_message(&self, target_id: ConnectorId, message: Message) {

        let target = self.get_component_public(target_id);

        target.inbox.insert_message(message);

        let should_wake_up = target.sleeping

            .compare_exchange(true, false, Ordering::SeqCst, Ordering::Acquire)

            .is_ok();

        if should_wake_up {

            let key = unsafe{ ConnectorKey::from_id(target_id) };

            self.push_work(key);

        }

    }

    // --- Creating/retrieving/destroying components

    /// Creates an initially sleeping application connector.

    fn create_interface_component(&self, component: ConnectorApplication) -> ConnectorKey {

        // Initialize as sleeping, as it will be scheduled by the programmer.

        let mut lock = self.connectors.write().unwrap();

        let key = lock.create(ConnectorVariant::Native(Box::new(component)), true);

        self.increment_active_components();

        return key;

    }

    /// Creates a new PDL component. This function just creates the component.

    /// If you create it initially awake, then you must add it to the work

    /// queue. Other aspects of correctness (i.e. setting initial ports) are

    /// relinquished to the caller!

    pub(crate) fn create_pdl_component(&self, connector: ConnectorPDL, initially_sleeping: bool) -> ConnectorKey {

        // Create as not sleeping, as we'll schedule it immediately

        let key = {

            let mut lock = self.connectors.write().unwrap();

            lock.create(ConnectorVariant::UserDefined(connector), initially_sleeping)

        };

        self.increment_active_components();

        return key;

    }

    #[inline]

    pub(crate) fn get_component_private(&self, connector_key: &ConnectorKey) -> &'static mut ScheduledConnector {

        let lock = self.connectors.read().unwrap();

        return lock.get_private(connector_key);

    }

    #[inline]

    pub(crate) fn get_component_public(&self, connector_id: ConnectorId) -> &'static ConnectorPublic {

        let lock = self.connectors.read().unwrap();

        return lock.get_public(connector_id);

    }

    pub(crate) fn destroy_component(&self, connector_key: ConnectorKey) {

        let mut lock = self.connectors.write().unwrap();

        lock.destroy(connector_key);

        self.decrement_active_components();

    }

    // --- Managing exit condition

    #[inline]

    pub(crate) fn increment_active_interfaces(&self) {

        let _old_num = self.active_interfaces.fetch_add(1, Ordering::SeqCst);

        println!("DEBUG: Incremented active interfaces to {}", _old_num + 1);

        debug_assert_ne!(_old_num, 0); // once it hits 0, it stays zero

    }

    pub(crate) fn decrement_active_interfaces(&self) {

        let old_num = self.active_interfaces.fetch_sub(1, Ordering::SeqCst);

        println!("DEBUG: Decremented active interfaces to {}", old_num - 1);

        debug_assert!(old_num > 0);

        if old_num == 1 { // such that active interfaces is now 0

            let num_connectors = self.active_connectors.load(Ordering::Acquire);

            if num_connectors == 0 {

                self.signal_for_shutdown();

            }

        }

    }

    #[inline]

    fn increment_active_components(&self) {

        let _old_num = self.active_connectors.fetch_add(1, Ordering::SeqCst);

        println!("DEBUG: Incremented components to {}", _old_num + 1);

    }

    fn decrement_active_components(&self) {

        let old_num = self.active_connectors.fetch_sub(1, Ordering::SeqCst);

        println!("DEBUG: Decremented components to {}", old_num - 1);

        debug_assert!(old_num > 0);

        if old_num == 1 { // such that we have no more active connectors (for now!)

            let num_interfaces = self.active_interfaces.load(Ordering::Acquire);

            if num_interfaces == 0 {

                self.signal_for_shutdown();

            }

        }

    }

    #[inline]

    fn signal_for_shutdown(&self) {

        debug_assert_eq!(self.active_interfaces.load(Ordering::Acquire), 0);

        debug_assert_eq!(self.active_connectors.load(Ordering::Acquire), 0);

        println!("DEBUG: Signaling for shutdown");

        let _lock = self.connector_queue.lock().unwrap();

        let should_signal = self.should_exit

            .compare_exchange(false, true, Ordering::SeqCst, Ordering::Acquire)

            .is_ok();

        if should_signal {

            println!("DEBUG: Notifying all waiting schedulers");

            self.scheduler_notifier.notify_all();

        }

    }

}

// TODO: Come back to this at some point

unsafe impl Send for RuntimeInner {}

unsafe impl Sync for RuntimeInner {}

// -----------------------------------------------------------------------------

// ConnectorStore

// -----------------------------------------------------------------------------

struct ConnectorStore {

    // Freelist storage of connectors. Storage should be pointer-stable as

    // someone might be mutating the vector while we're executing one of the

    // connectors.

    connectors: RawVec<*mut ScheduledConnector>,

    free: Vec<usize>,

}

impl ConnectorStore {

    fn with_capacity(capacity: usize) -> Self {

        Self {

            connectors: RawVec::with_capacity(capacity),

            free: Vec::with_capacity(capacity),

        }

    }

    /// Retrieves public part of connector - accessible by many threads at once.

    fn get_public(&self, id: ConnectorId) -> &'static ConnectorPublic {

        unsafe {

            let connector = self.connectors.get(id.0 as usize);

            debug_assert!(!connector.is_null());

            return &(**connector).public;

        }

use std::sync::Arc;

use super::*;

use crate::{PortId, ProtocolDescription};

use crate::common::Id;

use crate::protocol::eval::*;

const NUM_INSTANCES: u32 = 10;  // number of test instances constructed

const NUM_LOOPS: u32 = 10;      // number of loops within a single test (not used by all tests)

fn create_runtime(pdl: &str) -> Runtime {

    let protocol = ProtocolDescription::parse(pdl.as_bytes()).expect("parse pdl");

    let runtime = Runtime::new(NUM_THREADS, protocol);

    return runtime;

}

fn run_test_in_runtime<F: Fn(&mut ApplicationInterface)>(pdl: &str, constructor: F) {

    let protocol = ProtocolDescription::parse(pdl.as_bytes())

        .expect("parse PDL");

    let runtime = Runtime::new(NUM_THREADS, protocol);

    let mut api = runtime.create_interface();

    for _ in 0..NUM_INSTANCES {

        constructor(&mut api);

    }

    // Wait until done :)

}

#[test]

fn test_put_and_get() {

    const CODE: &'static str = "

    primitive putter(out<bool> sender, u32 loops) {

        u32 index = 0;

        while (index < loops) {

            synchronous {

                print(\"putting!\");

                put(sender, true);

            }

            index += 1;

        }

    }

    primitive getter(in<bool> receiver, u32 loops) {

        u32 index = 0;

        while (index < loops) {

            synchronous {

                print(\"getting!\");

                auto result = get(receiver);

                assert(result);

            }

            index += 1;

        }

    }

    ";

    run_test_in_runtime(CODE, |api| {

        let channel = api.create_channel();

        api.create_connector("", "putter", ValueGroup::new_stack(vec![

            Value::Output(PortId(Id{ connector_id: 0, u32_suffix: channel.putter_id.index })),

            Value::UInt32(NUM_LOOPS)

        ])).expect("create putter");

        api.create_connector("", "getter", ValueGroup::new_stack(vec![

            Value::Input(PortId(Id{ connector_id: 0, u32_suffix: channel.getter_id.index })),

            Value::UInt32(NUM_LOOPS)

        ])).expect("create getter");

    });

}