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

const MAX_NODE_SIZE: usize = 1024;

enum Insert {
    Node(Rope),
    Split(Rope, Rope),
}

#[derive(Debug, Clone)]
struct Leaf {
    unsafe_buf: Option<Rc<[u8]>>,
    start: u16,
    end: u16,
    lines: u16,
}

impl Leaf {
    fn bytes(&self) -> &[u8] {
        match &self.unsafe_buf {
            None => &[],
            Some(buf) => &buf[usize::from(self.start)..usize::from(self.end)],
        }
    }
}

#[derive(Debug, Clone)]
struct Branch {
    children: Rc<[Rope]>,
    len: usize,
    lines: usize,
}

#[derive(Debug, Clone)]
enum Node {
    Leaf(Leaf),
    Branch(Branch),
}

#[derive(Debug, Clone)]
pub struct Rope(Node);

impl Rope {
    pub fn len(&self) -> usize {
        match self {
            Rope(Node::Leaf(l)) => l.bytes().len(),
            Rope(Node::Branch(b)) => b.len,
        }
    }

    pub fn lines(&self) -> usize {
        match self {
            Rope(Node::Leaf(l)) => usize::from(l.lines),
            Rope(Node::Branch(b)) => b.lines,
        }
    }

    fn leaf(buf: &[u8]) -> Rope {
        if buf.len() > MAX_NODE_SIZE {
            panic!("buffer too long for leaf!");
        }
        if buf.is_empty() {
            return Rope(Node::Leaf(Leaf {
                unsafe_buf: None,
                start: 0,
                end: 0,
                lines: 0,
            }));
        }
        let mut lines = 0;
        for &b in buf.iter() {
            if b == b'\n' {
                lines += 1;
            }
        }
        Rope(Node::Leaf(Leaf {
            unsafe_buf: Some(Rc::from(buf)),
            start: 0,
            end: u16::try_from(buf.len()).expect("buf.len() <= MAX_NODE_SIZE"),
            lines,
        }))
    }

    fn branch(children: &[Rope]) -> Rope {
        let mut len = 0;
        let mut lines = 0;
        for c in children.iter() {
            len += c.len();
            lines += c.lines();
        }
        Rope(Node::Branch(Branch {
            children: Rc::from(children),
            len,
            lines,
        }))
    }

    fn concat_height(&self, self_height: usize, other: &Rope, other_height: usize) -> Insert {
        if let (Rope(Node::Leaf(l)), Rope(Node::Leaf(r))) = (self, other) {
            if self.len() + other.len() <= MAX_NODE_SIZE {
                let mut buf = vec![0; self.len() + other.len()];
                buf[..self.len()].copy_from_slice(l.bytes());
                buf[self.len()..].copy_from_slice(r.bytes());
                return Insert::Node(Rope::leaf(&buf));
            }
            if self.len() < MAX_NODE_SIZE / 2 || other.len() < MAX_NODE_SIZE / 2 {
                let mut buf = vec![0; self.len() + other.len()];
                buf[..self.len()].copy_from_slice(l.bytes());
                buf[self.len()..].copy_from_slice(r.bytes());
                return Insert::Split(
                    Rope::leaf(&buf[..buf.len() / 2]),
                    Rope::leaf(&buf[buf.len() / 2..]),
                );
            }
            return Insert::Split(self.clone(), other.clone());
        }
        if self_height == other_height {
            return Insert::Split(self.clone(), other.clone());
        }
        if self_height > other_height {
            let Rope(Node::Branch(b)) = self else {
                panic!("self_height is at least 1, so it's a branch");
            };
            match b.children[b.children.len() - 1].concat_height(
                self_height - 1,
                other,
                other_height,
            ) {
                Insert::Node(new_child) => {
                    let mut new_children = b.children.to_vec();
                    new_children[b.children.len() - 1] = new_child;
                    return Insert::Node(Rope::branch(&new_children));
                }
                Insert::Split(child1, child2) => {
                    if b.children.len() == 2 {
                        return Insert::Node(Rope::branch(&[
                            b.children[0].clone(),
                            child1,
                            child2,
                        ]));
                    }
                    return Insert::Split(
                        Rope::branch(&[b.children[0].clone(), b.children[1].clone()]),
                        Rope::branch(&[child1, child2]),
                    );
                }
            }
        }
        let Rope(Node::Branch(b)) = other else {
            panic!("other_height is at least 1, so it's a branch");
        };
        match self.concat_height(self_height, &b.children[0], other_height - 1) {
            Insert::Node(new_child) => {
                let mut new_children = b.children.to_vec();
                new_children[0] = new_child;
                Insert::Node(Rope::branch(&new_children))
            }
            Insert::Split(child1, child2) => {
                if b.children.len() == 2 {
                    return Insert::Node(Rope::branch(&[child1, child2, b.children[1].clone()]));
                }
                Insert::Split(
                    Rope::branch(&[child1, child2]),
                    Rope::branch(&[b.children[1].clone(), b.children[2].clone()]),
                )
            }
        }
    }

    fn height(&self) -> usize {
        match self {
            Rope(Node::Leaf(_)) => 0,
            Rope(Node::Branch(b)) => 1 + b.children[0].height(),
        }
    }

    pub fn concat(&self, other: &Rope) -> Rope {
        if self.len() == 0 {
            return other.clone();
        }
        if other.len() == 0 {
            return self.clone();
        }
        match self.concat_height(self.height(), other, other.height()) {
            Insert::Node(r) => r,
            Insert::Split(l, r) => Rope::branch(&[l, r]),
        }
    }

    pub fn read(r: &mut dyn Read) -> Result<Rope, std::io::Error> {
        let mut buf = vec![0; MAX_NODE_SIZE];
        let mut rope = Rope::leaf(&[]);
        loop {
            let n = r.read(&mut buf)?;
            if n == 0 {
                break;
            }
            rope = rope.concat(&Rope::leaf(&buf[..n]));
        }
        Ok(rope)
    }

    pub fn new(buf: &[u8]) -> Rope {
        let mut buf = buf;
        Rope::read(&mut buf).expect("reading from a slice should never fail")
    }

    fn is_char_boundary(&self, i: usize) -> bool {
        match self {
            Rope(Node::Leaf(l)) => {
                return i < l.bytes().len() && l.bytes()[i] & 0xc0 != 0x80;
            }
            Rope(Node::Branch(b)) => {
                let mut i = i;
                for c in b.children.iter() {
                    if i < c.len() {
                        return c.is_char_boundary(i);
                    }
                    i -= c.len();
                }
                false
            }
        }
    }

    pub fn floor_char_boundary(&self, i: usize) -> usize {
        for i in (0..i + 1).rev() {
            if self.is_char_boundary(i) {
                return i;
            }
        }
        0
    }

    pub fn ceil_char_boundary(&self, i: usize) -> usize {
        for i in i..self.len() {
            if self.is_char_boundary(i) {
                return i;
            }
        }
        self.len()
    }

    fn line_end(&self, n: usize) -> usize {
        match self {
            Rope(Node::Leaf(l)) => {
                let mut count = 0;
                for (i, &c) in l.bytes().iter().enumerate() {
                    if c != b'\n' {
                        continue;
                    }
                    if count == n {
                        return i;
                    }
                    count += 1;
                }
            }
            Rope(Node::Branch(b)) => {
                let mut n = n;
                let mut offset = 0;
                for c in b.children.iter() {
                    if n < c.lines() {
                        return offset + c.line_end(n);
                    }
                    n -= c.lines();
                    offset += c.len();
                }
            }
        }
        self.len()
    }

    pub fn line_start(&self, n: usize) -> usize {
        if n == 0 {
            return 0;
        }
        self.line_end(n - 1) + 1
    }

    pub fn line(&self, n: usize) -> Rope {
        if n > self.lines() {
            panic!("Line index {} out of range 0..={}", n, self.lines());
        }
        self.slice(self.line_start(n), self.line_end(n)).clone()
    }

    pub fn print(&self, w: &mut dyn Write) -> Result<(), std::io::Error> {
        match self {
            Rope(Node::Leaf(l)) => {
                w.write_all(l.bytes())?;
                Ok(())
            }
            Rope(Node::Branch(b)) => {
                for c in b.children.iter() {
                    c.print(w)?;
                }
                Ok(())
            }
        }
    }

    pub fn slice(&self, start: usize, end: usize) -> Rope {
        if start > self.len() {
            panic!(
                "Slice start index {} out of range 0..={}",
                start,
                self.len()
            );
        }
        if end > self.len() {
            panic!("Slice end index {} out of range 0..={}", end, self.len());
        }
        if start > end {
            panic!("Slice start index {} greater than end index {}", start, end);
        }
        if start == 0 && end == self.len() {
            return self.clone();
        }
        match self {
            Rope(Node::Leaf(l)) => {
                let mut lines = 0;
                for &b in l.bytes()[start..end].iter() {
                    if b == b'\n' {
                        lines += 1;
                    }
                }
                Rope(Node::Leaf(Leaf {
                    unsafe_buf: l.unsafe_buf.clone(),
                    start: l.start + u16::try_from(start).expect("start <= self.len()"),
                    end: l.start + u16::try_from(end).expect("end <= self.len()"),
                    lines,
                }))
            }
            Rope(Node::Branch(b)) => {
                let mut s = Rope::new(&[]);
                let mut start = start;
                let mut end = end;
                for c in b.children.iter() {
                    if start >= c.len() {
                        start -= c.len();
                        end -= c.len();
                        continue;
                    }
                    s = s.concat(&c.slice(start, std::cmp::min(end, c.len())));
                    if end <= c.len() {
                        break;
                    }
                    start = 0;
                    end -= c.len();
                }
                s
            }
        }
    }

    pub fn byte(&self, i: usize) -> u8 {
        if i >= self.len() {
            panic!("Index {} out of range 0..{}", i, self.len());
        }
        match self {
            Rope(Node::Leaf(l)) => l.bytes()[i],
            Rope(Node::Branch(b)) => {
                let mut i = i;
                for c in b.children.iter() {
                    if i < c.len() {
                        return c.byte(i);
                    }
                    i -= c.len();
                }
                panic!("unreachable");
            }
        }
    }
}