use std::io::{Read, Write};

use msg::{MaelstromBody, MaelstromBodyOr, Message, MessageHeader};
use serde::{Deserialize, Serialize};
use serde_json::Deserializer;

pub mod msg;

pub trait Handler {
    type Body: Serialize + for<'a> Deserialize<'a>;

    fn init(node_id: String, node_ids: Vec<String>, msg_id: usize) -> Self;
    fn handle(
        &mut self,
        header: MessageHeader,
        body: Self::Body,
        writer: &mut MsgWriter<impl Write>,
    ) -> ();
}

pub struct MsgWriter<W> {
    node_id: String,
    writer: W,
}

impl<W: Write> MsgWriter<W> {
    pub fn new(node_id: String, writer: W) -> Self {
        Self { node_id, writer }
    }

    pub fn write<T: Serialize>(&mut self, dst: String, msg: &T) {
        let msg = Message {
            header: MessageHeader {
                src: self.node_id.clone(),
                dst,
            },
            body: MaelstromBodyOr::Other { inner: msg },
        };
        serde_json::to_writer(&mut self.writer, &msg).unwrap();
        self.writer.write(&[b'\n']).unwrap();
    }
}

pub fn run_with<T: Handler>(mut reader: impl Read, mut writer: impl Write) {
    let (mut handler, mut msg_writer) = init_handler::<T, _, _>(&mut reader, &mut writer);

    let deser = Deserializer::from_reader(reader);
    for msg in deser.into_iter::<Message<T::Body>>() {
        let msg = msg.unwrap();
        match msg.body {
            MaelstromBodyOr::Other { inner } => {
                handler.handle(msg.header, inner, &mut msg_writer);
            }
            _ => todo!(),
        };
    }
}

pub fn init_handler<T: Handler, R: Read, W: Write>(reader: R, writer: W) -> (T, MsgWriter<W>) {
    let deser = Deserializer::from_reader(reader);
    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");
        }
    };

    let mut writer = MsgWriter::new(node_id.clone(), writer);

    writer.write(
        msg.header.src,
        &MaelstromBody::InitOk {
            msg_id: 0,
            in_reply_to: msg_id,
        },
    );

    (T::init(node_id, node_ids, msg_id), writer)
}