use std::io;

use common::{msg::*, send_msg, Handler};
use serde::{Deserialize, Serialize};
use serde_json::Deserializer;

#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum EchoBody {
    #[serde(rename = "echo")]
    Echo { msg_id: usize, echo: String },

    #[serde(rename = "echo_ok")]
    EchoOk {
        msg_id: usize,
        in_reply_to: usize,
        echo: String,
    },
}
fn main() {
    let mut stdout = io::stdout();
    let deser = Deserializer::from_reader(io::stdin());
    let mut deser = deser.into_iter::<Message<()>>();
    let Some(msg) = deser.next() else {
        panic!("stream ended before init message");
    };
    let Ok(msg) = msg else {
            panic!("{}", msg.unwrap_err());
    };

    let (node_id, node_ids, msg_id) = match msg.body {
        MaelstromBodyOr::MaelstromBody {
            inner:
                MaelstromBody::Init {
                    node_id,
                    node_ids,
                    msg_id,
                },
        } => (node_id, node_ids, msg_id),
        _ => {
            panic!("expected init message to be first message");
        }
    };

    send_msg(
        &mut stdout,
        &Message {
            header: msg.header.flip(),
            body: MaelstromBodyOr::MaelstromBody::<()> {
                inner: MaelstromBody::InitOk {
                    msg_id: 0,
                    in_reply_to: msg_id,
                },
            },
        },
    );

    let mut handler = EchoHandler::init(node_id, node_ids, msg_id);

    drop(deser);

    let deser = Deserializer::from_reader(io::stdin());
    for msg in deser.into_iter::<Message<EchoBody>>() {
        let msg = msg.unwrap();
        match msg.body {
            MaelstromBodyOr::Other { inner } => {
                if let Some(out) = handler.handle(msg.header, inner) {
                    send_msg(&mut stdout, &out);
                }
            }
            _ => todo!(),
        };
    }
}

pub struct EchoHandler {
    next_msg_id: usize,
}

impl Handler for EchoHandler {
    type Body = EchoBody;

    fn init(_node_id: String, _node_ids: Vec<String>, _msg_id: usize) -> Self {
        EchoHandler { next_msg_id: 1 }
    }

    fn handle(&mut self, header: MessageHeader, body: Self::Body) -> Option<Message<EchoBody>> {
        match body {
            EchoBody::Echo { msg_id, echo } => {
                let msg = Message {
                    header: header.flip(),
                    body: MaelstromBodyOr::Other {
                        inner: EchoBody::EchoOk {
                            msg_id: self.next_msg_id,
                            in_reply_to: msg_id,
                            echo,
                        },
                    },
                };

                self.next_msg_id += 1;

                Some(msg)
            }
            EchoBody::EchoOk { .. } => None,
        }
    }
}