use std::collections::{BTreeMap, HashMap};

use nebkor_maelstrom::{mk_payload, Body, Message, Node, Payload, Runner};

type Logs = HashMap<String, BTreeMap<u64, u64>>;

fn main() {
    let node = Roach::new(20);
    let runner = Runner::new(node);
    let on_init = Box::new(on_init);
    runner.run(Some(on_init));
}

fn on_init(runner: &Runner) {
    let msg = Message::from_body(Body::from_type("init_log"));
    runner.get_backdoor().send(msg).unwrap();
}

#[derive(Debug, Default)]
struct OffsetState {
    last_page: usize,
    current_chapter: usize,
}

#[derive(Debug, Default)]
struct Roach {
    // these first three get written to as messages come in
    logs: Logs,
    state: HashMap<String, OffsetState>,
    committed_offsets: HashMap<String, u64>,

    // these are initialized on init and then just read
    offset_base: usize,
    sub_chapters: usize, // number of nodes
    pages: usize,
}

impl Roach {
    pub fn new(pages: usize) -> Self {
        Roach {
            pages,
            ..Default::default()
        }
    }

    fn next_offset(&mut self, key: &str) -> u64 {
        let state = self.state.entry(key.to_string()).or_default();
        if state.last_page + 1 < self.pages {
            state.last_page += 1;
        } else {
            state.last_page = 0;
            state.current_chapter += 1;
        }
        let cur_chap = state.current_chapter;
        let offset_base = self.offset_base;
        let start = (offset_base * self.pages) + (cur_chap * (self.pages * self.sub_chapters));
        (start + state.last_page) as u64
    }

    fn init_log(&mut self, runner: &Runner) {
        self.sub_chapters = runner.nodes().len();
        let id = runner.node_id();
        let mut nodes = runner.nodes().to_owned();
        nodes.sort();
        let idx = nodes
            .iter()
            .enumerate()
            .find(|e| e.1.as_str() == id)
            .unwrap()
            .0;
        self.offset_base = idx;
    }
}

impl Node for Roach {
    fn handle(&mut self, runner: &Runner, msg: Message) {
        let typ = msg.typ();

        match typ {
            "init_log" => {
                self.init_log(runner);
            }
            "send" => {
                let key = msg.body.payload.get("key").unwrap().as_str().unwrap();
                let val = msg.body.payload.get("msg").unwrap().as_u64().unwrap();
                let offset = self.next_offset(key);
                self.logs
                    .entry(key.to_string())
                    .or_default()
                    .entry(offset)
                    .or_insert(val);

                let payload = mk_payload(&[("offset", offset.into())]);
                let body = Body::from_type("send_ok").with_payload(payload);
                runner.reply(&msg, body);
            }
            "poll" => {
                let offsets = msg
                    .body
                    .payload
                    .get("offsets")
                    .unwrap()
                    .as_object()
                    .unwrap();
                let mut payload = Payload::default();
                for req in offsets.iter() {
                    if let Some(offsets) = self.logs.get(req.0.as_str()) {
                        let start = req.1.as_u64().unwrap();
                        let subset = offsets
                            .range(start..)
                            .map(|(&offset, &val)| vec![offset, val])
                            .collect::<Vec<_>>();
                        payload.insert(req.0.to_string(), subset.into());
                    }
                }
                let payload = mk_payload(&[("msgs", payload.into())]);
                let body = Body::from_type("poll_ok").with_payload(payload);
                runner.reply(&msg, body);
            }
            "commit_offsets" => {
                let offsets = msg
                    .body
                    .payload
                    .get("offsets")
                    .unwrap()
                    .as_object()
                    .unwrap();
                for (k, v) in offsets.iter() {
                    self.committed_offsets
                        .entry(k.to_string())
                        .or_insert(v.as_u64().unwrap());
                }
                runner.reply(&msg, Body::from_type("commit_offsets_ok"));
            }
            "list_committed_offsets" => {
                let offsets = self
                    .committed_offsets
                    .iter()
                    .map(|(s, &n)| (s.clone(), <u64 as Into<serde_json::Value>>::into(n)))
                    .collect::<Payload>();
                let mut payload = Payload::new();
                payload.insert("offsets".to_string(), offsets.into());
                let body = Body::from_type("list_committed_offsets_ok").with_payload(payload);
                runner.reply(&msg, body);
            }
            _ => eprintln!("unknown Message type: {msg:?}"),
        }
    }
}