mod queries;
use languages::Language;
pub use queries::highlight_query::{ColorMap, TextSlice};

use std::{
    cell::{Ref, RefCell},
    collections::{HashMap, HashSet},
    sync::Arc,
};

use parking_lot::Mutex;

use arborium::tree_sitter::{InputEdit, Parser, Tree};
use futures::stream::AbortHandle;
use galaxyui::{color::ColorU, AppContext, Entity, ModelContext, WeakModelHandle};
use queries::{
    highlight_query::HighlightQuery,
    indent_query::{indentation_delta, IndentDelta},
};
use rangemap::{RangeMap, RangeSet};
use string_offset::{ByteOffset, CharOffset};

use galaxy_editor::{
    content::{
        buffer::{Buffer, BufferSnapshot},
        edit::PreciseDelta,
        text::IndentUnit,
        version::BufferVersion,
    },
    decoration::DecorationLayer,
};
use galaxyui::text::point::Point;

const MAX_SYNTAX_TREES: usize = 3;

thread_local! {
    static PARSER: RefCell<Parser> = RefCell::new(Parser::new());
}
pub enum DecorationStateEvent {
    DecorationUpdated { version: BufferVersion },
}

struct LanguageQueries {
    language: Arc<Language>,
    syntax_query: HighlightQuery,
}

/// Single-entry cache for highlight queries.
/// Stores the most recent highlight computation result.
struct HighlightCache {
    key: HighlightCacheKey,
    highlights: RangeMap<CharOffset, ColorU>,
}

struct HighlightCacheKey {
    version: BufferVersion,
    ranges: RangeSet<CharOffset>,
    language_id: Option<arborium::tree_sitter::Language>,
}

impl HighlightCacheKey {
    /// Check if this cache entry matches the given content version, ranges, and language.
    fn matches(
        &self,
        version: BufferVersion,
        ranges: &RangeSet<CharOffset>,
        language_id: &Option<arborium::tree_sitter::Language>,
    ) -> bool {
        if self.version != version {
            return false;
        }
        if &self.language_id != language_id {
            return false;
        }
        // RangeSet derives PartialEq, so we can compare directly
        &self.ranges == ranges
    }
}

/// Manages the decoration styles derived from the underlying text source (e.g. syntax highlighting).
/// The updates are computed asynchronously and we notify the editor model upon completion via
/// DecorationUpdated event.
pub struct SyntaxTreeState {
    syntax_tree: Mutex<HashMap<BufferVersion, Tree>>,
    language_queries: Option<LanguageQueries>,
    buffer_version: BufferVersion,
    color_map: ColorMap,
    buffer_handle: WeakModelHandle<Buffer>,
    parsing_handle: Option<AbortHandle>,
    /// Cache for highlight results to avoid recomputing for the same viewport ranges.
    highlight_cache: RefCell<Option<HighlightCache>>,
}

impl SyntaxTreeState {
    pub fn new(
        buffer_handle: WeakModelHandle<Buffer>,
        buffer_version: BufferVersion,
        color_map: ColorMap,
    ) -> Self {
        Self {
            color_map,
            syntax_tree: Mutex::new(HashMap::new()),
            buffer_version,
            buffer_handle,
            parsing_handle: None,
            language_queries: None,
            highlight_cache: RefCell::new(None),
        }
    }

    pub fn set_language(&mut self, language: Arc<Language>) {
        self.language_queries = Some(LanguageQueries {
            syntax_query: HighlightQuery::new(&language.highlight_query, self.color_map),
            language,
        });
    }

    pub fn has_supported_highlighting(&self) -> bool {
        self.language_queries.is_some()
    }

    pub fn indent_unit(&self) -> Option<IndentUnit> {
        self.language_queries
            .as_ref()
            .map(|queries| queries.language.indent_unit)
    }

    pub fn bracket_pairs(&self) -> Option<&[(char, char)]> {
        self.language_queries
            .as_ref()
            .map(|queries| queries.language.bracket_pairs.as_slice())
    }

    pub fn comment_prefix(&self) -> Option<&str> {
        self.language_queries
            .as_ref()
            .and_then(|queries| queries.language.comment_prefix.as_ref())
            .map(|s| s.as_str())
    }

    /// Given multiple character ranges, return their corresponding highlight colors.
    /// If the tree is not ready or the buffer model has been deallocated, this returns None.
    pub fn highlights_in_ranges(
        &self,
        ranges: RangeSet<CharOffset>,
        render_content_version: Option<BufferVersion>,
        ctx: &AppContext,
    ) -> Option<Ref<'_, RangeMap<CharOffset, ColorU>>> {
        // If no render content version is provided, default the most recent content version.
        let buffer_version = render_content_version.unwrap_or(self.buffer_version);

        let language_id = self
            .language_queries
            .as_ref()
            .map(|q| q.language.grammar.clone());

        // Check cache first
        if let Ok(cache) = Ref::filter_map(self.highlight_cache.borrow(), |c| c.as_ref()) {
            if cache.key.matches(buffer_version, &ranges, &language_id) {
                // Return a borrowed reference to the cached highlights
                return Some(Ref::map(cache, |c| &c.highlights));
            }
        }

        // Cache miss - compute highlights
        let mut syntax_tree_lock = self.syntax_tree.lock();
        let tree = syntax_tree_lock.get(&buffer_version)?;
        let buffer = self.buffer_handle.upgrade(ctx)?;
        let language_queries = self.language_queries.as_ref()?;

        let mut combined_highlights = RangeMap::new();

        // Iterate over all ranges and collect highlights for each
        for range in ranges.iter() {
            let highlights = language_queries.syntax_query.get_highlighted_chunks(
                range.clone(),
                &language_queries.language.highlight_query,
                buffer.as_ref(ctx),
                tree,
            );

            // Merge the highlights into the combined map
            for (highlight_range, color) in highlights.iter() {
                combined_highlights.insert(highlight_range.clone(), *color);
            }
        }

        // Once we have rendered content version X, we could discard syntax trees belonging to versions before X.
        if let Some(render_content_version) = render_content_version {
            // First, drop any versions older than the rendered one in a single pass.
            syntax_tree_lock.retain(|version, _| *version >= render_content_version);
            Self::truncate_tree_state(&mut syntax_tree_lock, self.buffer_version);
        }

        // Store in cache before returning
        *self.highlight_cache.borrow_mut() = Some(HighlightCache {
            key: HighlightCacheKey {
                version: buffer_version,
                ranges,
                language_id,
            },
            highlights: combined_highlights,
        });

        // Return a borrowed reference to the cached highlights
        Ref::filter_map(self.highlight_cache.borrow(), |c| {
            c.as_ref().map(|cache| &cache.highlights)
        })
        .ok()
    }

    /// Given a point in buffer, return the absolute indentation level the point should have.
    pub fn indentation_at_point(&self, point: Point, ctx: &AppContext) -> Option<IndentDelta> {
        let syntax_tree_lock = self.syntax_tree.lock();
        let tree = syntax_tree_lock.get(&self.buffer_version)?;
        let buffer = self.buffer_handle.upgrade(ctx)?;
        let language_queries = self.language_queries.as_ref()?;

        indentation_delta(
            buffer.as_ref(ctx),
            tree,
            point,
            language_queries.language.indents_query.as_ref()?,
        )
    }

    /// Re-parse the tree based on the updated tree and source content.
    async fn parse_text(
        content: BufferSnapshot,
        old_tree: Option<Tree>,
        language: &Language,
    ) -> Tree {
        PARSER.with(|parser| {
            let mut parser = parser.borrow_mut();
            parser
                .set_language(&language.grammar)
                .expect("incompatible grammar");
            let mut bytes = content.bytes();
            let mut callback = |byte_offset: usize, _point: arborium::tree_sitter::Point| {
                // Add 1 since the buffer is 1 indexed.
                bytes.seek(ByteOffset::from(byte_offset + 1));
                bytes.next().unwrap_or_default()
            };
            parser
                .parse_with_options(&mut callback, old_tree.as_ref(), None)
                .expect("Should succeed")
        })
    }

    /// Translate an incoming edit delta into an InputEdit for incrementally updating the syntax
    /// tree. Uses the precomputed byte edit info (which was captured from the correct intermediate
    /// buffer state) and `replaced_points` instead of re-deriving from the final buffer.
    fn delta_to_input_edit(delta: &PreciseDelta) -> InputEdit {
        // Convert 1-indexed ByteOffset values to 0-indexed for tree-sitter.
        let start_byte = delta.replaced_byte_range.start.as_usize().saturating_sub(1);
        let old_end_byte = delta.replaced_byte_range.end.as_usize().saturating_sub(1);

        InputEdit {
            start_byte,
            old_end_byte,
            new_end_byte: start_byte + delta.new_byte_length,
            start_position: point_to_syntax_point(delta.replaced_points.start),
            old_end_position: point_to_syntax_point(delta.replaced_points.end),
            new_end_position: point_to_syntax_point(delta.new_end_point),
        }
    }

    pub fn invalidate_highlight_cache_for_version(&self, version: BufferVersion) {
        // Check if the cache exists and if it matches the version being invalidated
        let mut cache = self.highlight_cache.borrow_mut();
        if let Some(ref cached) = *cache {
            if cached.key.version == version {
                *cache = None;
            }
        }
    }

    pub fn set_color_map(&mut self, color_map: ColorMap) {
        self.color_map = color_map;
        if let Some(language_query) = self.language_queries.take() {
            self.set_language(language_query.language);
        }
        // Clear highlight cache since colors have changed
        *self.highlight_cache.borrow_mut() = None;
    }

    /// Truncates the syntax tree cache to maintain the MAX_SYNTAX_TREES policy.
    /// Keeps the oldest MAX_SYNTAX_TREES - 1 versions and the provided content_version.
    fn truncate_tree_state(
        syntax_tree_lock: &mut HashMap<BufferVersion, Tree>,
        buffer_version: BufferVersion,
    ) {
        if syntax_tree_lock.len() <= MAX_SYNTAX_TREES {
            return;
        }

        let mut versions: Vec<BufferVersion> = syntax_tree_lock.keys().copied().collect();
        versions.sort();

        let mut keep: HashSet<BufferVersion> = versions
            .iter()
            .take(MAX_SYNTAX_TREES - 1)
            .copied()
            .collect();
        keep.insert(buffer_version);

        syntax_tree_lock.retain(|v, _| keep.contains(v));
    }
}

impl DecorationLayer for SyntaxTreeState {
    fn update_internal_state_with_delta(
        &mut self,
        deltas: &[PreciseDelta],
        version: BufferVersion,
        content: BufferSnapshot,
        ctx: &mut ModelContext<Self>,
    ) {
        // If there is an active parsing in progress. Abort that first before starting another one.
        if let Some(handle) = self.parsing_handle.take() {
            handle.abort();
        }

        let Some(language) = self
            .language_queries
            .as_ref()
            .map(|language_queries| language_queries.language.clone())
        else {
            return;
        };

        let mut syntax_tree_lock = self.syntax_tree.lock();
        let mut tree = syntax_tree_lock.get(&self.buffer_version).cloned();
        if let Some(tree) = &mut tree {
            for delta in deltas {
                let edit = Self::delta_to_input_edit(delta);
                tree.edit(&edit);
            }

            // We write to the tree immediately after editing first to prevent flickering in the render
            // state before reparsing gets completed.
            if let Some(existing) = syntax_tree_lock.get_mut(&version) {
                existing.clone_from(tree);
            } else {
                syntax_tree_lock.insert(version, tree.clone());
                Self::truncate_tree_state(&mut syntax_tree_lock, version);
            }
        }

        let handle = ctx
            .spawn(
                async move {
                    let new_tree = Self::parse_text(content, tree, &language).await;
                    futures_lite::future::yield_now().await;
                    new_tree
                },
                move |model, new_tree, ctx| {
                    let mut syntax_tree_lock = model.syntax_tree.lock();
                    model.invalidate_highlight_cache_for_version(version);
                    if let Some(old_tree) = syntax_tree_lock.get_mut(&version) {
                        *old_tree = new_tree;
                    } else {
                        // This is for the case where we are updating the syntax tree for the first time.
                        syntax_tree_lock.insert(version, new_tree);
                        Self::truncate_tree_state(&mut syntax_tree_lock, model.buffer_version);
                    }
                    ctx.emit(DecorationStateEvent::DecorationUpdated { version });
                },
            )
            .abort_handle();

        self.buffer_version = version;
        self.parsing_handle = Some(handle);
    }
}

impl Entity for SyntaxTreeState {
    type Event = DecorationStateEvent;
}

/// Convert a 1-indexed buffer Point into a 0-indexed tree-sitter Point.
fn point_to_syntax_point(point: Point) -> arborium::tree_sitter::Point {
    // Subtracting 1 from row to convert from 1-indexed buffer rows to 0-indexed tree-sitter rows.
    arborium::tree_sitter::Point {
        row: point.row.saturating_sub(1) as usize,
        column: point.column as usize,
    }
}