use std::io::Write;

#[derive(Default)]
pub struct SparseBuffer {
    chunks: Vec<(usize, Vec<u8>)>,
}

impl SparseBuffer {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn write(&mut self, addr: usize, bytes: &[u8]) {
        if self.chunks.is_empty() {
            self.chunks.push((addr, Vec::from(bytes)));
            return;
        }

        enum InsertLoc {
            Head,
            Append(usize),
            Tail,
        }

        let mut loc = None;
        for (i, (at, _ch)) in self.chunks.iter().enumerate() {
            if *at > addr {
                if i == 0 {
                    loc = Some(InsertLoc::Head);
                } else {
                    loc = Some(InsertLoc::Append(i - 1));
                }
                break;
            }
        }

        let loc = if let Some(l) = loc {
            l
        } else {
            InsertLoc::Tail
        };

        match loc {
            InsertLoc::Head => self.write_head(addr, bytes),
            InsertLoc::Append(i) => self.write_after(i, addr, bytes),
            InsertLoc::Tail => self.write_tail(addr, bytes),
        };
    }

    fn write_head(&mut self, addr: usize, bytes: &[u8]) {
        self.chunks.insert(0, (addr, vec![]));
        self.write_after(0, addr, bytes);
    }

    fn write_after(&mut self, index: usize, addr: usize, bytes: &[u8]) {
        if index == self.chunks.len() - 1 {
            self.write_tail(addr, bytes);
            return;
        }

        let end_addr = addr + bytes.len();

        // This means we have at least two chunks.
        // The written area can:
        // - fit within the chunk
        // - extend the chunk, but still end before the second's start address
        // - extend the chunk, overflowing into one or more following chunks

        let (a, slice) = self.chunks.iter().nth(index + 1).unwrap();
        let second_start = *a;
        let _second_len = slice.len();

        let (a, slice) = self.chunks.get_mut(index).unwrap();
        let first_addr = *a;
        let first_len = slice.len();

        if end_addr <= first_addr + first_len {
            (&mut slice[(addr - first_addr) as usize..]).write(bytes);
        } else if end_addr <= second_start {
            slice.truncate((addr - first_addr) as usize);
            slice.extend_from_slice(bytes);
        } else {
            // overflows into one or more chunks
            slice.truncate((addr - first_addr) as usize);
            slice.extend_from_slice(&bytes[..(second_start - addr) as usize]);

            // recurse
            self.write_after(index + 1, second_start, &bytes[(second_start - addr) as usize..]);
        }
    }

    fn write_tail(&mut self, addr: usize, bytes: &[u8]) {
        let (a, slice) = self.chunks.last_mut().unwrap();

        let last_addr = *a;
        let last_len = slice.len();
        let end_addr = addr + bytes.len();

        assert!(addr >= last_addr);

        if end_addr <= last_addr + last_len {
            // Entirely contained within the last chunk
            (&mut slice[((addr - last_addr) as usize)..]).write(bytes);
        } else if addr > last_addr + last_len {
            self.chunks.push((addr, Vec::from(bytes)));
        } else {
            // The write slice starts within the last chunk, but extends past its end.
            slice.truncate((addr - last_addr) as usize);
            slice.extend_from_slice(bytes);
        }
    }

    pub fn coalesce(&mut self) {
        let mut prev = Option::<(usize, Vec<u8>)>::None;
        let mut merged = vec![];
        for chunk in self.chunks.drain(..) {
            if let Some(prevchunk) = &mut prev {
                if chunk.0 == prevchunk.0 + prevchunk.1.len() {
                    prevchunk.1.extend_from_slice(&chunk.1[..]);
                    continue;
                } else {
                    merged.push(prev.take().unwrap());
                }
            }
            prev = Some(chunk);
        }
        if let Some(prevchunk) = prev {
            merged.push(prevchunk);
        }
        self.chunks = merged;
    }
}

#[cfg(test)]
mod tests {
    use super::SparseBuffer;

    #[test]
    fn test_empty() {
        let mut sparse = SparseBuffer::new();

        sparse.write(100, &[0, 1, 2, 3, 4]);
        assert_eq!(vec![(100, vec![0, 1, 2, 3, 4])], sparse.chunks);
    }

    #[test]
    fn test_append_sparse() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        sparse.write(100, &[7, 8, 9, 10]);
        assert_eq!(vec![
            (0, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
            (100, vec![7, 8, 9, 10])
        ], sparse.chunks);
    }

    #[test]
    fn test_append_inside_last() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        sparse.write(7, &[70, 80, 90]);
        assert_eq!(vec![
            (0, vec![0, 1, 2, 3, 4, 5, 6, 70, 80, 90, 10]),
        ], sparse.chunks);
    }

    #[test]
    fn test_append_extend_last() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        sparse.write(5, &[50, 60, 70, 80, 90, 100, 110, 120, 130, 140]);
        assert_eq!(vec![(0, vec![0, 1, 2, 3, 4, 50, 60, 70, 80, 90, 100, 110, 120, 130, 140])], sparse.chunks);
    }

    #[test]
    fn test_prepend_sparse() {
        let mut sparse = SparseBuffer::new();

        sparse.write(100, &[0, 1, 2, 3, 4]);
        sparse.write(10, &[70, 80, 90, 100]);

        assert_eq!(vec![
            (10, vec![70, 80, 90, 100]),
            (100, vec![0, 1, 2, 3, 4])
        ], sparse.chunks);
    }

    #[test]
    fn test_within_first() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
        sparse.write(4, &[40, 50, 60]);

        assert_eq!(vec![
            (0, vec![0, 1, 2, 3, 40, 50, 60, 7, 8, 9]),
        ], sparse.chunks);
    }

    #[test]
    fn test_grows_first() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2, 3, 4, 5]);
        sparse.write(10, &[10, 11, 12]);
        sparse.write(4, &[40, 50, 60]);

        assert_eq!(vec![
            (0, vec![0, 1, 2, 3, 40, 50, 60]),
            (10, vec![10, 11, 12]),
        ], sparse.chunks);
    }

    #[test]
    fn test_grows_first2() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2, 3, 4, 5]);
        sparse.write(10, &[10, 11, 12]);
        sparse.write(4, &[40, 50, 60, 70, 80, 90]);

        assert_eq!(vec![
            (0, vec![0, 1, 2, 3, 40, 50, 60, 70, 80, 90]),
            (10, vec![10, 11, 12]),
        ], sparse.chunks);
    }

    #[test]
    fn test_overflow_first() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2, 3, 4, 5]);
        sparse.write(10, &[10, 11, 12, 13, 14, 15]);
        sparse.write(4, &[40, 50, 60, 70, 80, 90, 100, 110, 120]);

        assert_eq!(vec![
            (0, vec![0, 1, 2, 3, 40, 50, 60, 70, 80, 90]),
            (10, vec![100, 110, 120, 13, 14, 15]),
        ], sparse.chunks);
    }

    #[test]
    fn test_join_tail() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2]);
        sparse.write(3, &[3, 4, 5]);
        sparse.write(6, &[6, 7, 8]);

        assert_eq!(vec![
            (0, vec![0, 1, 2, 3, 4, 5, 6, 7, 8]),
        ], sparse.chunks);
    }

    #[test]
    fn test_overflow_multiple() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2]);
        sparse.write(4, &[4, 5, 6]);
        sparse.write(8, &[8, 9, 10]);

        assert_eq!(vec![
            (0, vec![0, 1, 2]),
            (4, vec![4, 5, 6]),
            (8, vec![8, 9, 10]),
        ], sparse.chunks);

        sparse.write(2, &[20, 30, 40, 50, 60, 70, 80, 90]);

        assert_eq!(vec![
            (0, vec![0, 1, 20, 30]),
            (4, vec![40, 50, 60, 70]),
            (8, vec![80, 90, 10]),
        ], sparse.chunks);
    }

    #[test]
    fn test_overflow_multiple2() {
        let mut sparse = SparseBuffer::new();

        sparse.write(0, &[0, 1, 2]);
        sparse.write(4, &[4, 5, 6]);
        sparse.write(8, &[8, 9, 10]);

        assert_eq!(vec![
            (0, vec![0, 1, 2]),
            (4, vec![4, 5, 6]),
            (8, vec![8, 9, 10]),
        ], sparse.chunks);
        sparse.coalesce();

        // no change, as expected
        assert_eq!(vec![
            (0, vec![0, 1, 2]),
            (4, vec![4, 5, 6]),
            (8, vec![8, 9, 10]),
        ], sparse.chunks);

        sparse.write(2, &[20, 30, 40, 50, 60, 70, 80, 90, 100, 110]);

        assert_eq!(vec![
            (0, vec![0, 1, 20, 30]),
            (4, vec![40, 50, 60, 70]),
            (8, vec![80, 90, 100, 110]),
        ], sparse.chunks);

        // join contiguous
        sparse.coalesce();
        assert_eq!(vec![
            (0, vec![0, 1, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110]),
        ], sparse.chunks);
    }
}