use crate::assemble_excerpts::assemble_excerpts;
use anyhow::Result;
use collections::HashMap;
use futures::{FutureExt, StreamExt as _, channel::mpsc, future};
use gpui::{App, AppContext, AsyncApp, Context, Entity, EventEmitter, Task, WeakEntity};
use language::{Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, Point, Rope, ToOffset as _};
use project::{LocationLink, Project, ProjectPath};
use serde::{Serialize, Serializer};
use smallvec::SmallVec;
use std::{
    collections::hash_map,
    ops::Range,
    sync::Arc,
    time::{Duration, Instant},
};
use util::{RangeExt as _, ResultExt};

mod assemble_excerpts;
#[cfg(test)]
mod edit_prediction_context_tests;
mod excerpt;
#[cfg(test)]
mod fake_definition_lsp;

pub use cloud_llm_client::predict_edits_v3::Line;
pub use excerpt::{EditPredictionExcerpt, EditPredictionExcerptOptions, EditPredictionExcerptText};

const IDENTIFIER_LINE_COUNT: u32 = 3;

pub struct RelatedExcerptStore {
    project: WeakEntity<Project>,
    related_files: Vec<RelatedFile>,
    cache: HashMap<Identifier, Arc<CacheEntry>>,
    update_tx: mpsc::UnboundedSender<(Entity<Buffer>, Anchor)>,
    identifier_line_count: u32,
}

pub enum RelatedExcerptStoreEvent {
    StartedRefresh,
    FinishedRefresh {
        cache_hit_count: usize,
        cache_miss_count: usize,
        mean_definition_latency: Duration,
        max_definition_latency: Duration,
    },
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct Identifier {
    pub name: String,
    pub range: Range<Anchor>,
}

enum DefinitionTask {
    CacheHit(Arc<CacheEntry>),
    CacheMiss(Task<Result<Option<Vec<LocationLink>>>>),
}

#[derive(Debug)]
struct CacheEntry {
    definitions: SmallVec<[CachedDefinition; 1]>,
}

#[derive(Clone, Debug)]
struct CachedDefinition {
    path: ProjectPath,
    buffer: Entity<Buffer>,
    anchor_range: Range<Anchor>,
}

#[derive(Clone, Debug, Serialize)]
pub struct RelatedFile {
    #[serde(serialize_with = "serialize_project_path")]
    pub path: ProjectPath,
    #[serde(skip)]
    pub buffer: WeakEntity<Buffer>,
    pub excerpts: Vec<RelatedExcerpt>,
    pub max_row: u32,
}

impl RelatedFile {
    pub fn merge_excerpts(&mut self) {
        self.excerpts.sort_unstable_by(|a, b| {
            a.point_range
                .start
                .cmp(&b.point_range.start)
                .then(b.point_range.end.cmp(&a.point_range.end))
        });

        let mut index = 1;
        while index < self.excerpts.len() {
            if self.excerpts[index - 1]
                .point_range
                .end
                .cmp(&self.excerpts[index].point_range.start)
                .is_ge()
            {
                let removed = self.excerpts.remove(index);
                if removed
                    .point_range
                    .end
                    .cmp(&self.excerpts[index - 1].point_range.end)
                    .is_gt()
                {
                    self.excerpts[index - 1].point_range.end = removed.point_range.end;
                    self.excerpts[index - 1].anchor_range.end = removed.anchor_range.end;
                }
            } else {
                index += 1;
            }
        }
    }
}

#[derive(Clone, Debug, Serialize)]
pub struct RelatedExcerpt {
    #[serde(skip)]
    pub anchor_range: Range<Anchor>,
    #[serde(serialize_with = "serialize_point_range")]
    pub point_range: Range<Point>,
    #[serde(serialize_with = "serialize_rope")]
    pub text: Rope,
}

fn serialize_project_path<S: Serializer>(
    project_path: &ProjectPath,
    serializer: S,
) -> Result<S::Ok, S::Error> {
    project_path.path.serialize(serializer)
}

fn serialize_rope<S: Serializer>(rope: &Rope, serializer: S) -> Result<S::Ok, S::Error> {
    rope.to_string().serialize(serializer)
}

fn serialize_point_range<S: Serializer>(
    range: &Range<Point>,
    serializer: S,
) -> Result<S::Ok, S::Error> {
    [
        [range.start.row, range.start.column],
        [range.end.row, range.end.column],
    ]
    .serialize(serializer)
}

const DEBOUNCE_DURATION: Duration = Duration::from_millis(100);

impl EventEmitter<RelatedExcerptStoreEvent> for RelatedExcerptStore {}

impl RelatedExcerptStore {
    pub fn new(project: &Entity<Project>, cx: &mut Context<Self>) -> Self {
        let (update_tx, mut update_rx) = mpsc::unbounded::<(Entity<Buffer>, Anchor)>();
        cx.spawn(async move |this, cx| {
            let executor = cx.background_executor().clone();
            while let Some((mut buffer, mut position)) = update_rx.next().await {
                let mut timer = executor.timer(DEBOUNCE_DURATION).fuse();
                loop {
                    futures::select_biased! {
                        next = update_rx.next() => {
                            if let Some((new_buffer, new_position)) = next {
                                buffer = new_buffer;
                                position = new_position;
                                timer = executor.timer(DEBOUNCE_DURATION).fuse();
                            } else {
                                return anyhow::Ok(());
                            }
                        }
                        _ = timer => break,
                    }
                }

                Self::fetch_excerpts(this.clone(), buffer, position, cx).await?;
            }
            anyhow::Ok(())
        })
        .detach_and_log_err(cx);

        RelatedExcerptStore {
            project: project.downgrade(),
            update_tx,
            related_files: Vec::new(),
            cache: Default::default(),
            identifier_line_count: IDENTIFIER_LINE_COUNT,
        }
    }

    pub fn set_identifier_line_count(&mut self, count: u32) {
        self.identifier_line_count = count;
    }

    pub fn refresh(&mut self, buffer: Entity<Buffer>, position: Anchor, _: &mut Context<Self>) {
        self.update_tx.unbounded_send((buffer, position)).ok();
    }

    pub fn related_files(&self) -> &[RelatedFile] {
        &self.related_files
    }

    async fn fetch_excerpts(
        this: WeakEntity<Self>,
        buffer: Entity<Buffer>,
        position: Anchor,
        cx: &mut AsyncApp,
    ) -> Result<()> {
        let (project, snapshot, identifier_line_count) = this.read_with(cx, |this, cx| {
            (
                this.project.upgrade(),
                buffer.read(cx).snapshot(),
                this.identifier_line_count,
            )
        })?;
        let Some(project) = project else {
            return Ok(());
        };

        let file = snapshot.file().cloned();
        if let Some(file) = &file {
            log::debug!("retrieving_context buffer:{}", file.path().as_unix_str());
        }

        this.update(cx, |_, cx| {
            cx.emit(RelatedExcerptStoreEvent::StartedRefresh);
        })?;

        let identifiers = cx
            .background_spawn(async move {
                identifiers_for_position(&snapshot, position, identifier_line_count)
            })
            .await;

        let async_cx = cx.clone();
        let start_time = Instant::now();
        let futures = this.update(cx, |this, cx| {
            identifiers
                .into_iter()
                .filter_map(|identifier| {
                    let task = if let Some(entry) = this.cache.get(&identifier) {
                        DefinitionTask::CacheHit(entry.clone())
                    } else {
                        DefinitionTask::CacheMiss(
                            this.project
                                .update(cx, |project, cx| {
                                    project.definitions(&buffer, identifier.range.start, cx)
                                })
                                .ok()?,
                        )
                    };

                    let cx = async_cx.clone();
                    let project = project.clone();
                    Some(async move {
                        match task {
                            DefinitionTask::CacheHit(cache_entry) => {
                                Some((identifier, cache_entry, None))
                            }
                            DefinitionTask::CacheMiss(task) => {
                                let locations = task.await.log_err()??;
                                let duration = start_time.elapsed();
                                cx.update(|cx| {
                                    (
                                        identifier,
                                        Arc::new(CacheEntry {
                                            definitions: locations
                                                .into_iter()
                                                .filter_map(|location| {
                                                    process_definition(location, &project, cx)
                                                })
                                                .collect(),
                                        }),
                                        Some(duration),
                                    )
                                })
                                .ok()
                            }
                        }
                    })
                })
                .collect::<Vec<_>>()
        })?;

        let mut cache_hit_count = 0;
        let mut cache_miss_count = 0;
        let mut mean_definition_latency = Duration::ZERO;
        let mut max_definition_latency = Duration::ZERO;
        let mut new_cache = HashMap::default();
        new_cache.reserve(futures.len());
        for (identifier, entry, duration) in future::join_all(futures).await.into_iter().flatten() {
            new_cache.insert(identifier, entry);
            if let Some(duration) = duration {
                cache_miss_count += 1;
                mean_definition_latency += duration;
                max_definition_latency = max_definition_latency.max(duration);
            } else {
                cache_hit_count += 1;
            }
        }
        mean_definition_latency /= cache_miss_count.max(1) as u32;

        let (new_cache, related_files) = rebuild_related_files(new_cache, cx).await?;

        if let Some(file) = &file {
            log::debug!(
                "finished retrieving context buffer:{}, latency:{:?}",
                file.path().as_unix_str(),
                start_time.elapsed()
            );
        }

        this.update(cx, |this, cx| {
            this.cache = new_cache;
            this.related_files = related_files;
            cx.emit(RelatedExcerptStoreEvent::FinishedRefresh {
                cache_hit_count,
                cache_miss_count,
                mean_definition_latency,
                max_definition_latency,
            });
        })?;

        anyhow::Ok(())
    }
}

async fn rebuild_related_files(
    new_entries: HashMap<Identifier, Arc<CacheEntry>>,
    cx: &mut AsyncApp,
) -> Result<(HashMap<Identifier, Arc<CacheEntry>>, Vec<RelatedFile>)> {
    let mut snapshots = HashMap::default();
    for entry in new_entries.values() {
        for definition in &entry.definitions {
            if let hash_map::Entry::Vacant(e) = snapshots.entry(definition.buffer.entity_id()) {
                definition
                    .buffer
                    .read_with(cx, |buffer, _| buffer.parsing_idle())?
                    .await;
                e.insert(
                    definition
                        .buffer
                        .read_with(cx, |buffer, _| buffer.snapshot())?,
                );
            }
        }
    }

    Ok(cx
        .background_spawn(async move {
            let mut files = Vec::<RelatedFile>::new();
            let mut ranges_by_buffer = HashMap::<_, Vec<Range<Point>>>::default();
            let mut paths_by_buffer = HashMap::default();
            for entry in new_entries.values() {
                for definition in &entry.definitions {
                    let Some(snapshot) = snapshots.get(&definition.buffer.entity_id()) else {
                        continue;
                    };
                    paths_by_buffer.insert(definition.buffer.entity_id(), definition.path.clone());
                    ranges_by_buffer
                        .entry(definition.buffer.clone())
                        .or_default()
                        .push(definition.anchor_range.to_point(snapshot));
                }
            }

            for (buffer, ranges) in ranges_by_buffer {
                let Some(snapshot) = snapshots.get(&buffer.entity_id()) else {
                    continue;
                };
                let Some(project_path) = paths_by_buffer.get(&buffer.entity_id()) else {
                    continue;
                };
                let excerpts = assemble_excerpts(snapshot, ranges);
                files.push(RelatedFile {
                    path: project_path.clone(),
                    buffer: buffer.downgrade(),
                    excerpts,
                    max_row: snapshot.max_point().row,
                });
            }

            files.sort_by_key(|file| file.path.clone());
            (new_entries, files)
        })
        .await)
}

const MAX_TARGET_LEN: usize = 128;

fn process_definition(
    location: LocationLink,
    project: &Entity<Project>,
    cx: &mut App,
) -> Option<CachedDefinition> {
    let buffer = location.target.buffer.read(cx);
    let anchor_range = location.target.range;
    let file = buffer.file()?;
    let worktree = project.read(cx).worktree_for_id(file.worktree_id(cx), cx)?;
    if worktree.read(cx).is_single_file() {
        return None;
    }

    // If the target range is large, it likely means we requested the definition of an entire module.
    // For individual definitions, the target range should be small as it only covers the symbol.
    let buffer = location.target.buffer.read(cx);
    let target_len = anchor_range.to_offset(&buffer).len();
    if target_len > MAX_TARGET_LEN {
        return None;
    }

    Some(CachedDefinition {
        path: ProjectPath {
            worktree_id: file.worktree_id(cx),
            path: file.path().clone(),
        },
        buffer: location.target.buffer,
        anchor_range,
    })
}

/// Gets all of the identifiers that are present in the given line, and its containing
/// outline items.
fn identifiers_for_position(
    buffer: &BufferSnapshot,
    position: Anchor,
    identifier_line_count: u32,
) -> Vec<Identifier> {
    let offset = position.to_offset(buffer);
    let point = buffer.offset_to_point(offset);

    // Search for identifiers on lines adjacent to the cursor.
    let start = Point::new(point.row.saturating_sub(identifier_line_count), 0);
    let end = Point::new(point.row + identifier_line_count + 1, 0).min(buffer.max_point());
    let line_range = start..end;
    let mut ranges = vec![line_range.to_offset(&buffer)];

    // Search for identifiers mentioned in headers/signatures of containing outline items.
    let outline_items = buffer.outline_items_as_offsets_containing(offset..offset, false, None);
    for item in outline_items {
        if let Some(body_range) = item.body_range(&buffer) {
            ranges.push(item.range.start..body_range.start.to_offset(&buffer));
        } else {
            ranges.push(item.range.clone());
        }
    }

    ranges.sort_by(|a, b| a.start.cmp(&b.start).then(b.end.cmp(&a.end)));
    ranges.dedup_by(|a, b| {
        if a.start <= b.end {
            b.start = b.start.min(a.start);
            b.end = b.end.max(a.end);
            true
        } else {
            false
        }
    });

    let mut identifiers = Vec::new();
    let outer_range =
        ranges.first().map_or(0, |r| r.start)..ranges.last().map_or(buffer.len(), |r| r.end);

    let mut captures = buffer
        .syntax
        .captures(outer_range.clone(), &buffer.text, |grammar| {
            grammar
                .highlights_config
                .as_ref()
                .map(|config| &config.query)
        });

    for range in ranges {
        captures.set_byte_range(range.start..outer_range.end);

        let mut last_range = None;
        while let Some(capture) = captures.peek() {
            let node_range = capture.node.byte_range();
            if node_range.start > range.end {
                break;
            }
            let config = captures.grammars()[capture.grammar_index]
                .highlights_config
                .as_ref();

            if let Some(config) = config
                && config.identifier_capture_indices.contains(&capture.index)
                && range.contains_inclusive(&node_range)
                && Some(&node_range) != last_range.as_ref()
            {
                let name = buffer.text_for_range(node_range.clone()).collect();
                identifiers.push(Identifier {
                    range: buffer.anchor_after(node_range.start)
                        ..buffer.anchor_before(node_range.end),
                    name,
                });
                last_range = Some(node_range);
            }

            captures.advance();
        }
    }

    identifiers
}