mod create_file_parser;
mod edit_parser;
#[cfg(test)]
mod evals;
mod streaming_fuzzy_matcher;

use crate::{Template, Templates};
use anyhow::Result;
use assistant_tool::ActionLog;
use create_file_parser::{CreateFileParser, CreateFileParserEvent};
use edit_parser::{EditParser, EditParserEvent, EditParserMetrics};
use futures::{
    Stream, StreamExt,
    channel::mpsc::{self, UnboundedReceiver},
    pin_mut,
    stream::BoxStream,
};
use gpui::{AppContext, AsyncApp, Entity, Task};
use language::{Anchor, Buffer, BufferSnapshot, LineIndent, Point, TextBufferSnapshot};
use language_model::{
    LanguageModel, LanguageModelCompletionError, LanguageModelRequest, LanguageModelRequestMessage,
    LanguageModelToolChoice, MessageContent, Role,
};
use project::{AgentLocation, Project};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::{cmp, iter, mem, ops::Range, path::PathBuf, pin::Pin, sync::Arc, task::Poll};
use streaming_diff::{CharOperation, StreamingDiff};
use streaming_fuzzy_matcher::StreamingFuzzyMatcher;
use util::debug_panic;

#[derive(Serialize)]
struct CreateFilePromptTemplate {
    path: Option<PathBuf>,
    edit_description: String,
}

impl Template for CreateFilePromptTemplate {
    const TEMPLATE_NAME: &'static str = "create_file_prompt.hbs";
}

#[derive(Serialize)]
struct EditFilePromptTemplate {
    path: Option<PathBuf>,
    edit_description: String,
}

impl Template for EditFilePromptTemplate {
    const TEMPLATE_NAME: &'static str = "edit_file_prompt.hbs";
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum EditAgentOutputEvent {
    ResolvingEditRange(Range<Anchor>),
    UnresolvedEditRange,
    Edited,
}

#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
pub struct EditAgentOutput {
    pub raw_edits: String,
    pub parser_metrics: EditParserMetrics,
}

#[derive(Clone)]
pub struct EditAgent {
    model: Arc<dyn LanguageModel>,
    action_log: Entity<ActionLog>,
    project: Entity<Project>,
    templates: Arc<Templates>,
}

impl EditAgent {
    pub fn new(
        model: Arc<dyn LanguageModel>,
        project: Entity<Project>,
        action_log: Entity<ActionLog>,
        templates: Arc<Templates>,
    ) -> Self {
        EditAgent {
            model,
            project,
            action_log,
            templates,
        }
    }

    pub fn overwrite(
        &self,
        buffer: Entity<Buffer>,
        edit_description: String,
        conversation: &LanguageModelRequest,
        cx: &mut AsyncApp,
    ) -> (
        Task<Result<EditAgentOutput>>,
        mpsc::UnboundedReceiver<EditAgentOutputEvent>,
    ) {
        let this = self.clone();
        let (events_tx, events_rx) = mpsc::unbounded();
        let conversation = conversation.clone();
        let output = cx.spawn(async move |cx| {
            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;
            let path = cx.update(|cx| snapshot.resolve_file_path(cx, true))?;
            let prompt = CreateFilePromptTemplate {
                path,
                edit_description,
            }
            .render(&this.templates)?;
            let new_chunks = this.request(conversation, prompt, cx).await?;

            let (output, mut inner_events) = this.overwrite_with_chunks(buffer, new_chunks, cx);
            while let Some(event) = inner_events.next().await {
                events_tx.unbounded_send(event).ok();
            }
            output.await
        });
        (output, events_rx)
    }

    fn overwrite_with_chunks(
        &self,
        buffer: Entity<Buffer>,
        edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
        cx: &mut AsyncApp,
    ) -> (
        Task<Result<EditAgentOutput>>,
        mpsc::UnboundedReceiver<EditAgentOutputEvent>,
    ) {
        let (output_events_tx, output_events_rx) = mpsc::unbounded();
        let (parse_task, parse_rx) = Self::parse_create_file_chunks(edit_chunks, cx);
        let this = self.clone();
        let task = cx.spawn(async move |cx| {
            this.action_log
                .update(cx, |log, cx| log.buffer_created(buffer.clone(), cx))?;
            this.overwrite_with_chunks_internal(buffer, parse_rx, output_events_tx, cx)
                .await?;
            parse_task.await
        });
        (task, output_events_rx)
    }

    async fn overwrite_with_chunks_internal(
        &self,
        buffer: Entity<Buffer>,
        mut parse_rx: UnboundedReceiver<Result<CreateFileParserEvent>>,
        output_events_tx: mpsc::UnboundedSender<EditAgentOutputEvent>,
        cx: &mut AsyncApp,
    ) -> Result<()> {
        cx.update(|cx| {
            buffer.update(cx, |buffer, cx| buffer.set_text("", cx));
            self.action_log.update(cx, |log, cx| {
                log.buffer_edited(buffer.clone(), cx);
            });
            self.project.update(cx, |project, cx| {
                project.set_agent_location(
                    Some(AgentLocation {
                        buffer: buffer.downgrade(),
                        position: language::Anchor::MAX,
                    }),
                    cx,
                )
            });
            output_events_tx
                .unbounded_send(EditAgentOutputEvent::Edited)
                .ok();
        })?;

        while let Some(event) = parse_rx.next().await {
            match event? {
                CreateFileParserEvent::NewTextChunk { chunk } => {
                    cx.update(|cx| {
                        buffer.update(cx, |buffer, cx| buffer.append(chunk, cx));
                        self.action_log
                            .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
                        self.project.update(cx, |project, cx| {
                            project.set_agent_location(
                                Some(AgentLocation {
                                    buffer: buffer.downgrade(),
                                    position: language::Anchor::MAX,
                                }),
                                cx,
                            )
                        });
                    })?;
                    output_events_tx
                        .unbounded_send(EditAgentOutputEvent::Edited)
                        .ok();
                }
            }
        }

        Ok(())
    }

    pub fn edit(
        &self,
        buffer: Entity<Buffer>,
        edit_description: String,
        conversation: &LanguageModelRequest,
        cx: &mut AsyncApp,
    ) -> (
        Task<Result<EditAgentOutput>>,
        mpsc::UnboundedReceiver<EditAgentOutputEvent>,
    ) {
        let this = self.clone();
        let (events_tx, events_rx) = mpsc::unbounded();
        let conversation = conversation.clone();
        let output = cx.spawn(async move |cx| {
            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;
            let path = cx.update(|cx| snapshot.resolve_file_path(cx, true))?;
            let prompt = EditFilePromptTemplate {
                path,
                edit_description,
            }
            .render(&this.templates)?;
            let edit_chunks = this.request(conversation, prompt, cx).await?;
            this.apply_edit_chunks(buffer, edit_chunks, events_tx, cx)
                .await
        });
        (output, events_rx)
    }

    async fn apply_edit_chunks(
        &self,
        buffer: Entity<Buffer>,
        edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
        output_events: mpsc::UnboundedSender<EditAgentOutputEvent>,
        cx: &mut AsyncApp,
    ) -> Result<EditAgentOutput> {
        self.action_log
            .update(cx, |log, cx| log.buffer_read(buffer.clone(), cx))?;

        let (output, edit_events) = Self::parse_edit_chunks(edit_chunks, cx);
        let mut edit_events = edit_events.peekable();
        while let Some(edit_event) = Pin::new(&mut edit_events).peek().await {
            // Skip events until we're at the start of a new edit.
            let Ok(EditParserEvent::OldTextChunk { .. }) = edit_event else {
                edit_events.next().await.unwrap()?;
                continue;
            };

            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;

            // Resolve the old text in the background, updating the agent
            // location as we keep refining which range it corresponds to.
            let (resolve_old_text, mut old_range) =
                Self::resolve_old_text(snapshot.text.clone(), edit_events, cx);
            while let Ok(old_range) = old_range.recv().await {
                if let Some(old_range) = old_range {
                    let old_range = snapshot.anchor_before(old_range.start)
                        ..snapshot.anchor_before(old_range.end);
                    self.project.update(cx, |project, cx| {
                        project.set_agent_location(
                            Some(AgentLocation {
                                buffer: buffer.downgrade(),
                                position: old_range.end,
                            }),
                            cx,
                        );
                    })?;
                    output_events
                        .unbounded_send(EditAgentOutputEvent::ResolvingEditRange(old_range))
                        .ok();
                }
            }

            let (edit_events_, resolved_old_text) = resolve_old_text.await?;
            edit_events = edit_events_;

            // If we can't resolve the old text, restart the loop waiting for a
            // new edit (or for the stream to end).
            let Some(resolved_old_text) = resolved_old_text else {
                output_events
                    .unbounded_send(EditAgentOutputEvent::UnresolvedEditRange)
                    .ok();
                continue;
            };

            // Compute edits in the background and apply them as they become
            // available.
            let (compute_edits, edits) =
                Self::compute_edits(snapshot, resolved_old_text, edit_events, cx);
            let mut edits = edits.ready_chunks(32);
            while let Some(edits) = edits.next().await {
                if edits.is_empty() {
                    continue;
                }

                // Edit the buffer and report edits to the action log as part of the
                // same effect cycle, otherwise the edit will be reported as if the
                // user made it.
                cx.update(|cx| {
                    let max_edit_end = buffer.update(cx, |buffer, cx| {
                        buffer.edit(edits.iter().cloned(), None, cx);
                        let max_edit_end = buffer
                            .summaries_for_anchors::<Point, _>(
                                edits.iter().map(|(range, _)| &range.end),
                            )
                            .max()
                            .unwrap();
                        buffer.anchor_before(max_edit_end)
                    });
                    self.action_log
                        .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));
                    self.project.update(cx, |project, cx| {
                        project.set_agent_location(
                            Some(AgentLocation {
                                buffer: buffer.downgrade(),
                                position: max_edit_end,
                            }),
                            cx,
                        );
                    });
                })?;
                output_events
                    .unbounded_send(EditAgentOutputEvent::Edited)
                    .ok();
            }

            edit_events = compute_edits.await?;
        }

        output.await
    }

    fn parse_edit_chunks(
        chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
        cx: &mut AsyncApp,
    ) -> (
        Task<Result<EditAgentOutput>>,
        UnboundedReceiver<Result<EditParserEvent>>,
    ) {
        let (tx, rx) = mpsc::unbounded();
        let output = cx.background_spawn(async move {
            pin_mut!(chunks);

            let mut parser = EditParser::new();
            let mut raw_edits = String::new();
            while let Some(chunk) = chunks.next().await {
                match chunk {
                    Ok(chunk) => {
                        raw_edits.push_str(&chunk);
                        for event in parser.push(&chunk) {
                            tx.unbounded_send(Ok(event))?;
                        }
                    }
                    Err(error) => {
                        tx.unbounded_send(Err(error.into()))?;
                    }
                }
            }
            Ok(EditAgentOutput {
                raw_edits,
                parser_metrics: parser.finish(),
            })
        });
        (output, rx)
    }

    fn parse_create_file_chunks(
        chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,
        cx: &mut AsyncApp,
    ) -> (
        Task<Result<EditAgentOutput>>,
        UnboundedReceiver<Result<CreateFileParserEvent>>,
    ) {
        let (tx, rx) = mpsc::unbounded();
        let output = cx.background_spawn(async move {
            pin_mut!(chunks);

            let mut parser = CreateFileParser::new();
            let mut raw_edits = String::new();
            while let Some(chunk) = chunks.next().await {
                match chunk {
                    Ok(chunk) => {
                        raw_edits.push_str(&chunk);
                        for event in parser.push(Some(&chunk)) {
                            tx.unbounded_send(Ok(event))?;
                        }
                    }
                    Err(error) => {
                        tx.unbounded_send(Err(error.into()))?;
                    }
                }
            }
            // Send final events with None to indicate completion
            for event in parser.push(None) {
                tx.unbounded_send(Ok(event))?;
            }
            Ok(EditAgentOutput {
                raw_edits,
                parser_metrics: EditParserMetrics::default(),
            })
        });
        (output, rx)
    }

    fn resolve_old_text<T>(
        snapshot: TextBufferSnapshot,
        mut edit_events: T,
        cx: &mut AsyncApp,
    ) -> (
        Task<Result<(T, Option<ResolvedOldText>)>>,
        async_watch::Receiver<Option<Range<usize>>>,
    )
    where
        T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,
    {
        let (old_range_tx, old_range_rx) = async_watch::channel(None);
        let task = cx.background_spawn(async move {
            let mut matcher = StreamingFuzzyMatcher::new(snapshot);
            while let Some(edit_event) = edit_events.next().await {
                let EditParserEvent::OldTextChunk { chunk, done } = edit_event? else {
                    break;
                };

                old_range_tx.send(matcher.push(&chunk))?;
                if done {
                    break;
                }
            }

            let old_range = matcher.finish();
            old_range_tx.send(old_range.clone())?;
            if let Some(old_range) = old_range {
                let line_indent =
                    LineIndent::from_iter(matcher.query_lines().first().unwrap().chars());
                Ok((
                    edit_events,
                    Some(ResolvedOldText {
                        range: old_range,
                        indent: line_indent,
                    }),
                ))
            } else {
                Ok((edit_events, None))
            }
        });

        (task, old_range_rx)
    }

    fn compute_edits<T>(
        snapshot: BufferSnapshot,
        resolved_old_text: ResolvedOldText,
        mut edit_events: T,
        cx: &mut AsyncApp,
    ) -> (
        Task<Result<T>>,
        UnboundedReceiver<(Range<Anchor>, Arc<str>)>,
    )
    where
        T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,
    {
        let (edits_tx, edits_rx) = mpsc::unbounded();
        let compute_edits = cx.background_spawn(async move {
            let buffer_start_indent = snapshot
                .line_indent_for_row(snapshot.offset_to_point(resolved_old_text.range.start).row);
            let indent_delta = if buffer_start_indent.tabs > 0 {
                IndentDelta::Tabs(
                    buffer_start_indent.tabs as isize - resolved_old_text.indent.tabs as isize,
                )
            } else {
                IndentDelta::Spaces(
                    buffer_start_indent.spaces as isize - resolved_old_text.indent.spaces as isize,
                )
            };

            let old_text = snapshot
                .text_for_range(resolved_old_text.range.clone())
                .collect::<String>();
            let mut diff = StreamingDiff::new(old_text);
            let mut edit_start = resolved_old_text.range.start;
            let mut new_text_chunks =
                Self::reindent_new_text_chunks(indent_delta, &mut edit_events);
            let mut done = false;
            while !done {
                let char_operations = if let Some(new_text_chunk) = new_text_chunks.next().await {
                    diff.push_new(&new_text_chunk?)
                } else {
                    done = true;
                    mem::take(&mut diff).finish()
                };

                for op in char_operations {
                    match op {
                        CharOperation::Insert { text } => {
                            let edit_start = snapshot.anchor_after(edit_start);
                            edits_tx.unbounded_send((edit_start..edit_start, Arc::from(text)))?;
                        }
                        CharOperation::Delete { bytes } => {
                            let edit_end = edit_start + bytes;
                            let edit_range =
                                snapshot.anchor_after(edit_start)..snapshot.anchor_before(edit_end);
                            edit_start = edit_end;
                            edits_tx.unbounded_send((edit_range, Arc::from("")))?;
                        }
                        CharOperation::Keep { bytes } => edit_start += bytes,
                    }
                }
            }

            drop(new_text_chunks);
            anyhow::Ok(edit_events)
        });

        (compute_edits, edits_rx)
    }

    fn reindent_new_text_chunks(
        delta: IndentDelta,
        mut stream: impl Unpin + Stream<Item = Result<EditParserEvent>>,
    ) -> impl Stream<Item = Result<String>> {
        let mut buffer = String::new();
        let mut in_leading_whitespace = true;
        let mut done = false;
        futures::stream::poll_fn(move |cx| {
            while !done {
                let (chunk, is_last_chunk) = match stream.poll_next_unpin(cx) {
                    Poll::Ready(Some(Ok(EditParserEvent::NewTextChunk { chunk, done }))) => {
                        (chunk, done)
                    }
                    Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err))),
                    Poll::Pending => return Poll::Pending,
                    _ => return Poll::Ready(None),
                };

                buffer.push_str(&chunk);

                let mut indented_new_text = String::new();
                let mut start_ix = 0;
                let mut newlines = buffer.match_indices('\n').peekable();
                loop {
                    let (line_end, is_pending_line) = match newlines.next() {
                        Some((ix, _)) => (ix, false),
                        None => (buffer.len(), true),
                    };
                    let line = &buffer[start_ix..line_end];

                    if in_leading_whitespace {
                        if let Some(non_whitespace_ix) = line.find(|c| delta.character() != c) {
                            // We found a non-whitespace character, adjust
                            // indentation based on the delta.
                            let new_indent_len =
                                cmp::max(0, non_whitespace_ix as isize + delta.len()) as usize;
                            indented_new_text
                                .extend(iter::repeat(delta.character()).take(new_indent_len));
                            indented_new_text.push_str(&line[non_whitespace_ix..]);
                            in_leading_whitespace = false;
                        } else if is_pending_line {
                            // We're still in leading whitespace and this line is incomplete.
                            // Stop processing until we receive more input.
                            break;
                        } else {
                            // This line is entirely whitespace. Push it without indentation.
                            indented_new_text.push_str(line);
                        }
                    } else {
                        indented_new_text.push_str(line);
                    }

                    if is_pending_line {
                        start_ix = line_end;
                        break;
                    } else {
                        in_leading_whitespace = true;
                        indented_new_text.push('\n');
                        start_ix = line_end + 1;
                    }
                }
                buffer.replace_range(..start_ix, "");

                // This was the last chunk, push all the buffered content as-is.
                if is_last_chunk {
                    indented_new_text.push_str(&buffer);
                    buffer.clear();
                    done = true;
                }

                if !indented_new_text.is_empty() {
                    return Poll::Ready(Some(Ok(indented_new_text)));
                }
            }

            Poll::Ready(None)
        })
    }

    async fn request(
        &self,
        mut conversation: LanguageModelRequest,
        prompt: String,
        cx: &mut AsyncApp,
    ) -> Result<BoxStream<'static, Result<String, LanguageModelCompletionError>>> {
        let mut messages_iter = conversation.messages.iter_mut();
        if let Some(last_message) = messages_iter.next_back() {
            if last_message.role == Role::Assistant {
                let old_content_len = last_message.content.len();
                last_message
                    .content
                    .retain(|content| !matches!(content, MessageContent::ToolUse(_)));
                let new_content_len = last_message.content.len();

                // We just removed pending tool uses from the content of the
                // last message, so it doesn't make sense to cache it anymore
                // (e.g., the message will look very different on the next
                // request). Thus, we move the flag to the message prior to it,
                // as it will still be a valid prefix of the conversation.
                if old_content_len != new_content_len && last_message.cache {
                    if let Some(prev_message) = messages_iter.next_back() {
                        last_message.cache = false;
                        prev_message.cache = true;
                    }
                }

                if last_message.content.is_empty() {
                    conversation.messages.pop();
                }
            } else {
                debug_panic!(
                    "Last message must be an Assistant tool calling! Got {:?}",
                    last_message.content
                );
            }
        }

        conversation.messages.push(LanguageModelRequestMessage {
            role: Role::User,
            content: vec![MessageContent::Text(prompt)],
            cache: false,
        });

        // Include tools in the request so that we can take advantage of
        // caching when ToolChoice::None is supported.
        let mut tool_choice = None;
        let mut tools = Vec::new();
        if !conversation.tools.is_empty()
            && self
                .model
                .supports_tool_choice(LanguageModelToolChoice::None)
        {
            tool_choice = Some(LanguageModelToolChoice::None);
            tools = conversation.tools.clone();
        }

        let request = LanguageModelRequest {
            thread_id: conversation.thread_id,
            prompt_id: conversation.prompt_id,
            mode: conversation.mode,
            messages: conversation.messages,
            tool_choice,
            tools,
            stop: Vec::new(),
            temperature: None,
        };

        Ok(self.model.stream_completion_text(request, cx).await?.stream)
    }
}

struct ResolvedOldText {
    range: Range<usize>,
    indent: LineIndent,
}

#[derive(Copy, Clone, Debug)]
enum IndentDelta {
    Spaces(isize),
    Tabs(isize),
}

impl IndentDelta {
    fn character(&self) -> char {
        match self {
            IndentDelta::Spaces(_) => ' ',
            IndentDelta::Tabs(_) => '\t',
        }
    }

    fn len(&self) -> isize {
        match self {
            IndentDelta::Spaces(n) => *n,
            IndentDelta::Tabs(n) => *n,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use fs::FakeFs;
    use futures::stream;
    use gpui::{AppContext, TestAppContext};
    use indoc::indoc;
    use language_model::fake_provider::FakeLanguageModel;
    use project::{AgentLocation, Project};
    use rand::prelude::*;
    use rand::rngs::StdRng;
    use std::cmp;

    #[gpui::test(iterations = 100)]
    async fn test_empty_old_text(cx: &mut TestAppContext, mut rng: StdRng) {
        let agent = init_test(cx).await;
        let buffer = cx.new(|cx| {
            Buffer::local(
                indoc! {"
                    abc
                    def
                    ghi
                "},
                cx,
            )
        });
        let (apply, _events) = agent.edit(
            buffer.clone(),
            String::new(),
            &LanguageModelRequest::default(),
            &mut cx.to_async(),
        );
        cx.run_until_parked();

        simulate_llm_output(
            &agent,
            indoc! {"
                <old_text></old_text>
                <new_text>jkl</new_text>
                <old_text>def</old_text>
                <new_text>DEF</new_text>
            "},
            &mut rng,
            cx,
        );
        apply.await.unwrap();

        pretty_assertions::assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            indoc! {"
                abc
                DEF
                ghi
            "}
        );
    }

    #[gpui::test(iterations = 100)]
    async fn test_indentation(cx: &mut TestAppContext, mut rng: StdRng) {
        let agent = init_test(cx).await;
        let buffer = cx.new(|cx| {
            Buffer::local(
                indoc! {"
                    lorem
                            ipsum
                            dolor
                            sit
                "},
                cx,
            )
        });
        let (apply, _events) = agent.edit(
            buffer.clone(),
            String::new(),
            &LanguageModelRequest::default(),
            &mut cx.to_async(),
        );
        cx.run_until_parked();

        simulate_llm_output(
            &agent,
            indoc! {"
                <old_text>
                    ipsum
                    dolor
                    sit
                </old_text>
                <new_text>
                    ipsum
                    dolor
                    sit
                amet
                </new_text>
            "},
            &mut rng,
            cx,
        );
        apply.await.unwrap();

        pretty_assertions::assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            indoc! {"
                lorem
                        ipsum
                        dolor
                        sit
                    amet
            "}
        );
    }

    #[gpui::test(iterations = 100)]
    async fn test_dependent_edits(cx: &mut TestAppContext, mut rng: StdRng) {
        let agent = init_test(cx).await;
        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
        let (apply, _events) = agent.edit(
            buffer.clone(),
            String::new(),
            &LanguageModelRequest::default(),
            &mut cx.to_async(),
        );
        cx.run_until_parked();

        simulate_llm_output(
            &agent,
            indoc! {"
                <old_text>
                def
                </old_text>
                <new_text>
                DEF
                </new_text>

                <old_text>
                DEF
                </old_text>
                <new_text>
                DeF
                </new_text>
            "},
            &mut rng,
            cx,
        );
        apply.await.unwrap();

        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abc\nDeF\nghi"
        );
    }

    #[gpui::test(iterations = 100)]
    async fn test_old_text_hallucination(cx: &mut TestAppContext, mut rng: StdRng) {
        let agent = init_test(cx).await;
        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
        let (apply, _events) = agent.edit(
            buffer.clone(),
            String::new(),
            &LanguageModelRequest::default(),
            &mut cx.to_async(),
        );
        cx.run_until_parked();

        simulate_llm_output(
            &agent,
            indoc! {"
                <old_text>
                jkl
                </old_text>
                <new_text>
                mno
                </new_text>

                <old_text>
                abc
                </old_text>
                <new_text>
                ABC
                </new_text>
            "},
            &mut rng,
            cx,
        );
        apply.await.unwrap();

        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "ABC\ndef\nghi"
        );
    }

    #[gpui::test]
    async fn test_edit_events(cx: &mut TestAppContext) {
        let agent = init_test(cx).await;
        let model = agent.model.as_fake();
        let project = agent
            .action_log
            .read_with(cx, |log, _| log.project().clone());
        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi\njkl", cx));

        let mut async_cx = cx.to_async();
        let (apply, mut events) = agent.edit(
            buffer.clone(),
            String::new(),
            &LanguageModelRequest::default(),
            &mut async_cx,
        );
        cx.run_until_parked();

        model.stream_last_completion_response("<old_text>a");
        cx.run_until_parked();
        assert_eq!(drain_events(&mut events), vec![]);
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abc\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            None
        );

        model.stream_last_completion_response("bc</old_text>");
        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
                cx,
                |buffer, _| buffer.anchor_before(Point::new(0, 0))
                    ..buffer.anchor_before(Point::new(0, 3))
            ))]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abc\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 3)))
            })
        );

        model.stream_last_completion_response("<new_text>abX");
        cx.run_until_parked();
        assert_eq!(drain_events(&mut events), [EditAgentOutputEvent::Edited]);
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXc\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 3)))
            })
        );

        model.stream_last_completion_response("cY");
        cx.run_until_parked();
        assert_eq!(drain_events(&mut events), [EditAgentOutputEvent::Edited]);
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXcY\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
            })
        );

        model.stream_last_completion_response("</new_text>");
        model.stream_last_completion_response("<old_text>hall");
        cx.run_until_parked();
        assert_eq!(drain_events(&mut events), vec![]);
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXcY\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
            })
        );

        model.stream_last_completion_response("ucinated old</old_text>");
        model.stream_last_completion_response("<new_text>");
        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::UnresolvedEditRange]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXcY\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
            })
        );

        model.stream_last_completion_response("hallucinated new</new_");
        model.stream_last_completion_response("text>");
        cx.run_until_parked();
        assert_eq!(drain_events(&mut events), vec![]);
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXcY\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
            })
        );

        model.stream_last_completion_response("<old_text>\nghi\nj");
        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
                cx,
                |buffer, _| buffer.anchor_before(Point::new(2, 0))
                    ..buffer.anchor_before(Point::new(2, 3))
            ))]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXcY\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
            })
        );

        model.stream_last_completion_response("kl</old_text>");
        model.stream_last_completion_response("<new_text>");
        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
                cx,
                |buffer, _| buffer.anchor_before(Point::new(2, 0))
                    ..buffer.anchor_before(Point::new(3, 3))
            ))]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXcY\ndef\nghi\njkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(3, 3)))
            })
        );

        model.stream_last_completion_response("GHI</new_text>");
        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::Edited]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXcY\ndef\nGHI"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
            })
        );

        model.end_last_completion_stream();
        apply.await.unwrap();
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "abXcY\ndef\nGHI"
        );
        assert_eq!(drain_events(&mut events), vec![]);
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
            })
        );
    }

    #[gpui::test]
    async fn test_overwrite_events(cx: &mut TestAppContext) {
        let agent = init_test(cx).await;
        let project = agent
            .action_log
            .read_with(cx, |log, _| log.project().clone());
        let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
        let (chunks_tx, chunks_rx) = mpsc::unbounded();
        let (apply, mut events) = agent.overwrite_with_chunks(
            buffer.clone(),
            chunks_rx.map(|chunk: &str| Ok(chunk.to_string())),
            &mut cx.to_async(),
        );

        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::Edited]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            ""
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: language::Anchor::MAX
            })
        );

        chunks_tx.unbounded_send("```\njkl\n").unwrap();
        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::Edited]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "jkl"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: language::Anchor::MAX
            })
        );

        chunks_tx.unbounded_send("mno\n").unwrap();
        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::Edited]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "jkl\nmno"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: language::Anchor::MAX
            })
        );

        chunks_tx.unbounded_send("pqr\n```").unwrap();
        cx.run_until_parked();
        assert_eq!(
            drain_events(&mut events),
            vec![EditAgentOutputEvent::Edited]
        );
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "jkl\nmno\npqr"
        );
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: language::Anchor::MAX
            })
        );

        drop(chunks_tx);
        apply.await.unwrap();
        assert_eq!(
            buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
            "jkl\nmno\npqr"
        );
        assert_eq!(drain_events(&mut events), vec![]);
        assert_eq!(
            project.read_with(cx, |project, _| project.agent_location()),
            Some(AgentLocation {
                buffer: buffer.downgrade(),
                position: language::Anchor::MAX
            })
        );
    }

    #[gpui::test(iterations = 100)]
    async fn test_indent_new_text_chunks(mut rng: StdRng) {
        let chunks = to_random_chunks(&mut rng, "    abc\n  def\n      ghi");
        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
            Ok(EditParserEvent::NewTextChunk {
                chunk: chunk.clone(),
                done: index == chunks.len() - 1,
            })
        }));
        let indented_chunks =
            EditAgent::reindent_new_text_chunks(IndentDelta::Spaces(2), new_text_chunks)
                .collect::<Vec<_>>()
                .await;
        let new_text = indented_chunks
            .into_iter()
            .collect::<Result<String>>()
            .unwrap();
        assert_eq!(new_text, "      abc\n    def\n        ghi");
    }

    #[gpui::test(iterations = 100)]
    async fn test_outdent_new_text_chunks(mut rng: StdRng) {
        let chunks = to_random_chunks(&mut rng, "\t\t\t\tabc\n\t\tdef\n\t\t\t\t\t\tghi");
        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
            Ok(EditParserEvent::NewTextChunk {
                chunk: chunk.clone(),
                done: index == chunks.len() - 1,
            })
        }));
        let indented_chunks =
            EditAgent::reindent_new_text_chunks(IndentDelta::Tabs(-2), new_text_chunks)
                .collect::<Vec<_>>()
                .await;
        let new_text = indented_chunks
            .into_iter()
            .collect::<Result<String>>()
            .unwrap();
        assert_eq!(new_text, "\t\tabc\ndef\n\t\t\t\tghi");
    }

    #[gpui::test(iterations = 100)]
    async fn test_random_indents(mut rng: StdRng) {
        let len = rng.gen_range(1..=100);
        let new_text = util::RandomCharIter::new(&mut rng)
            .with_simple_text()
            .take(len)
            .collect::<String>();
        let new_text = new_text
            .split('\n')
            .map(|line| format!("{}{}", " ".repeat(rng.gen_range(0..=8)), line))
            .collect::<Vec<_>>()
            .join("\n");
        let delta = IndentDelta::Spaces(rng.gen_range(-4..=4));

        let chunks = to_random_chunks(&mut rng, &new_text);
        let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
            Ok(EditParserEvent::NewTextChunk {
                chunk: chunk.clone(),
                done: index == chunks.len() - 1,
            })
        }));
        let reindented_chunks = EditAgent::reindent_new_text_chunks(delta, new_text_chunks)
            .collect::<Vec<_>>()
            .await;
        let actual_reindented_text = reindented_chunks
            .into_iter()
            .collect::<Result<String>>()
            .unwrap();
        let expected_reindented_text = new_text
            .split('\n')
            .map(|line| {
                if let Some(ix) = line.find(|c| c != ' ') {
                    let new_indent = cmp::max(0, ix as isize + delta.len()) as usize;
                    format!("{}{}", " ".repeat(new_indent), &line[ix..])
                } else {
                    line.to_string()
                }
            })
            .collect::<Vec<_>>()
            .join("\n");
        assert_eq!(actual_reindented_text, expected_reindented_text);
    }

    fn to_random_chunks(rng: &mut StdRng, input: &str) -> Vec<String> {
        let chunk_count = rng.gen_range(1..=cmp::min(input.len(), 50));
        let mut chunk_indices = (0..input.len()).choose_multiple(rng, chunk_count);
        chunk_indices.sort();
        chunk_indices.push(input.len());

        let mut chunks = Vec::new();
        let mut last_ix = 0;
        for chunk_ix in chunk_indices {
            chunks.push(input[last_ix..chunk_ix].to_string());
            last_ix = chunk_ix;
        }
        chunks
    }

    fn simulate_llm_output(
        agent: &EditAgent,
        output: &str,
        rng: &mut StdRng,
        cx: &mut TestAppContext,
    ) {
        let executor = cx.executor();
        let chunks = to_random_chunks(rng, output);
        let model = agent.model.clone();
        cx.background_spawn(async move {
            for chunk in chunks {
                executor.simulate_random_delay().await;
                model.as_fake().stream_last_completion_response(chunk);
            }
            model.as_fake().end_last_completion_stream();
        })
        .detach();
    }

    async fn init_test(cx: &mut TestAppContext) -> EditAgent {
        cx.update(settings::init);
        cx.update(Project::init_settings);
        let project = Project::test(FakeFs::new(cx.executor()), [], cx).await;
        let model = Arc::new(FakeLanguageModel::default());
        let action_log = cx.new(|_| ActionLog::new(project.clone()));
        EditAgent::new(model, project, action_log, Templates::new())
    }

    fn drain_events(
        stream: &mut UnboundedReceiver<EditAgentOutputEvent>,
    ) -> Vec<EditAgentOutputEvent> {
        let mut events = Vec::new();
        while let Ok(Some(event)) = stream.try_next() {
            events.push(event);
        }
        events
    }
}