#include <time.h>

#include "../../reowolf.h"

#include "../utility.c"

int main(int argc, char** argv) {

	int i, proto_components;

	proto_components = atoi(argv[1]);

	printf("proto_components: %d\n", proto_components);

	const unsigned char pdl[] = 

	"primitive trivial_loop() {   "

	"    while(true) synchronous{}"

	"}                            "

	;

	Arc_ProtocolDescription * pd = protocol_description_parse(pdl, sizeof(pdl)-1);

	char logpath[] = "./bench_4.txt";

	Connector * c = connector_new_logging(pd, logpath, sizeof(logpath)-1);

	for (i=0; i<proto_components; i++) {

		printf("Error str `%s`\n", reowolf_error_peek(NULL));

	}

	connector_connect(c, -1);

	printf("Error str `%s`\n", reowolf_error_peek(NULL));

	clock_t begin = clock();

	for (i=0; i<1000000; i++) {

		connector_sync(c, -1);

	}

	clock_t end = clock();

	double time_spent = (double)(end - begin) / CLOCKS_PER_SEC;

	printf("Time taken: %f\n", time_spent);

	return 0;

}

) -> *mut Arc<ProtocolDescription> {

    StoredError::tl_clear();

    match ProtocolDescription::parse(&*slice_from_raw_parts(pdl, pdl_len)) {

        Ok(new) => Box::into_raw(Box::new(Arc::new(new))),

        Err(err) => {

            StoredError::tl_bytes_store(err.as_bytes());

            std::ptr::null_mut()

        }

    }

}

/// Destroys the given initialized protocol description and frees its resources.

#[no_mangle]

pub unsafe extern "C" fn protocol_description_destroy(pd: *mut Arc<ProtocolDescription>) {

    drop(Box::from_raw(pd))

}

/// Given an initialized protocol description, initializes `out` with a clone which can be independently created or destroyed.

#[no_mangle]

pub unsafe extern "C" fn protocol_description_clone(

    pd: &Arc<ProtocolDescription>,

) -> *mut Arc<ProtocolDescription> {

    Box::into_raw(Box::new(pd.clone()))

}

///////////////////// CONNECTOR //////////////////////////

#[no_mangle]

pub unsafe extern "C" fn connector_new_logging_with_id(

    pd: &Arc<ProtocolDescription>,

    path_ptr: *const u8,

    path_len: usize,

    connector_id: ConnectorId,

) -> *mut Connector {

    StoredError::tl_clear();

    let path_bytes = &*slice_from_raw_parts(path_ptr, path_len);

    let path_str = match std::str::from_utf8(path_bytes) {

        Ok(path_str) => path_str,

        Err(err) => {

            StoredError::tl_debug_store(&err);

            return std::ptr::null_mut();

        }

    };

    match std::fs::File::create(path_str) {

        Ok(file) => {

            let file_logger = Box::new(FileLogger::new(connector_id, file));

            let c = Connector::new(file_logger, pd.clone(), connector_id);

            Box::into_raw(Box::new(c))

        }

        Err(err) => {

            StoredError::tl_debug_store(&err);

            std::ptr::null_mut()

        }

    }

}

#[no_mangle]

pub unsafe extern "C" fn connector_new_logging(

    pd: &Arc<ProtocolDescription>,

    path_ptr: *const u8,

    path_len: usize,

) -> *mut Connector {

    connector_new_logging_with_id(pd, path_ptr, path_len, Connector::random_id())

}

#[no_mangle]

pub unsafe extern "C" fn connector_print_debug(connector: &mut Connector) {

    println!("Debug print dump {:#?}", connector);

}

/// Initializes `out` with a new connector using the given protocol description as its configuration.

/// The connector uses the given (internal) connector ID.

#[no_mangle]

pub unsafe extern "C" fn connector_new(pd: &Arc<ProtocolDescription>) -> *mut Connector {

    let c = Connector::new(Box::new(DummyLogger), pd.clone(), Connector::random_id());

    Box::into_raw(Box::new(c))

}

/// Destroys the given a pointer to the connector on the heap, freeing its resources.

/// Usable in {setup, communication} states.

#[no_mangle]

pub unsafe extern "C" fn connector_destroy(connector: *mut Connector) {

    drop(Box::from_raw(connector))

}

/// Given an initialized connector in setup or connecting state,

/// - Creates a new directed port pair with logical channel putter->getter,

/// - adds the ports to the native component's interface,

/// - and returns them using the given out pointers.

/// Usable in {setup, communication} states.

#[no_mangle]

pub unsafe extern "C" fn connector_add_port_pair(

    connector: &mut Connector,

    out_putter: *mut PortId,

    out_getter: *mut PortId,

) {

    let [o, i] = connector.new_port_pair();

}

/// Given

/// - an initialized connector in setup or connecting state,

/// - a string slice for the component's identifier in the connector's configured protocol description,

/// - a set of ports (represented as a slice; duplicates are ignored) in the native component's interface,

/// the connector creates a new (internal) protocol component C, such that the set of native ports are moved to C.

/// Usable in {setup, communication} states.

#[no_mangle]

pub unsafe extern "C" fn connector_add_component(

    connector: &mut Connector,

    ident_ptr: *const u8,

    ident_len: usize,

    ports_ptr: *const PortId,

    ports_len: usize,

) -> c_int {

    StoredError::tl_clear();

    match connector.add_component(

        &*slice_from_raw_parts(ident_ptr, ident_len),

        &*slice_from_raw_parts(ports_ptr, ports_len),

    ) {

        Ok(()) => RW_OK,

        Err(err) => {

            StoredError::tl_debug_store(&err);

            RW_TL_ERR

        }

    }

}

/// Given

/// - an initialized connector in setup or connecting state,

/// - a utf-8 encoded socket address,

/// - the logical polarity of P,

/// - the "physical" polarity in {Active, Passive} of the endpoint through which P's peer will be discovered,

/// returns P, a port newly added to the native interface.

#[no_mangle]

pub unsafe extern "C" fn connector_add_net_port(

    connector: &mut Connector,

    port: *mut PortId,

    addr: FfiSocketAddr,

    port_polarity: Polarity,

    endpoint_polarity: EndpointPolarity,

) -> c_int {

    StoredError::tl_clear();

    match connector.new_net_port(port_polarity, addr.into(), endpoint_polarity) {

        Ok(p) => {

            if !port.is_null() {

                port.write(p);

            }

            RW_OK

        }

        Err(err) => {

            StoredError::tl_debug_store(&err);

            RW_TL_ERR

        }

    }

}

/// Given

/// - an initialized connector in setup or connecting state,

/// - a utf-8 encoded BIND socket addresses (i.e., "local"),

/// - a utf-8 encoded CONNECT socket addresses (i.e., "peer"),

/// returns [P, G] via out pointers [putter, getter],

/// - where P is a Putter port that sends messages into the socket

/// - where G is a Getter port that recvs messages from the socket

#[no_mangle]

pub unsafe extern "C" fn connector_add_udp_mediator_component(

    connector: &mut Connector,

    putter: *mut PortId,

    getter: *mut PortId,

    local_addr: FfiSocketAddr,

    peer_addr: FfiSocketAddr,

) -> c_int {

    StoredError::tl_clear();

    match connector.new_udp_mediator_component(local_addr.into(), peer_addr.into()) {

        Ok([p, g]) => {

            if !putter.is_null() {

                putter.write(p);

            }

            if !getter.is_null() {

                getter.write(g);

            }

            RW_OK

        }

        Err(err) => {

            StoredError::tl_debug_store(&err);

            RW_TL_ERR

        }

    }

}

/// Connects this connector to the distributed system of connectors reachable through endpoints,

/// Usable in setup state, and changes the state to communication.

#[no_mangle]

pub unsafe extern "C" fn connector_connect(

    connector: &mut Connector,

    }

    // drain old_local to new_local, visiting all new additions to old_local

    pub(crate) fn iter_new_local_make_old(&mut self) -> impl Iterator<Item = Predicate> + '_ {

        let Self { old_local, new_local, .. } = self;

        new_local.drain().map(move |local| {

            // rely on invariant: empty intersection between old and new local sets

            assert!(old_local.insert(local.clone()));

            local

        })

    }

    // insert a solution for the given subtree ID,

    // AND update new_local to include any solutions that become 

    // possible as a result of this new addition  

    pub(crate) fn submit_and_digest_subtree_solution(

        &mut self,

        cu: &mut impl CuUndecided,

        subtree_id: SubtreeId,

        predicate: Predicate,

    ) {

        log!(cu.logger(), "++ new component solution {:?} {:?}", subtree_id, &predicate);

        let Self { subtree_solutions, new_local, old_local, subtree_id_to_index } = self;

        let index = subtree_id_to_index[&subtree_id];

        let was_new = subtree_solutions[index].insert(predicate.clone());

        if was_new {

            // This is a newly-added solution! update new_local

            // consider ALL consistent combinations of one element from each solution set

            // to our right or left in the solution-set vector

            // but with THIS PARTICULAR predicate from our own index.

            let left = 0..index;

            let right = (index + 1)..subtree_solutions.len();

            // iterator over SETS of solutions, one for every component except `subtree_id` (me)

            let set_visitor = left.chain(right).map(|index| &subtree_solutions[index]);

            // Recursively enumerate all solutions matching the description above,

            Self::elaborate_into_new_local_rec(cu, predicate, set_visitor, old_local, new_local);

        }

    }

    // Recursively build local solutions for this connector,

    // see `submit_and_digest_subtree_solution`

    fn elaborate_into_new_local_rec<'a, 'b>(

        cu: &mut impl CuUndecided,

        partial: Predicate,

        mut set_visitor: impl Iterator<Item = &'b HashSet<Predicate>> + Clone,

        old_local: &'b HashSet<Predicate>,

        new_local: &'a mut HashSet<Predicate>,

    ) {

        if let Some(set) = set_visitor.next() {

            // incomplete solution. keep recursively creating combined solutions

            for pred in set.iter() {

                if let Some(elaborated) = pred.union_with(&partial) {

                    Self::elaborate_into_new_local_rec(

                        cu,

                        elaborated,

                        set_visitor.clone(),

                        old_local,

                        new_local,

                    )

                }

            }

        } else {

            // recursive stop condition. This is a solution for this connector...

            if !old_local.contains(&partial) {

                // ... and it hasn't been found before

                log!(cu.logger(), "storing NEW LOCAL SOLUTION {:?}", &partial);

                new_local.insert(partial);

            }

        }

    }

}

impl NonsyncProtoContext<'_> {

    // Facilitates callback from the component to the connector runtime,

    // creating a new component and changing the given port's ownership to that

    // of the new component.

    pub(crate) fn new_component(&mut self, moved_ports: HashSet<PortId>, state: ComponentState) {

        // Sanity check! The moved ports are owned by this component to begin with

        for port in moved_ports.iter() {

            assert_eq!(

                self.proto_component_id,

                self.ips.port_info.map.get(port).unwrap().owner

            );

        }

        // Create the new component, and schedule it to be run

        let new_cid = self.ips.id_manager.new_component_id();

        log!(

            self.logger,

            "Component {:?} added new component {:?} with state {:?}, moving ports {:?}",

            self.proto_component_id,

            new_cid,

            &state,

            &moved_ports

        );

        self.unrun_components.push((new_cid, state));

        // Update the ownership of the moved ports

        for port in moved_ports.iter() {

            self.ips.port_info.map.get_mut(port).unwrap().owner = new_cid;

        }

    }

    // Facilitates callback from the component to the connector runtime,

    // creating a new port-pair connected by an memory channel

    pub(crate) fn new_port_pair(&mut self) -> [PortId; 2] {

        // adds two new associated ports, related to each other, and exposed to the proto component

        let mut new_cid_fn = || self.ips.id_manager.new_port_id();

        let [o, i] = [new_cid_fn(), new_cid_fn()];

        self.ips.port_info.map.insert(

            o,

            PortInfo {

                route: Route::LocalComponent,

                peer: Some(i),

                polarity: Putter,

                owner: self.proto_component_id,

            },

        );

        self.ips.port_info.map.insert(

            i,

            PortInfo {

                route: Route::LocalComponent,

                peer: Some(o),

                polarity: Getter,

                owner: self.proto_component_id,

            },

        );

        log!(

            self.logger,

            "Component {:?} port pair (out->in) {:?} -> {:?}",

            self.proto_component_id,

            o,

            i

        );

        [o, i]

    }

}

impl SyncProtoContext<'_> {

    // The component calls the runtime back, inspecting whether it's associated

    // preidcate has already determined a (speculative) value for the given port's firing variable.

    pub(crate) fn is_firing(&mut self, port: PortId) -> Option<bool> {

        let var = self.rctx.ips.port_info.spec_var_for(port);

        self.predicate.query(var).map(SpecVal::is_firing)

    }

    // The component calls the runtime back, trying to inspect a port's message

    pub(crate) fn read_msg(&mut self, port: PortId) -> Option<&Payload> {

        let maybe_msg = self.branch_inner.inbox.get(&port);

        if maybe_msg.is_some() {

            // Make a note that this component has received

            // this port's message 1+ times this round

            self.branch_inner.did_put_or_get.insert(port);

        }

        maybe_msg

    }

    // NOT CURRENTLY USED

    // Once this component has injected a new nondeterministic branch with

    // SyncBlocker::NondetChoice, this is how the component retrieves it.

    // (Two step process necessary to get around mutable access rules,

    //  as injection of the nondeterministic choice modifies the

    //  branch predicate, forks the branch, etc.)

    pub(crate) fn take_choice(&mut self) -> Option<u16> {

        self.branch_inner.untaken_choice.take()

    }

}

impl<'a, K: Eq + Hash + 'static, V: 'static> CyclicDrainer<'a, K, V> {

    fn new(

        input: &'a mut HashMap<K, V>,

        swap: &'a mut HashMap<K, V>,

        output: &'a mut HashMap<K, V>,

    ) -> Self {

        Self { input, inner: CyclicDrainerInner { swap, output } }

    }

    // This hides the ugliness of facilitating a memory-safe cyclic drain.

    // A "drain" would refer to a procedure that empties the input and populates the output.

    // It's "cyclic" because the processing function can also populate the input.

    // Making this memory safe requires an additional temporary storage, such that

    // the input can safely be drained and populated concurrently.

    fn cyclic_drain<E>(

        self,

        mut func: impl FnMut(K, V, CyclicDrainerInner<'_, K, V>) -> Result<(), E>,

    ) -> Result<(), E> {

        let Self { input, inner: CyclicDrainerInner { swap, output } } = self;

        while !input.is_empty() {

            for (k, v) in input.drain() {

                // func is the user-provided callback function, which consumes an element

                // as its drained from the input

                func(k, v, CyclicDrainerInner { swap, output })?

            }

            std::mem::swap(input, swap);

        }

        Ok(())

    }

}

impl<'a, K: Eq + Hash, V> CyclicDrainerInner<'a, K, V> {

    // Add this key-value pair to be yielded by the drainer later

    fn add_input(&mut self, k: K, v: V) {

        self.swap.insert(k, v);

    }

    // Add this key-value pair as an output of the drainer

    fn add_output(&mut self, k: K, v: V) {

        self.output.insert(k, v);

    }

}

// NetEndpoint annotated with the ID of the port that receives payload

// messages received through the endpoint. This approach assumes that NetEndpoints

// DO NOT multiplex port->port channels, and so a mapping such as this is possible.

// As a result, the messages themselves don't need to carry the PortID with them.

#[derive(Debug)]

struct NetEndpointExt {

    net_endpoint: NetEndpoint,

    getter_for_incoming: PortId,

}

// Endpoint for a "raw" UDP endpoint. Corresponds to the "Udp Mediator Component"

// described in the literature.

// It acts as an endpoint by receiving messages via the poller etc. (managed by EndpointManager),

// It acts as a native component by managing a (speculative) set of payload messages (an outbox,

//  protecting the peer on the other side of the network).

#[derive(Debug)]

struct UdpEndpointExt {

    sock: UdpSocket, // already bound and connected

    received_this_round: bool,

    outgoing_payloads: HashMap<Predicate, Payload>,

    getter_for_incoming: PortId,

}

// Meta-data for the connector: its role in the consensus tree.

#[derive(Debug)]

struct Neighborhood {

    parent: Option<usize>,

    children: VecSet<usize>,

}

// Manages the connector's ID, and manages allocations for connector/port IDs.

#[derive(Debug, Clone)]

struct IdManager {

    connector_id: ConnectorId,

    port_suffix_stream: U32Stream,

    component_suffix_stream: U32Stream,

}

// Newtype wrapper around a byte buffer, used for UDP mediators to receive incoming datagrams.

struct UdpInBuffer {

    byte_vec: Vec<u8>,

}

// A generator of speculative variables. Created on-demand during the synchronous round

// by the IdManager.

#[derive(Debug)]

struct SpecVarStream {

    connector_id: ConnectorId,

    port_suffix_stream: U32Stream,

}

// Manages the messy state of the various endpoints, pollers, buffers, etc.

#[derive(Debug)]

struct EndpointManager {

    // invariants:

    // 1. net and udp endpoints are registered with poll with tokens computed with TargetToken::into

    // 2. Events is empty

    poll: Poll,

    events: Events,

    delayed_messages: Vec<(usize, Msg)>,

    undelayed_messages: Vec<(usize, Msg)>, // ready to yield

    net_endpoint_store: EndpointStore<NetEndpointExt>,

    udp_endpoint_store: EndpointStore<UdpEndpointExt>,

    udp_in_buffer: UdpInBuffer,

}

// A storage of endpoints, which keeps track of which components have raised

// an event during poll(), signifying that they need to be checked for new incoming data

#[derive(Debug)]

struct EndpointStore<T> {

    endpoint_exts: Vec<T>,

    polled_undrained: VecSet<usize>,

}

// The information associated with a port identifier, designed for local storage.

#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]

struct PortInfo {

    owner: ComponentId,

    peer: Option<PortId>,

    polarity: Polarity,

    route: Route,

}

// Similar to `PortInfo`, but designed for communication during the setup procedure.

#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]

struct MyPortInfo {

    polarity: Polarity,

    port: PortId,

    owner: ComponentId,

}

// Newtype around port info map, allowing the implementation of some

// useful methods

#[derive(Default, Debug, Clone, serde::Serialize, serde::Deserialize)]

struct PortInfoMap {

    map: HashMap<PortId, PortInfo>,

}

// A convenient substructure for containing port info and the ID manager.

// Houses the bulk of the connector's persistent state between rounds.

// It turns out several situations require access to both things.

#[derive(Debug, Clone)]

struct IdAndPortState {

    port_info: PortInfoMap,

    id_manager: IdManager,

}

// A component's setup-phase-specific data

#[derive(Debug)]

struct ConnectorCommunication {

    round_index: usize,

    endpoint_manager: EndpointManager,

    neighborhood: Neighborhood,

    native_batches: Vec<NativeBatch>,

    round_result: Result<Option<RoundEndedNative>, SyncError>,

}

// A component's data common to both setup and communication phases

#[derive(Debug)]

struct ConnectorUnphased {

    proto_description: Arc<ProtocolDescription>,

    proto_components: HashMap<ComponentId, ComponentState>,

    logger: Box<dyn Logger>,

    ips: IdAndPortState,

    native_component_id: ComponentId,

}

// A connector's phase-specific data

#[derive(Debug)]

enum ConnectorPhased {

    Setup(Box<ConnectorSetup>),

    Communication(Box<ConnectorCommunication>),

}

// A connector's setup-phase-specific data

#[derive(Debug)]

struct ConnectorSetup {

    net_endpoint_setups: Vec<NetEndpointSetup>,

    udp_endpoint_setups: Vec<UdpEndpointSetup>,

}

// A newtype wrapper for a map from speculative variable to speculative value

// A missing mapping corresponds with "unspecified".

#[derive(Default, Clone, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize)]

struct Predicate {

    assigned: BTreeMap<SpecVar, SpecVal>,

}

// Identifies a child of this connector in the _solution tree_.

// Each connector creates its own local solutions for the consensus procedure during `sync`,

// from the solutions of its children. Those children are either locally-managed components,

// (which are leaves in the solution tree), or other connectors reachable through the given

// network endpoint (which are internal nodes in the solution tree).

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize)]

enum SubtreeId {

    LocalComponent(ComponentId),

    NetEndpoint { index: usize },

}

// An accumulation of the connector's knowledge of all (a) the local solutions its children

// in the solution tree have found, and (b) its own solutions derivable from those of its children.

// This structure starts off each round with an empty set, and accumulates solutions as they are found

// by local components, or received over the network in control messages.

// IMPORTANT: solutions, once found, don't go away until the end of the round. That is to

// say that these sets GROW until the round is over, and all solutions are reset.

#[derive(Debug)]

struct SolutionStorage {

    // invariant: old_local U new_local solutions are those that can be created from

    // the UNION of one element from each set in `subtree_solution`.

    // invariant is maintained by potentially populating new_local whenever subtree_solutions is populated.

    old_local: HashSet<Predicate>, // already sent to this connector's parent OR decided

    new_local: HashSet<Predicate>, // not yet sent to this connector's parent OR decided

    // this pair acts as SubtreeId -> HashSet<Predicate> which is friendlier to iteration

    subtree_solutions: Vec<HashSet<Predicate>>,

    subtree_id_to_index: HashMap<SubtreeId, usize>,

}

// Stores the transient data of a synchronous round.

// Some of it is for bookkeeping, and the rest is a temporary mirror of fields of

// `ConnectorUnphased`, such that any changes are safely contained within RoundCtx,

// and can be undone if the round fails.

struct RoundCtx {

    solution_storage: SolutionStorage,

    spec_var_stream: SpecVarStream,

    payload_inbox: Vec<(PortId, SendPayloadMsg)>,

    deadline: Option<Instant>,

    ips: IdAndPortState,

}

// A trait intended to limit the access of the ConnectorUnphased structure

// such that we don't accidentally modify any important component/port data

// while the results of the round are undecided. Why? Any actions during Connector::sync

// are _speculative_ until the round is decided, and we need a safe way of rolling

// back any changes.

trait CuUndecided {

    fn logger(&mut self) -> &mut dyn Logger;

    fn proto_description(&self) -> &ProtocolDescription;

    fn native_component_id(&self) -> ComponentId;

    fn logger_and_protocol_description(&mut self) -> (&mut dyn Logger, &ProtocolDescription);

}

// Represents a set of synchronous port operations that the native component