use collections::BTreeMap;
use gpui::HighlightStyle;
use language::{Chunk, LanguageAwareStyling};
use multi_buffer::{MultiBufferChunks, MultiBufferOffset, MultiBufferSnapshot};
use std::{cmp, ops::Range};

use crate::display_map::{HighlightKey, SemanticTokensHighlights, TextHighlights};

pub struct CustomHighlightsChunks<'a> {
    buffer_chunks: MultiBufferChunks<'a>,
    buffer_chunk: Option<Chunk<'a>>,
    offset: MultiBufferOffset,
    multibuffer_snapshot: &'a MultiBufferSnapshot,

    highlight_endpoints: Vec<HighlightEndpoint>,
    active_highlights: BTreeMap<HighlightKey, HighlightStyle>,
    text_highlights: Option<&'a TextHighlights>,
    semantic_token_highlights: Option<&'a SemanticTokensHighlights>,
}

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
struct HighlightEndpoint {
    offset: MultiBufferOffset,
    tag: HighlightKey,
    style: Option<HighlightStyle>,
}

impl<'a> CustomHighlightsChunks<'a> {
    #[ztracing::instrument(skip_all)]
    pub fn new(
        range: Range<MultiBufferOffset>,
        language_aware: LanguageAwareStyling,
        text_highlights: Option<&'a TextHighlights>,
        semantic_token_highlights: Option<&'a SemanticTokensHighlights>,
        multibuffer_snapshot: &'a MultiBufferSnapshot,
    ) -> Self {
        let mut highlight_endpoints = Vec::new();
        create_highlight_endpoints(
            &range,
            text_highlights,
            semantic_token_highlights,
            multibuffer_snapshot,
            &mut highlight_endpoints,
        );
        Self {
            buffer_chunks: multibuffer_snapshot.chunks(range.clone(), language_aware),
            buffer_chunk: None,
            offset: range.start,
            text_highlights,
            highlight_endpoints,
            active_highlights: Default::default(),
            multibuffer_snapshot,
            semantic_token_highlights,
        }
    }

    #[ztracing::instrument(skip_all)]
    pub fn seek(&mut self, new_range: Range<MultiBufferOffset>) {
        create_highlight_endpoints(
            &new_range,
            self.text_highlights,
            self.semantic_token_highlights,
            self.multibuffer_snapshot,
            &mut self.highlight_endpoints,
        );
        self.offset = new_range.start;
        self.buffer_chunks.seek(new_range);
        self.buffer_chunk.take();
        self.active_highlights.clear()
    }
}

fn create_highlight_endpoints(
    range: &Range<MultiBufferOffset>,
    text_highlights: Option<&TextHighlights>,
    semantic_token_highlights: Option<&SemanticTokensHighlights>,
    buffer: &MultiBufferSnapshot,
    highlight_endpoints: &mut Vec<HighlightEndpoint>,
) {
    highlight_endpoints.clear();
    if let Some(text_highlights) = text_highlights {
        let start = buffer.anchor_after(range.start);
        let end = buffer.anchor_after(range.end);
        let mut text_highlights_scratch = Vec::new();

        for (&tag, text_highlights) in text_highlights.iter() {
            let style = text_highlights.0;
            let ranges = &text_highlights.1;

            let start_ix = ranges
                .binary_search_by(|probe| probe.end.cmp(&start, buffer).then(cmp::Ordering::Less))
                .unwrap_or_else(|i| i);
            let end_ix = ranges[start_ix..]
                .binary_search_by(|probe| {
                    probe.start.cmp(&end, buffer).then(cmp::Ordering::Greater)
                })
                .unwrap_or_else(|i| i);

            let ranges_ = &ranges[start_ix..][..end_ix];
            text_highlights_scratch.clear();
            text_highlights_scratch.reserve(ranges_.len());
            highlight_endpoints.reserve(2 * ranges_.len());

            let mut iter = ranges_.iter();
            buffer.summaries_for_anchors_cb(
                ranges_.iter().map(|r| &r.start),
                |start: MultiBufferOffset| {
                    text_highlights_scratch.push((start, iter.next().unwrap().end));
                },
            );
            text_highlights_scratch.sort_by(|a, b| a.1.cmp(&b.1, buffer));
            let mut iter = text_highlights_scratch.iter();
            buffer.summaries_for_anchors_cb(
                text_highlights_scratch.iter().map(|(_, end)| end),
                |end: MultiBufferOffset| {
                    let start = iter.next().unwrap().0;
                    if start == end {
                        return;
                    }
                    highlight_endpoints.push(HighlightEndpoint {
                        offset: start,
                        tag,
                        style: Some(style),
                    });
                    highlight_endpoints.push(HighlightEndpoint {
                        offset: end,
                        tag,
                        style: None,
                    });
                },
            );
        }
    }
    if let Some(semantic_token_highlights) = semantic_token_highlights {
        let start = buffer.anchor_after(range.start);
        let end = buffer.anchor_after(range.end);
        let mut semantic_highlights_scratch = Vec::new();
        for buffer_id in buffer.buffer_ids_for_range(range.clone()) {
            let Some((semantic_token_highlights, interner)) =
                semantic_token_highlights.get(&buffer_id)
            else {
                continue;
            };
            let start_ix = semantic_token_highlights
                .binary_search_by(|probe| {
                    probe
                        .range
                        .end
                        .cmp(&start, buffer)
                        .then(cmp::Ordering::Less)
                })
                .unwrap_or_else(|i| i);
            let end_ix = semantic_token_highlights[start_ix..]
                .binary_search_by(|probe| {
                    probe
                        .range
                        .start
                        .cmp(&end, buffer)
                        .then(cmp::Ordering::Greater)
                })
                .unwrap_or_else(|i| i);

            let ranges_ = &semantic_token_highlights[start_ix..][..end_ix];
            semantic_highlights_scratch.clear();
            semantic_highlights_scratch.reserve(ranges_.len());
            highlight_endpoints.reserve(2 * ranges_.len());

            let mut iter = ranges_.iter();
            buffer.summaries_for_anchors_cb(
                ranges_.iter().map(|token| &token.range.start),
                |start: MultiBufferOffset| {
                    semantic_highlights_scratch.push((start, iter.next().unwrap()));
                },
            );
            semantic_highlights_scratch.sort_by(|a, b| a.1.range.end.cmp(&b.1.range.end, buffer));
            let mut iter = semantic_highlights_scratch.iter();
            buffer.summaries_for_anchors_cb(
                semantic_highlights_scratch
                    .iter()
                    .map(|(_, token)| &token.range.end),
                |end: MultiBufferOffset| {
                    let (start, token) = iter.next().unwrap();
                    if *start == end {
                        return;
                    }
                    highlight_endpoints.push(HighlightEndpoint {
                        offset: *start,
                        tag: HighlightKey::SemanticToken,
                        style: Some(interner[token.style]),
                    });
                    highlight_endpoints.push(HighlightEndpoint {
                        offset: end,
                        tag: HighlightKey::SemanticToken,
                        style: None,
                    });
                },
            );
        }
    }
    highlight_endpoints.sort_by(|a, b| a.cmp(b).reverse());
}

impl<'a> Iterator for CustomHighlightsChunks<'a> {
    type Item = Chunk<'a>;

    #[ztracing::instrument(skip_all)]
    fn next(&mut self) -> Option<Self::Item> {
        let mut next_highlight_endpoint = MultiBufferOffset(usize::MAX);
        while let Some(endpoint) = self.highlight_endpoints.last().copied() {
            if endpoint.offset <= self.offset {
                if let Some(style) = endpoint.style {
                    self.active_highlights.insert(endpoint.tag, style);
                } else {
                    self.active_highlights.remove(&endpoint.tag);
                }
                self.highlight_endpoints.pop();
            } else {
                next_highlight_endpoint = endpoint.offset;
                break;
            }
        }

        let chunk = match &mut self.buffer_chunk {
            Some(it) => it,
            slot => slot.insert(self.buffer_chunks.next()?),
        };
        while chunk.text.is_empty() {
            *chunk = self.buffer_chunks.next()?;
        }

        let split_idx = chunk.text.len().min(next_highlight_endpoint - self.offset);
        let (prefix, suffix) = chunk.text.split_at(split_idx);
        self.offset += prefix.len();

        let mask = 1u128.unbounded_shl(split_idx as u32).wrapping_sub(1);
        let chars = chunk.chars & mask;
        let tabs = chunk.tabs & mask;
        let newlines = chunk.newlines & mask;
        let mut prefix = Chunk {
            text: prefix,
            chars,
            tabs,
            newlines,
            ..chunk.clone()
        };

        chunk.chars = chunk.chars.unbounded_shr(split_idx as u32);
        chunk.tabs = chunk.tabs.unbounded_shr(split_idx as u32);
        chunk.newlines = chunk.newlines.unbounded_shr(split_idx as u32);
        chunk.text = suffix;
        if !self.active_highlights.is_empty() {
            prefix.highlight_style = self
                .active_highlights
                .values()
                .copied()
                .reduce(|acc, active_highlight| acc.highlight(active_highlight));
        }
        Some(prefix)
    }
}

impl PartialOrd for HighlightEndpoint {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for HighlightEndpoint {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.offset
            .cmp(&other.offset)
            .then_with(|| self.style.is_some().cmp(&other.style.is_some()))
            .then_with(|| self.tag.cmp(&other.tag))
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use super::*;
    use crate::MultiBuffer;
    use gpui::App;
    use rand::prelude::*;
    use util::RandomCharIter;

    #[gpui::test(iterations = 100)]
    fn test_random_chunk_bitmaps(cx: &mut App, mut rng: StdRng) {
        // Generate random buffer using existing test infrastructure
        let len = rng.random_range(10..10000);
        let buffer = if rng.random() {
            let text = RandomCharIter::new(&mut rng).take(len).collect::<String>();
            MultiBuffer::build_simple(&text, cx)
        } else {
            MultiBuffer::build_random(&mut rng, cx)
        };

        let buffer_snapshot = buffer.read(cx).snapshot(cx);

        // Create random highlights
        let mut highlights = sum_tree::TreeMap::default();
        let highlight_count = rng.random_range(1..10);

        for _i in 0..highlight_count {
            let style = HighlightStyle {
                color: Some(gpui::Hsla {
                    h: rng.random::<f32>(),
                    s: rng.random::<f32>(),
                    l: rng.random::<f32>(),
                    a: 1.0,
                }),
                ..Default::default()
            };

            let mut ranges = Vec::new();
            let range_count = rng.random_range(1..10);
            let text = buffer_snapshot.text();
            for _ in 0..range_count {
                if buffer_snapshot.len() == MultiBufferOffset(0) {
                    continue;
                }

                let mut start = rng.random_range(
                    MultiBufferOffset(0)..=buffer_snapshot.len().saturating_sub_usize(10),
                );

                while !text.is_char_boundary(start.0) {
                    start = start.saturating_sub_usize(1);
                }

                let end_end = buffer_snapshot.len().min(start + 100usize);
                let mut end = rng.random_range(start..=end_end);
                while !text.is_char_boundary(end.0) {
                    end = end.saturating_sub_usize(1);
                }

                if start < end {
                    start = end;
                }
                let start_anchor = buffer_snapshot.anchor_before(start);
                let end_anchor = buffer_snapshot.anchor_after(end);
                ranges.push(start_anchor..end_anchor);
            }

            highlights.insert(HighlightKey::Editor, Arc::new((style, ranges)));
        }

        // Get all chunks and verify their bitmaps
        let chunks = CustomHighlightsChunks::new(
            MultiBufferOffset(0)..buffer_snapshot.len(),
            LanguageAwareStyling {
                tree_sitter: false,
                diagnostics: false,
            },
            None,
            None,
            &buffer_snapshot,
        );

        for chunk in chunks {
            let chunk_text = chunk.text;
            let chars_bitmap = chunk.chars;
            let tabs_bitmap = chunk.tabs;

            // Check empty chunks have empty bitmaps
            if chunk_text.is_empty() {
                assert_eq!(
                    chars_bitmap, 0,
                    "Empty chunk should have empty chars bitmap"
                );
                assert_eq!(tabs_bitmap, 0, "Empty chunk should have empty tabs bitmap");
                continue;
            }

            // Verify that chunk text doesn't exceed 128 bytes
            assert!(
                chunk_text.len() <= 128,
                "Chunk text length {} exceeds 128 bytes",
                chunk_text.len()
            );

            // Verify chars bitmap
            let char_indices = chunk_text
                .char_indices()
                .map(|(i, _)| i)
                .collect::<Vec<_>>();

            for byte_idx in 0..chunk_text.len() {
                let should_have_bit = char_indices.contains(&byte_idx);
                let has_bit = chars_bitmap & (1 << byte_idx) != 0;

                if has_bit != should_have_bit {
                    eprintln!("Chunk text bytes: {:?}", chunk_text.as_bytes());
                    eprintln!("Char indices: {:?}", char_indices);
                    eprintln!("Chars bitmap: {:#b}", chars_bitmap);
                    assert_eq!(
                        has_bit, should_have_bit,
                        "Chars bitmap mismatch at byte index {} in chunk {:?}. Expected bit: {}, Got bit: {}",
                        byte_idx, chunk_text, should_have_bit, has_bit
                    );
                }
            }

            // Verify tabs bitmap
            for (byte_idx, byte) in chunk_text.bytes().enumerate() {
                let is_tab = byte == b'\t';
                let has_bit = tabs_bitmap & (1 << byte_idx) != 0;

                if has_bit != is_tab {
                    eprintln!("Chunk text bytes: {:?}", chunk_text.as_bytes());
                    eprintln!("Tabs bitmap: {:#b}", tabs_bitmap);
                    assert_eq!(
                        has_bit, is_tab,
                        "Tabs bitmap mismatch at byte index {} in chunk {:?}. Byte: {:?}, Expected bit: {}, Got bit: {}",
                        byte_idx, chunk_text, byte as char, is_tab, has_bit
                    );
                }
            }
        }
    }
}