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;
            let mut content_end = 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) {
                    content_end = column;
                } else if content_end > start {
                    // A word break must not leave only whitespace on this visual row.
                    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)]
#[path = "../tests/unit/wrap.rs"]
mod tests;