use distributed_systems_flyio::{
    types::{
        BroadcastMessage, EchoMessage, GenerateMessage, InitMessage, Message, ReadMessage,
        TopologyMessage,
    },
    Malestorm,
};
use std::rc::Rc;

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

    program.register(
        "init",
        // TODO: Replace error string with a type
        |mut msg: Message, program: Rc<Malestorm>| -> 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))?;

            program.set_node_id(message_body.node_id);
            program.set_nodes(message_body.nodes);

            msg.body.message_type = "init_ok".into();

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

    program.register(
        "echo",
        |mut msg: Message, program: Rc<Malestorm>| -> Result<(), String> {
            let message_body: EchoMessage = serde_json::from_value(msg.body.message_body.take())
                .map_err(|e| format!("error in parsing response body: {}", e))?;

            msg.body.message_type = "echo_ok".into();
            msg.body.message_body = serde_json::to_value(message_body).unwrap();

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

    program.register(
        "generate",
        |mut msg: Message, program: Rc<Malestorm>| -> Result<(), String> {
            msg.body.message_type = "generate_ok".into();
            msg.body.message_body = serde_json::to_value(GenerateMessage {
                id: program.generate_client_id(),
            })
            .unwrap();

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

    program.register(
        "broadcast",
        |mut msg: Message, program: Rc<Malestorm>| -> Result<(), String> {
            let body: BroadcastMessage = serde_json::from_value(msg.body.message_body.take())
                .map_err(|e| format!("error in parsing response body: {}", e))?;
            msg.body.message_type = "broadcast_ok".into();

            program.record_broadcast(body);

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

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

            let body = serde_json::to_value(ReadMessage {
                messages: program.read_broadcasts(),
            })
            .map_err(|e| format!("error in marshalling response: {}", e))?;
            msg.body.message_body = body;

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

    program.register(
        "topology",
        |mut msg: Message, program: Rc<Malestorm>| -> Result<(), String> {
            let topology: TopologyMessage = serde_json::from_value(msg.body.message_body.take())
                .map_err(|e| format!("error in  parsing response: {}", e))?;

            msg.body.message_type = "topology_ok".into();
            program.save_toplogy(topology);

            program
                .reply(msg)
                .map_err(|e| format!("broadcast: error in writing response: {}", e))
        },
    );
    program.run();
}