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

pub type InlineStack = Arc<Vec<Arc<String>>>;

#[derive(Debug, Clone)]
pub struct InlineStackMap {
    entries: Vec<(u32, InlineStack)>,
    empty: InlineStack,
}

impl InlineStackMap {
    pub fn empty() -> Self {
        Self {
            entries: vec![],
            empty: Arc::new(vec![]),
        }
    }
    pub fn builder() -> InlineStackMapBuilder {
        InlineStackMapBuilder {
            events: BTreeMap::new(),
        }
    }
    pub fn get(&self, address: u32) -> &InlineStack {
        match self.entries.binary_search_by_key(&address, |(a, _)| *a) {
            Ok(index) => self.entries.get(index),
            Err(after) => after.checked_sub(1).and_then(|i| self.entries.get(i)),
        }
        .map(|(_, s)| s)
        .unwrap_or(&self.empty)
    }
}

#[derive(Default)]
struct Event {
    end: usize,
    start: Vec<Arc<String>>,
}

pub struct InlineStackMapBuilder {
    events: BTreeMap<u32, Event>,
}

impl InlineStackMapBuilder {
    pub fn add(&mut self, start: u32, end: u32, name: Arc<String>) {
        self.events.entry(start).or_default().start.push(name);
        self.events.entry(end).or_default().end += 1;
    }

    pub fn build(self) -> InlineStackMap {
        let empty = Arc::new(vec![]);
        let mut entries = vec![];
        let mut stack_indexes = vec![];
        let mut stack = vec![];
        let mut string_dedup = HashMap::new();
        let mut stack_dedup = BTreeMap::new();
        stack_dedup.insert(vec![], empty.clone());

        for (address, event) in self.events {
            for _ in 0..event.end {
                stack.pop();
                stack_indexes.pop();
            }
            for call in event.start {
                if let Some(index) = string_dedup.get(&call) {
                    stack.push(call);
                    stack_indexes.push(*index);
                } else {
                    let index = string_dedup.len();
                    string_dedup.insert(call.clone(), index);
                    stack.push(call);
                    stack_indexes.push(index);
                }
            }
            if let Some(stack) = stack_dedup.get(&stack_indexes) {
                entries.push((address, stack.clone()));
            } else {
                let stack = Arc::new(stack.clone());
                stack_dedup.insert(stack_indexes.clone(), stack.clone());
                entries.push((address, stack));
            }
        }
        InlineStackMap { entries, empty }
    }
}