#![feature(async_closure)]

use distributed_systems_flyio::{
    types::{GrowCounterReadMessage, GrowCounterUpdateMessage, InitMessage, Message},
    Malestorm, MalestormIo,
};
use std::sync::{Arc, Mutex, RwLock};
use tokio::runtime::Runtime;

fn main() {
    let mut program = Malestorm::default();

    program.register(
        "init",
        // TODO: Replace error string with a type
        |mut msg: Message,
         program: Arc<RwLock<Malestorm>>,
         io: Arc<Mutex<MalestormIo>>|
         -> Result<(), String> {
            let message_body: InitMessage = serde_json::from_value(msg.body.message_body.take())
                .map_err(|e| format!("error in parsing response body: {}", e))?;

            let mut program = program.write().unwrap();
            program.set_node_id(message_body.node_id);
            program.set_nodes(message_body.nodes);
            msg.body.message_type = "init_ok".into();

            program
                .send(io, msg)
                .map_err(|e| format!("init: error in writing response: {}", e))
        },
    );

    program.register(
        "add",
        |mut msg: Message,
         program: Arc<RwLock<Malestorm>>,
         io: Arc<Mutex<MalestormIo>>|
         -> Result<(), String> {
            let body: GrowCounterUpdateMessage =
                serde_json::from_value(msg.body.message_body.take())
                    .map_err(|e| format!("error in parsing response body: {}", e))?;
            msg.body.message_type = "add_ok".into();

            let mut program = program.write().unwrap();
            let val = program.read_counter(&msg.src);
            program.write_counter(&msg.src, val + body.delta);

            program
                .send(io, msg.clone())
                .map_err(|e| format!("add: error in writing response: {}", e))
        },
    );

    program.register(
        "read",
        |mut msg: Message,
         program: Arc<RwLock<Malestorm>>,
         io: Arc<Mutex<MalestormIo>>|
         -> Result<(), String> {
            msg.body.message_type = "read_ok".into();

            let program = program.write().unwrap();
            let mut sum = 0;
            for node in program.get_nodes() {
                if node == msg.src {
                    sum += program.read_counter(&node);
                    continue;
                }
                // Sync first then add

                program.sync(&node, msg.clone(), io.clone());

                sum += program.read_counter(&msg.src);
            }

            msg.body.message_body =
                serde_json::to_value(GrowCounterReadMessage { value: sum }).unwrap();

            program
                .send(io, msg.clone())
                .map_err(|e| format!("read: error in writing response: {}", e))
        },
    );

    program.register(
        "counter_sync",
        |mut msg: Message,
         program: Arc<RwLock<Malestorm>>,
         io: Arc<Mutex<MalestormIo>>|
         -> Result<(), String> {
            msg.body.message_type = "counter_sync_ok".into();

            let program = program.write().unwrap();
            msg.body.message_body = serde_json::to_value(GrowCounterReadMessage {
                value: program.read_counter(&msg.src),
            })
            .unwrap();

            program
                .send(io, msg.clone())
                .map_err(|e| format!("read: error in writing response: {}", e))
        },
    );

    let io = MalestormIo::default();
    let runtime = Runtime::new().unwrap();

    program.run(runtime, io);
}