//! This module defines types that represent changes to source code flowing from
//! the server to the client.
//!
//! It can be viewed as a dual for [`Change`][vfs::Change].

use std::{collections::hash_map::Entry, fmt, iter, mem};

use crate::text_edit::{TextEdit, TextEditBuilder};
use crate::{SnippetCap, assists::Command, syntax_helpers::tree_diff::diff};
use base_db::AnchoredPathBuf;
use itertools::Itertools;
use macros::UpmapFromRaFixture;
use nohash_hasher::IntMap;
use rustc_hash::FxHashMap;
use span::FileId;
use stdx::never;
use syntax::{
    AstNode, SyntaxNode, TextRange, TextSize,
    syntax_editor::{SyntaxAnnotation, SyntaxEditor},
};

/// An annotation ID associated with an indel, to describe changes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, UpmapFromRaFixture)]
pub struct ChangeAnnotationId(u32);

impl fmt::Display for ChangeAnnotationId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

#[derive(Debug, Clone)]
pub struct ChangeAnnotation {
    pub label: String,
    pub needs_confirmation: bool,
    pub description: Option<String>,
}

#[derive(Default, Debug, Clone)]
pub struct SourceChange {
    pub source_file_edits: IntMap<FileId, (TextEdit, Option<SnippetEdit>)>,
    pub file_system_edits: Vec<FileSystemEdit>,
    pub is_snippet: bool,
    pub annotations: FxHashMap<ChangeAnnotationId, ChangeAnnotation>,
    next_annotation_id: u32,
}

impl SourceChange {
    pub fn from_text_edit(file_id: impl Into<FileId>, edit: TextEdit) -> Self {
        SourceChange {
            source_file_edits: iter::once((file_id.into(), (edit, None))).collect(),
            ..Default::default()
        }
    }

    pub fn insert_annotation(&mut self, annotation: ChangeAnnotation) -> ChangeAnnotationId {
        let id = ChangeAnnotationId(self.next_annotation_id);
        self.next_annotation_id += 1;
        self.annotations.insert(id, annotation);
        id
    }

    /// Inserts a [`TextEdit`] for the given [`FileId`]. This properly handles merging existing
    /// edits for a file if some already exist.
    pub fn insert_source_edit(&mut self, file_id: impl Into<FileId>, edit: TextEdit) {
        self.insert_source_and_snippet_edit(file_id.into(), edit, None)
    }

    /// Inserts a [`TextEdit`] and potentially a [`SnippetEdit`] for the given [`FileId`].
    /// This properly handles merging existing edits for a file if some already exist.
    fn insert_source_and_snippet_edit(
        &mut self,
        file_id: impl Into<FileId>,
        edit: TextEdit,
        snippet_edit: Option<SnippetEdit>,
    ) {
        match self.source_file_edits.entry(file_id.into()) {
            Entry::Occupied(mut entry) => {
                let value = entry.get_mut();
                never!(value.0.union(edit).is_err(), "overlapping edits for same file");
                never!(
                    value.1.is_some() && snippet_edit.is_some(),
                    "overlapping snippet edits for same file"
                );
                if value.1.is_none() {
                    value.1 = snippet_edit;
                }
            }
            Entry::Vacant(entry) => {
                entry.insert((edit, snippet_edit));
            }
        }
    }

    pub fn push_file_system_edit(&mut self, edit: FileSystemEdit) {
        self.file_system_edits.push(edit);
    }

    pub fn get_source_and_snippet_edit(
        &self,
        file_id: FileId,
    ) -> Option<&(TextEdit, Option<SnippetEdit>)> {
        self.source_file_edits.get(&file_id)
    }

    pub fn merge(mut self, other: SourceChange) -> SourceChange {
        self.extend(other.source_file_edits);
        self.file_system_edits.extend(other.file_system_edits);
        self.is_snippet |= other.is_snippet;
        self
    }
}

impl Extend<(FileId, TextEdit)> for SourceChange {
    fn extend<T: IntoIterator<Item = (FileId, TextEdit)>>(&mut self, iter: T) {
        self.extend(iter.into_iter().map(|(file_id, edit)| (file_id, (edit, None))))
    }
}

impl Extend<(FileId, (TextEdit, Option<SnippetEdit>))> for SourceChange {
    fn extend<T: IntoIterator<Item = (FileId, (TextEdit, Option<SnippetEdit>))>>(
        &mut self,
        iter: T,
    ) {
        iter.into_iter().for_each(|(file_id, (edit, snippet_edit))| {
            self.insert_source_and_snippet_edit(file_id, edit, snippet_edit)
        });
    }
}

impl FromIterator<(FileId, TextEdit)> for SourceChange {
    fn from_iter<T: IntoIterator<Item = (FileId, TextEdit)>>(iter: T) -> Self {
        let mut this = SourceChange::default();
        this.extend(iter);
        this
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnippetEdit(Vec<(u32, TextRange)>);

impl SnippetEdit {
    pub fn new(snippets: Vec<Snippet>) -> Self {
        let mut snippet_ranges = snippets
            .into_iter()
            .zip(1..)
            .with_position()
            .flat_map(|(position, (snippet, index))| {
                // The last/only snippet gets index 0.
                let index = if position.is_last { 0 } else { index };

                match snippet {
                    Snippet::Tabstop(pos) => vec![(index, TextRange::empty(pos))],
                    Snippet::Placeholder(range) => vec![(index, range)],
                    Snippet::PlaceholderGroup(ranges) => {
                        ranges.into_iter().map(|range| (index, range)).collect()
                    }
                }
            })
            .collect_vec();

        snippet_ranges.sort_by_key(|(_, range)| range.start());

        // Ensure that none of the ranges overlap
        let disjoint_ranges = snippet_ranges
            .iter()
            .zip(snippet_ranges.iter().skip(1))
            .all(|((_, left), (_, right))| left.end() <= right.start() || left == right);
        stdx::always!(disjoint_ranges);

        SnippetEdit(snippet_ranges)
    }

    /// Inserts all of the snippets into the given text.
    pub fn apply(&self, text: &mut String) {
        // Start from the back so that we don't have to adjust ranges
        for (index, range) in self.0.iter().rev() {
            if range.is_empty() {
                // is a tabstop
                text.insert_str(range.start().into(), &format!("${index}"));
            } else {
                // is a placeholder
                text.insert(range.end().into(), '}');
                text.insert_str(range.start().into(), &format!("${{{index}:"));
            }
        }
    }

    /// Gets the underlying snippet index + text range
    /// Tabstops are represented by an empty range, and placeholders use the range that they were given
    pub fn into_edit_ranges(self) -> Vec<(u32, TextRange)> {
        self.0
    }

    /// Escapes `\` and `$` so that they don't get interpreted as snippet-specific constructs.
    ///
    /// Note that we don't need to escape the other characters that can be escaped,
    /// because they wouldn't be treated as snippet-specific constructs without '$'.
    pub fn escape_snippet_bits(text: &mut String) {
        stdx::replace(text, '\\', "\\\\");
        stdx::replace(text, '$', "\\$");
    }
}

pub struct SourceChangeBuilder {
    edit: TextEditBuilder,
    pub file_id: FileId,
    pub source_change: SourceChange,
    pub command: Option<Command>,

    /// Keeps track of all edits performed on each file
    file_editors: FxHashMap<FileId, SyntaxEditor>,
    /// Keeps track of which annotations correspond to which snippets
    snippet_annotations: Vec<(AnnotationSnippet, SyntaxAnnotation)>,
}

impl SourceChangeBuilder {
    pub fn new(file_id: impl Into<FileId>) -> SourceChangeBuilder {
        SourceChangeBuilder {
            edit: TextEdit::builder(),
            file_id: file_id.into(),
            source_change: SourceChange::default(),
            command: None,
            file_editors: FxHashMap::default(),
            snippet_annotations: vec![],
        }
    }

    pub fn edit_file(&mut self, file_id: impl Into<FileId>) {
        self.commit();
        self.file_id = file_id.into();
    }

    pub fn make_editor(&self, node: &SyntaxNode) -> SyntaxEditor {
        SyntaxEditor::new(node.tree_top()).0
    }

    pub fn add_file_edits(&mut self, file_id: impl Into<FileId>, editor: SyntaxEditor) {
        match self.file_editors.entry(file_id.into()) {
            Entry::Occupied(mut entry) => entry.get_mut().merge(editor),
            Entry::Vacant(entry) => {
                entry.insert(editor);
            }
        }
    }

    pub fn make_placeholder_snippet(&mut self, _cap: SnippetCap) -> SyntaxAnnotation {
        self.add_snippet_annotation(AnnotationSnippet::Over)
    }

    pub fn make_tabstop_before(&mut self, _cap: SnippetCap) -> SyntaxAnnotation {
        self.add_snippet_annotation(AnnotationSnippet::Before)
    }

    pub fn make_tabstop_after(&mut self, _cap: SnippetCap) -> SyntaxAnnotation {
        self.add_snippet_annotation(AnnotationSnippet::After)
    }

    fn commit(&mut self) {
        // Apply syntax editor edits
        for (file_id, editor) in mem::take(&mut self.file_editors) {
            let edit_result = editor.finish();
            let mut snippet_edit = vec![];

            // Find snippet edits
            for (kind, annotation) in &self.snippet_annotations {
                let elements = edit_result.find_annotation(*annotation);

                let snippet = match (kind, elements) {
                    (AnnotationSnippet::Before, [element]) => {
                        Snippet::Tabstop(element.text_range().start())
                    }
                    (AnnotationSnippet::After, [element]) => {
                        Snippet::Tabstop(element.text_range().end())
                    }
                    (AnnotationSnippet::Over, [element]) => {
                        Snippet::Placeholder(element.text_range())
                    }
                    (AnnotationSnippet::Over, elements) if !elements.is_empty() => {
                        Snippet::PlaceholderGroup(
                            elements.iter().map(|it| it.text_range()).collect(),
                        )
                    }
                    _ => continue,
                };

                snippet_edit.push(snippet);
            }

            let mut edit = TextEdit::builder();
            diff(edit_result.old_root(), edit_result.new_root()).into_text_edit(&mut edit);
            let edit = edit.finish();

            let snippet_edit =
                if !snippet_edit.is_empty() { Some(SnippetEdit::new(snippet_edit)) } else { None };

            if !edit.is_empty() || snippet_edit.is_some() {
                self.source_change.insert_source_and_snippet_edit(file_id, edit, snippet_edit);
            }
        }

        // Apply mutable edits
        let edit = mem::take(&mut self.edit).finish();
        if !edit.is_empty() {
            self.source_change.insert_source_edit(self.file_id, edit);
        }
    }

    /// Remove specified `range` of text.
    pub fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }
    /// Append specified `text` at the given `offset`
    pub fn insert(&mut self, offset: TextSize, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }
    /// Replaces specified `range` of text with a given string.
    pub fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }
    pub fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
        diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
    }
    pub fn create_file(&mut self, dst: AnchoredPathBuf, content: impl Into<String>) {
        let file_system_edit = FileSystemEdit::CreateFile { dst, initial_contents: content.into() };
        self.source_change.push_file_system_edit(file_system_edit);
    }
    pub fn move_file(&mut self, src: impl Into<FileId>, dst: AnchoredPathBuf) {
        let file_system_edit = FileSystemEdit::MoveFile { src: src.into(), dst };
        self.source_change.push_file_system_edit(file_system_edit);
    }

    /// Triggers the parameter hint popup after the assist is applied
    pub fn trigger_parameter_hints(&mut self) {
        self.command = Some(Command::TriggerParameterHints);
    }

    /// Renames the item at the cursor position after the assist is applied
    pub fn rename(&mut self) {
        self.command = Some(Command::Rename);
    }

    fn add_snippet_annotation(&mut self, kind: AnnotationSnippet) -> SyntaxAnnotation {
        let annotation = SyntaxAnnotation::default();
        self.snippet_annotations.push((kind, annotation));
        self.source_change.is_snippet = true;
        annotation
    }

    pub fn finish(mut self) -> SourceChange {
        self.commit();

        // Only one file can have snippet edits
        stdx::never!(
            self.source_change
                .source_file_edits
                .iter()
                .filter(|(_, (_, snippet_edit))| snippet_edit.is_some())
                .at_most_one()
                .is_err()
        );

        self.source_change
    }
}

#[derive(Debug, Clone)]
pub enum FileSystemEdit {
    CreateFile { dst: AnchoredPathBuf, initial_contents: String },
    MoveFile { src: FileId, dst: AnchoredPathBuf },
    MoveDir { src: AnchoredPathBuf, src_id: FileId, dst: AnchoredPathBuf },
}

impl From<FileSystemEdit> for SourceChange {
    fn from(edit: FileSystemEdit) -> SourceChange {
        SourceChange {
            source_file_edits: Default::default(),
            file_system_edits: vec![edit],
            is_snippet: false,
            ..SourceChange::default()
        }
    }
}

pub enum Snippet {
    /// A tabstop snippet (e.g. `$0`).
    Tabstop(TextSize),
    /// A placeholder snippet (e.g. `${0:placeholder}`).
    Placeholder(TextRange),
    /// A group of placeholder snippets, e.g.
    ///
    /// ```ignore
    /// let ${0:new_var} = 4;
    /// fun(1, 2, 3, ${0:new_var});
    /// ```
    PlaceholderGroup(Vec<TextRange>),
}

enum AnnotationSnippet {
    /// Place a tabstop before an element
    Before,
    /// Place a tabstop after an element
    After,
    /// Place a placeholder snippet in place of the element(s)
    Over,
}
