use std::collections::BTreeMap;

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

use crate::tui::TAB_WIDTH;

/// Terminal-cell coordinates in the original source line, including its tab stops.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct Segment {
    pub start: usize,
    pub end: usize,
}

#[derive(Default)]
pub(crate) struct WrapCache {
    key: Option<(usize, u64)>,
    width: usize,
    rows: BTreeMap<usize, Vec<Segment>>,
}

impl WrapCache {
    pub(crate) fn prepare(&mut self, key: (usize, u64), width: usize) {
        let width = width.max(1);
        if self.key != Some(key) || self.width != width {
            self.rows.clear();
            self.key = Some(key);
            self.width = width;
        }
    }

    pub(crate) fn segments(&mut self, row: usize, text: &str) -> &[Segment] {
        let width = self.width.max(1);
        self.rows.entry(row).or_insert_with(|| {
            let mut segments = Vec::new();
            let mut start = 0;
            let mut column = 0;
            let mut word_boundary = 0;
            for grapheme in text.graphemes(true) {
                let cells = if grapheme == "\t" {
                    TAB_WIDTH - column % TAB_WIDTH
                } else {
                    UnicodeWidthStr::width(grapheme)
                };
                if column + cells - start > width && column > start {
                    let end = if word_boundary > start {
                        word_boundary
                    } else {
                        column
                    };
                    segments.push(Segment { start, end });
                    start = end;
                    if column + cells - start > width && column > start {
                        segments.push(Segment { start, end: column });
                        start = column;
                    }
                }
                if grapheme == "\t" {
                    // Tabs are expanded cells, so even a one-column pane preserves them.
                    while column + cells - start > width {
                        let end = start + width;
                        segments.push(Segment { start, end });
                        start = end;
                    }
                }
                column += cells;
                if grapheme.chars().all(char::is_whitespace) {
                    word_boundary = column;
                }
            }
            segments.push(Segment { start, end: column });
            // A full final segment has no spare cell for the insertion caret.
            if column - start >= width {
                segments.push(Segment {
                    start: column,
                    end: column,
                });
            }
            segments
        })
    }
}

pub(crate) fn caret_segment(segments: &[Segment], column: usize) -> usize {
    segments
        .partition_point(|segment| segment.start <= column)
        .saturating_sub(1)
}

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

    fn layout(text: &str, width: usize) -> Vec<Segment> {
        let mut cache = WrapCache::default();
        cache.prepare((1, 0), width);
        cache.segments(0, text).to_vec()
    }

    #[test]
    fn words_and_long_words_preserve_every_source_cell() {
        assert_eq!(
            layout("one two three", 8),
            vec![Segment { start: 0, end: 8 }, Segment { start: 8, end: 13 },]
        );
        let segments = layout("abcdefgh", 3);
        assert_eq!(
            segments,
            vec![
                Segment { start: 0, end: 3 },
                Segment { start: 3, end: 6 },
                Segment { start: 6, end: 8 },
            ]
        );
        assert_eq!(caret_segment(&segments, 3), 1);
        assert_eq!(caret_segment(&segments, 8), 2);
    }

    #[test]
    fn tabs_use_source_stops_and_unicode_graphemes_remain_intact() {
        assert_eq!(
            layout("ab\t界e\u{301} z", 4),
            vec![
                Segment { start: 0, end: 4 },
                Segment { start: 4, end: 8 },
                Segment { start: 8, end: 9 },
            ]
        );
        assert_eq!(
            layout("界界界", 3),
            vec![
                Segment { start: 0, end: 2 },
                Segment { start: 2, end: 4 },
                Segment { start: 4, end: 6 },
            ]
        );
        assert_eq!(
            layout("   \t ", 3),
            vec![Segment { start: 0, end: 3 }, Segment { start: 3, end: 5 },]
        );
    }

    #[test]
    fn narrow_tabs_preserve_all_expanded_cells_at_original_stops() {
        assert_eq!(
            layout("\tX", 2),
            [
                Segment { start: 0, end: 2 },
                Segment { start: 2, end: 4 },
                Segment { start: 4, end: 5 },
            ]
        );
        assert_eq!(
            layout("a\tX", 2),
            [
                Segment { start: 0, end: 1 },
                Segment { start: 1, end: 3 },
                Segment { start: 3, end: 5 },
                Segment { start: 5, end: 5 },
            ]
        );
        for source in ["\t", "a\t", "ab\tX", "\t\tX"] {
            let segments = layout(source, 1);
            assert!(
                segments
                    .iter()
                    .all(|segment| segment.end - segment.start <= 1)
            );
            assert!(segments.windows(2).all(|pair| pair[0].end == pair[1].start));
            let end = source.graphemes(true).fold(0, |column, grapheme| {
                column
                    + if grapheme == "\t" {
                        TAB_WIDTH - column % TAB_WIDTH
                    } else {
                        1
                    }
            });
            assert_eq!(segments.last().unwrap(), &Segment { start: end, end });
        }
    }

    #[test]
    fn full_lines_have_an_insertion_cell_and_cache_invalidation_is_lazy() {
        let mut cache = WrapCache::default();
        cache.prepare((7, 2), 4);
        assert_eq!(
            cache.segments(90, "abcd"),
            [Segment { start: 0, end: 4 }, Segment { start: 4, end: 4 },]
        );
        assert_eq!(caret_segment(cache.segments(90, "abcd"), 4), 1);
        assert_eq!(cache.rows.len(), 1);
        cache.prepare((7, 2), 4);
        assert_eq!(cache.rows.len(), 1);
        cache.prepare((7, 2), 8);
        assert!(cache.rows.is_empty());
        cache.segments(90, "abcd");
        cache.prepare((7, 3), 8);
        assert!(cache.rows.is_empty());
        assert_eq!(layout("", 0), [Segment { start: 0, end: 0 }]);
    }
}