#[macro_use]

mod macros;

mod common;

mod protocol;

mod runtime;

mod collections;

pub use common::{ConnectorId, EndpointPolarity, Payload, Polarity, PortId};

pub use protocol::ProtocolDescription;

pub use runtime::{error, Connector, DummyLogger, FileLogger, VecLogger};

// TODO: Remove when not benchmarking

pub use protocol::input_source::InputSource;

pub use protocol::ast::Heap;

#[cfg(feature = "ffi")]

pub mod ffi;

// Structure of module

mod branch;

mod native;

mod port;

mod scheduler;

mod consensus;

mod inbox;

#[cfg(test)] mod tests;

mod connector;

// 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 connector::{ConnectorPDL, ConnectorPublic, ConnectorScheduling};

use scheduler::{Scheduler, ComponentCtx, SchedulerCtx, ControlMessageHandler};

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

use inbox::Message;

use port::{ChannelId, Port, PortState};

/// 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, PartialOrd, Ord)]

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 ComponentCtx) -> 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 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

        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);

        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);

    }

    fn decrement_active_components(&self) {

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

        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);

        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 {

            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 {

            debug_assert!(!self.free.contains(&(id.0 as usize)));

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

            debug_assert!(!connector.is_null());

            return &(**connector).public;

        }

    }

    /// Retrieves private part of connector - accessible by one thread at a

    /// time.

    fn get_private(&self, key: &ConnectorKey) -> &'static mut ScheduledConnector {

        unsafe {

            debug_assert!(!self.free.contains(&(key.index as usize)));

            let connector = self.connectors.get_mut(key.index as usize);

            debug_assert!(!connector.is_null());

            return &mut (**connector);

        }

    }

    /// Creates a new connector. Caller should ensure ports are set up correctly

    /// and the connector is queued for execution if needed.

    fn create(&mut self, connector: ConnectorVariant, initially_sleeping: bool) -> ConnectorKey {

        let mut connector = ScheduledConnector {

            connector,

            public: ConnectorPublic::new(initially_sleeping),

            router: ControlMessageHandler::new(),

            shutting_down: false,

        };

        let index;

        let key;

        if self.free.is_empty() {

            // No free entries, allocate new entry

            index = self.connectors.len();

            key = ConnectorKey{ index: index as u32 };

            let connector = Box::into_raw(Box::new(connector));

            self.connectors.push(connector);

        } else {

            // Free spot available

            index = self.free.pop().unwrap();

            key = ConnectorKey{ index: index as u32 };

            unsafe {

                let target = self.connectors.get_mut(index);

                std::ptr::write(*target, connector);

            }

        }

        return key;

    }

    /// Destroys a connector. Caller should make sure it is not scheduled for

    /// execution. Otherwise one experiences "bad stuff" (tm).

    fn destroy(&mut self, key: ConnectorKey) {

        unsafe {

            let target = self.connectors.get_mut(key.index as usize);

            std::ptr::drop_in_place(*target);

            // Note: but not deallocating!

        }

        self.free.push(key.index as usize);

    }

}

impl Drop for ConnectorStore {

    fn drop(&mut self) {

        // Everything in the freelist already had its destructor called, so only

        // has to be deallocated

        for free_idx in self.free.iter().copied() {

            unsafe {

                let memory = self.connectors.get_mut(free_idx);

                let layout = std::alloc::Layout::for_value(&**memory);

                std::alloc::dealloc(*memory as *mut u8, layout);

                // mark as null for the remainder

                *memory = std::ptr::null_mut();

            }

        }

        // With the deallocated stuff marked as null, clear the remainder that

        // is not null

        for idx in 0..self.connectors.len() {

            unsafe {

                let memory = *self.connectors.get_mut(idx);

                if !memory.is_null() {

                    let _ = Box::from_raw(memory); // take care of deallocation, bit dirty, but meh

                }

            }

        }

    }

}

use std::collections::VecDeque;

use std::sync::Arc;

use std::sync::atomic::Ordering;

use super::{ScheduledConnector, RuntimeInner, ConnectorId, ConnectorKey};

use super::port::{Port, PortState, PortIdLocal};

use super::native::Connector;

use super::branch::{BranchId};

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

use super::inbox::{Message, DataMessage, ControlMessage, ControlContent};

// Because it contains pointers we're going to do a copy by value on this one

#[derive(Clone, Copy)]

pub(crate) struct SchedulerCtx<'a> {

    pub(crate) runtime: &'a RuntimeInner

}

pub(crate) struct Scheduler {

    runtime: Arc<RuntimeInner>,

    scheduler_id: u32,

}

impl Scheduler {

    pub fn new(runtime: Arc<RuntimeInner>, scheduler_id: u32) -> Self {

        return Self{ runtime, scheduler_id };

    }

    pub fn run(&mut self) {

        // Setup global storage and workspaces that are reused for every

        // connector that we run

        'thread_loop: loop {

            // Retrieve a unit of work

            self.debug("Waiting for work");

            let connector_key = self.runtime.wait_for_work();

            if connector_key.is_none() {

                // We should exit

                self.debug(" ... No more work, quitting");

                break 'thread_loop;

            }

            // We have something to do

            let connector_key = connector_key.unwrap();

            let connector_id = connector_key.downcast();

            self.debug_conn(connector_id, &format!(" ... Got work, running {}", connector_key.index));

            let scheduled = self.runtime.get_component_private(&connector_key);

            // Keep running until we should no longer immediately schedule the

            // connector.

            let mut cur_schedule = ConnectorScheduling::Immediate;

            while cur_schedule == ConnectorScheduling::Immediate {

                self.handle_inbox_messages(scheduled);

                // Run the main behaviour of the connector, depending on its

                // current state.

                if scheduled.shutting_down {

                    // Nothing to do. But we're stil waiting for all our pending

                    // control messages to be answered.

                    self.debug_conn(connector_id, &format!("Shutting down, {} Acks remaining", scheduled.router.num_pending_acks()));

                    if scheduled.router.num_pending_acks() == 0 {

                        // We're actually done, we can safely destroy the

                        // currently running connector

                        self.runtime.destroy_component(connector_key);

                        continue 'thread_loop;

                    } else {

                        cur_schedule = ConnectorScheduling::NotNow;

                    }

                } else {

                    self.debug_conn(connector_id, "Running ...");

                    let scheduler_ctx = SchedulerCtx{ runtime: &*self.runtime };

                    self.debug_conn(connector_id, "Finished running");

                    // Handle all of the output from the current run: messages to

                    // send and connectors to instantiate.

                    self.handle_changes_in_context(scheduled);

                    cur_schedule = new_schedule;

                }

            }

            // If here then the connector does not require immediate execution.

            // So enqueue it if requested, and otherwise put it in a sleeping

            // state.

            match cur_schedule {

                ConnectorScheduling::Immediate => unreachable!(),

                ConnectorScheduling::Later => {

                    // Simply queue it again later

                    self.runtime.push_work(connector_key);

                },

                ConnectorScheduling::NotNow => {

                    // Need to sleep, note that we are the only ones which are

                    // allows to set the sleeping state to `true`, and since

                    // we're running it must currently be `false`.

                    self.try_go_to_sleep(connector_key, scheduled);

                },

                ConnectorScheduling::Exit => {

                    // Prepare for exit. Set the shutdown flag and broadcast

                    // messages to notify peers of closing channels

                    scheduled.shutting_down = true;

                        if port.state != PortState::Closed {

                            let message = scheduled.router.prepare_closing_channel(

                                port.self_id, port.peer_id,

                                connector_id

                            );

                            self.debug_conn(connector_id, &format!("Sending message [ exit ] \n --- {:?}", message));

                            self.runtime.send_message(port.peer_connector, Message::Control(message));

                        }

                    }

                    if scheduled.router.num_pending_acks() == 0 {

                        self.runtime.destroy_component(connector_key);

                        continue 'thread_loop;

                    }

                    self.try_go_to_sleep(connector_key, scheduled);

                }

            }

        }

    }

    /// Receiving messages from the public inbox and handling them or storing

    /// them in the component's private inbox

    fn handle_inbox_messages(&mut self, scheduled: &mut ScheduledConnector) {

        while let Some(message) = scheduled.public.inbox.take_message() {

            // Check if the message has to be rerouted because we have moved the

            // target port to another component.

            self.debug_conn(connector_id, &format!("Handling message\n --- {:?}", message));

            if let Some(target_port) = Self::get_message_target_port(&message) {

                if let Some(other_component_id) = scheduled.router.should_reroute(target_port) {

                    self.debug_conn(connector_id, " ... Rerouting the message");

                    self.runtime.send_message(other_component_id, message);

                    continue;

                }

            }

            // If here, then we should handle the message

            self.debug_conn(connector_id, " ... Handling the message");

            match message {

                Message::Control(message) => {

                    match message.content {

                        ControlContent::PortPeerChanged(port_id, new_target_connector_id) => {

                            // Need to change port target

                            port.peer_connector = new_target_connector_id;

                            // Note: for simplicity we program the scheduler to always finish

                            // running a connector with an empty outbox. If this ever changes

                            // then accepting the "port peer changed" message implies we need

                            // to change the recipient of the message in the outbox.

                            // And respond with an Ack

                            let ack_message = Message::Control(ControlMessage {

                                id: message.id,

                                sending_component_id: connector_id,

                                content: ControlContent::Ack,

                            });

                            self.debug_conn(connector_id, &format!("Sending message [pp ack]\n --- {:?}", ack_message));

                            self.runtime.send_message(message.sending_component_id, ack_message);

                        },

                        ControlContent::CloseChannel(port_id) => {

                            // Mark the port as being closed

                            port.state = PortState::Closed;

                            // Send an Ack

                            let ack_message = Message::Control(ControlMessage {

                                id: message.id,

                                sending_component_id: connector_id,

                                content: ControlContent::Ack,

                            });

                            self.debug_conn(connector_id, &format!("Sending message [cc ack] \n --- {:?}", ack_message));

                            self.runtime.send_message(message.sending_component_id, ack_message);

                        },

                        ControlContent::Ack => {

                            scheduled.router.handle_ack(message.id);

                        },

                        ControlContent::Ping => {},

                    }

                },

                _ => {

                    // All other cases have to be handled by the component

                }

            }

        }

    }

    /// Handles changes to the context that were made by the component. This is

    /// the way (due to Rust's borrowing rules) that we bubble up changes in the

    /// component's state that the scheduler needs to know about (e.g. a message

    /// that the component wants to send, a port that has been added).

    fn handle_changes_in_context(&mut self, scheduled: &mut ScheduledConnector) {

        // Handling any messages that were sent

            self.debug_conn(connector_id, &format!("Sending message [outbox] \n --- {:?}", message));

            let target_component_id = match &message {

                Message::Data(content) => {

                    // Data messages are always sent to a particular port, and

                    // may end up being rerouted.

                    debug_assert_eq!(port_desc.peer_id, content.data_header.target_port);

                    if port_desc.state == PortState::Closed {

                        todo!("handle sending over a closed port")

                    }

                    port_desc.peer_connector

                },

                Message::Sync(content) => {

                    // Sync messages are always sent to a particular component,

                    // the sender must make sure it actually wants to send to

                    // the specified component (and is not using an inconsistent

                    // component ID associated with a port).

                    content.target_component_id

                },

                Message::Control(_) => {

                    unreachable!("component sending control messages directly");

                }

            };

            self.runtime.send_message(target_component_id, message);

        }

            match state_change {

                ComponentStateChange::CreatedComponent(component, initial_ports) => {

                    // Creating a new component. The creator needs to relinquish

                    // ownership of the ports that are given to the new

                    // component. All data messages that were intended for that

                    // port also needs to be transferred.

                    let new_key = self.runtime.create_pdl_component(component, false);

                    let new_connector = self.runtime.get_component_private(&new_key);

                    for port_id in initial_ports {

                        // Transfer messages associated with the transferred port

                        let mut message_idx = 0;

                            if Self::get_message_target_port(message) == Some(port_id) {

                                // Need to transfer this message

                            } else {

                                message_idx += 1;

                            }

                        }

                        // Transfer the port itself

                            .position(|v| v.self_id == port_id)

                            .unwrap();

                        // Notify the peer that the port has changed

                        let reroute_message = scheduled.router.prepare_reroute(

                        );

                        self.debug_conn(connector_id, &format!("Sending message [newcon]\n --- {:?}", reroute_message));

                        self.runtime.send_message(port.peer_connector, Message::Control(reroute_message));

                    }

                    // Schedule new connector to run

                    self.runtime.push_work(new_key);

                },

                ComponentStateChange::CreatedPort(port) => {

                },

                ComponentStateChange::ChangedPort(port_change) => {

                    if port_change.is_acquired {

                    } else {

                            .iter()

                            .position(|v| v.self_id == port_change.port.self_id)

                            .unwrap();

                    }

                }

            }

        }

        // Finally, check if we just entered or just left a sync region

                // Just entered sync region

            } else {

                // Just left sync region. So clear inbox

            }

        }

    }

    fn try_go_to_sleep(&self, connector_key: ConnectorKey, connector: &mut ScheduledConnector) {

        debug_assert_eq!(connector.public.sleeping.load(Ordering::Acquire), false);

        // This is the running connector, and only the running connector may

        // decide it wants to sleep again.

        connector.public.sleeping.store(true, Ordering::Release);

        // But due to reordering we might have received messages from peers who

        // did not consider us sleeping. If so, then we wake ourselves again.

        if !connector.public.inbox.is_empty() {

            // Try to wake ourselves up (needed because someone might be trying

            // the exact same atomic compare-and-swap at this point in time)

            let should_wake_up_again = connector.public.sleeping

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

                .is_ok();

            if should_wake_up_again {

                self.runtime.push_work(connector_key)

            }

        }

    }

    #[inline]

    fn get_message_target_port(message: &Message) -> Option<PortIdLocal> {

        match message {

            Message::Data(data) => return Some(data.data_header.target_port),

            Message::Sync(_) => {},

            Message::Control(control) => {

                match &control.content {

                    ControlContent::PortPeerChanged(port_id, _) => return Some(*port_id),

                    ControlContent::CloseChannel(port_id) => return Some(*port_id),

                    ControlContent::Ping | ControlContent::Ack => {},

                }

            },

        }

        return None

    }

    // TODO: Remove, this is debugging stuff

    fn debug(&self, message: &str) {

        // println!("DEBUG [thrd:{:02} conn:  ]: {}", self.scheduler_id, message);

    }

    fn debug_conn(&self, conn: ConnectorId, message: &str) {

        // println!("DEBUG [thrd:{:02} conn:{:02}]: {}", self.scheduler_id, conn.0, message);

    }

}

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

// ComponentCtx

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

enum ComponentStateChange {

    CreatedComponent(ConnectorPDL, Vec<PortIdLocal>),

    CreatedPort(Port),

    ChangedPort(ComponentPortChange),

}

#[derive(Clone)]

pub(crate) struct ComponentPortChange {

    pub is_acquired: bool, // otherwise: released

    pub port: Port,

}

/// The component context (better name may be invented). This was created

/// because part of the component's state is managed by the scheduler, and part

/// of it by the component itself. When the component starts a sync block or

/// exits a sync block the partially managed state by both component and

/// scheduler need to be exchanged.

pub(crate) struct ComponentCtx {

    // Mostly managed by the scheduler

    pub(crate) id: ConnectorId,

    ports: Vec<Port>,

    inbox_messages: Vec<Message>, // never control or ping messages

    inbox_len_read: usize,

    // Submitted by the component

    is_in_sync: bool,

    changed_in_sync: bool,

    outbox: VecDeque<Message>,

    state_changes: VecDeque<ComponentStateChange>,

    // Workspaces that may be used by components to (generally) prevent

    // allocations. Be a good scout and leave it empty after you've used it.

    // TODO: Move to scheduler ctx, this is the wrong place

    pub workspace_ports: Vec<PortIdLocal>,

    pub workspace_branches: Vec<BranchId>,

}

impl ComponentCtx {

    pub(crate) fn new_empty() -> Self {

        return Self{

            id: ConnectorId::new_invalid(),

            ports: Vec::new(),

            inbox_messages: Vec::new(),

            inbox_len_read: 0,

            is_in_sync: false,

            changed_in_sync: false,