use serde::{Deserialize, Serialize};
use smol::channel::Sender;

use crate::msg_id::MessageID;

#[derive(Debug, Serialize, Deserialize)]
pub struct Message<B> {
    #[serde(flatten)]
    pub header: MessageHeader,
    pub body: B,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct MessageHeader {
    pub src: String,
    #[serde(rename = "dest")]
    pub dst: String,
}

impl MessageHeader {
    pub fn flip(self) -> Self {
        Self {
            src: self.dst,
            dst: self.src,
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum MaelstromBody {
    #[serde(rename = "init")]
    Init {
        node_id: String,
        node_ids: Vec<String>,
        msg_id: MessageID,
    },
    #[serde(rename = "init_ok")]
    InitOk {
        msg_id: MessageID,
        in_reply_to: MessageID,
    },
}

pub struct Output<B> {
    node_id: String,
    channel: Sender<Message<B>>,
}

impl<B: Serialize + Clone> Output<B> {
    pub fn new(node_id: String, channel: Sender<Message<B>>) -> Self {
        Self { node_id, channel }
    }

    pub async fn send(&self, dst: &str, body: &B) {
        self.send_raw(Message {
            header: MessageHeader {
                src: self.node_id.clone(),
                dst: dst.to_string(),
            },
            body: body.clone(),
        })
        .await;
    }
    pub async fn send_raw(&self, msg: Message<B>) {
        self.channel.send(msg).await.unwrap();
    }
}