use super::*;
use crate::udiff::apply_diff_to_string;
use client::{RefreshLlmTokenListener, UserStore, test::FakeServer};
use clock::FakeSystemClock;
use clock::ReplicaId;
use cloud_api_types::{
    CreateLlmTokenResponse, LlmToken, Organization, OrganizationConfiguration,
    OrganizationEditPredictionConfiguration, OrganizationId, SubmitEditPredictionSettledBody,
    SubmitEditPredictionSettledResponse,
};
use cloud_llm_client::{
    EditPredictionRejectReason, EditPredictionRejection, RejectEditPredictionsBody,
    predict_edits_v3::{PredictEditsV3Request, PredictEditsV3Response},
};
use db::AppDatabase;
use settings::EditPredictionDataCollectionChoice;

use futures::{
    AsyncReadExt, FutureExt, StreamExt,
    channel::{mpsc, oneshot},
};
use gpui::App;
use gpui::{
    Entity, TestAppContext,
    http_client::{FakeHttpClient, Response},
};
use indoc::indoc;
use language::{
    Anchor, Buffer, Capability, CursorShape, Diagnostic, DiagnosticEntry, DiagnosticSet,
    DiagnosticSeverity, Operation, Point, Selection, SelectionGoal,
};

use lsp::LanguageServerId;
use parking_lot::Mutex;
use pretty_assertions::{assert_eq, assert_matches};
use project::{FakeFs, Project};
use serde_json::json;
use settings::SettingsStore;
use std::{ops::Range, path::Path, sync::Arc, time::Duration};
use util::{
    path,
    test::{TextRangeMarker, marked_text_ranges_by},
};
use uuid::Uuid;
use workspace::{AppState, CollaboratorId, MultiWorkspace};
use zeta_prompt::ZetaPromptInput;

use crate::{
    BufferEditPrediction, EDIT_PREDICTION_SETTLED_QUIESCENCE, EditPredictionId,
    EditPredictionJumpsFeatureFlag, EditPredictionStore, REJECT_REQUEST_DEBOUNCE,
};

#[gpui::test]
async fn test_current_state(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "1.txt": "Hello!\nHow\nBye\n",
            "2.txt": "Hola!\nComo\nAdios\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer1 = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("/root/1.txt"), cx).unwrap();
            project.set_active_path(Some(path.clone()), cx);
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot1 = buffer1.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot1.anchor_before(language::Point::new(1, 3));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_project(&project, cx);
        ep_store.register_buffer(&buffer1, &project, cx);
    });

    // Prediction for current file

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer1.clone(), position, cx)
    });
    let (request, respond_tx) = requests.predict.next().await.unwrap();

    respond_tx
        .send(model_response(
            &request,
            indoc! {r"
                --- a/root/1.txt
                +++ b/root/1.txt
                @@ ... @@
                 Hello!
                -How
                +How are you?
                 Bye
            "},
        ))
        .unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        let prediction = ep_store
            .prediction_at(&buffer1, None, &project, cx)
            .unwrap();
        assert_matches!(prediction, BufferEditPrediction::Local { .. });
    });

    ep_store.update(cx, |ep_store, cx| {
        ep_store.reject_current_prediction(EditPredictionRejectReason::Discarded, &project, cx);
    });

    // Prediction for diagnostic in another file

    let diagnostic = lsp::Diagnostic {
        range: lsp::Range::new(lsp::Position::new(1, 1), lsp::Position::new(1, 5)),
        severity: Some(lsp::DiagnosticSeverity::ERROR),
        message: "Sentence is incomplete".to_string(),
        ..Default::default()
    };

    project.update(cx, |project, cx| {
        project.lsp_store().update(cx, |lsp_store, cx| {
            lsp_store
                .update_diagnostics(
                    LanguageServerId(0),
                    lsp::PublishDiagnosticsParams {
                        uri: lsp::Uri::from_file_path(path!("/root/2.txt")).unwrap(),
                        diagnostics: vec![diagnostic],
                        version: None,
                    },
                    None,
                    language::DiagnosticSourceKind::Pushed,
                    &[],
                    cx,
                )
                .unwrap();
        });
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    respond_tx
        .send(model_response(
            &request,
            indoc! {r#"
                --- a/root/2.txt
                +++ b/root/2.txt
                @@ ... @@
                 Hola!
                -Como
                +Como estas?
                 Adios
            "#},
        ))
        .unwrap();
    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        let prediction = ep_store
            .prediction_at(&buffer1, None, &project, cx)
            .unwrap();
        assert_matches!(
            prediction,
            BufferEditPrediction::Jump { prediction } if prediction.snapshot.file().unwrap().full_path(cx) == Path::new(path!("root/2.txt"))
        );
    });

    let buffer2 = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/2.txt"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    ep_store.update(cx, |ep_store, cx| {
        let prediction = ep_store
            .prediction_at(&buffer2, None, &project, cx)
            .unwrap();
        assert_matches!(prediction, BufferEditPrediction::Local { .. });
    });
}

#[gpui::test]
async fn test_diagnostics_refresh_suppressed_while_following(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);

    cx.update(|cx| {
        cx.update_flags(
            false,
            vec![EditPredictionJumpsFeatureFlag::NAME.to_string()],
        );
    });

    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "1.txt": "Hello!\nHow\nBye\n",
            "2.txt": "Hola!\nComo\nAdios\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let app_state = cx.update(|cx| {
        let app_state = AppState::test(cx);
        AppState::set_global(app_state.clone(), cx);
        app_state
    });

    let multi_workspace =
        cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
    let workspace = multi_workspace
        .read_with(cx, |multi_workspace, _| multi_workspace.workspace().clone())
        .unwrap();
    cx.update(|cx| {
        AppState::set_global(workspace.read(cx).app_state().clone(), cx);
    });
    let _ = app_state;

    let buffer1 = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/1.txt"), cx).unwrap();
            project.set_active_path(Some(path.clone()), cx);
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot1 = buffer1.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot1.anchor_before(language::Point::new(1, 3));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_project(&project, cx);
        ep_store.register_buffer(&buffer1, &project, cx);
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer1.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    respond_tx
        .send(model_response(
            &request,
            indoc! {r"
                --- a/root/1.txt
                +++ b/root/1.txt
                @@ ... @@
                 Hello!
                -How
                +How are you?
                 Bye
            "},
        ))
        .unwrap();
    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        ep_store.reject_current_prediction(EditPredictionRejectReason::Discarded, &project, cx);
    });

    let _ = multi_workspace.update(cx, |multi_workspace, window, cx| {
        multi_workspace.workspace().update(cx, |workspace, cx| {
            workspace.start_following(CollaboratorId::Agent, window, cx);
        });
    });
    cx.run_until_parked();

    let diagnostic = lsp::Diagnostic {
        range: lsp::Range::new(lsp::Position::new(1, 1), lsp::Position::new(1, 5)),
        severity: Some(lsp::DiagnosticSeverity::ERROR),
        message: "Sentence is incomplete".to_string(),
        ..Default::default()
    };

    project.update(cx, |project, cx| {
        project.lsp_store().update(cx, |lsp_store, cx| {
            lsp_store
                .update_diagnostics(
                    LanguageServerId(0),
                    lsp::PublishDiagnosticsParams {
                        uri: lsp::Uri::from_file_path(path!("/root/2.txt")).unwrap(),
                        diagnostics: vec![diagnostic.clone()],
                        version: None,
                    },
                    None,
                    language::DiagnosticSourceKind::Pushed,
                    &[],
                    cx,
                )
                .unwrap();
        });
    });

    cx.run_until_parked();
    assert_no_predict_request_ready(&mut requests.predict);

    let _ = multi_workspace.update(cx, |multi_workspace, window, cx| {
        multi_workspace.workspace().update(cx, |workspace, cx| {
            workspace.unfollow(CollaboratorId::Agent, window, cx);
        });
    });
    cx.run_until_parked();

    project.update(cx, |project, cx| {
        project.lsp_store().update(cx, |lsp_store, cx| {
            lsp_store
                .update_diagnostics(
                    LanguageServerId(0),
                    lsp::PublishDiagnosticsParams {
                        uri: lsp::Uri::from_file_path(path!("/root/2.txt")).unwrap(),
                        diagnostics: vec![diagnostic],
                        version: None,
                    },
                    None,
                    language::DiagnosticSourceKind::Pushed,
                    &[],
                    cx,
                )
                .unwrap();
        });
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    respond_tx
        .send(model_response(
            &request,
            indoc! {r#"
                --- a/root/2.txt
                +++ b/root/2.txt
                @@ ... @@
                 Hola!
                -Como
                +Como estas?
                 Adios
            "#},
        ))
        .unwrap();
    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        let prediction = ep_store
            .prediction_at(&buffer1, None, &project, cx)
            .unwrap();
        assert_matches!(
            prediction,
            BufferEditPrediction::Jump { prediction } if prediction.snapshot.file().unwrap().full_path(cx) == Path::new(path!("root/2.txt"))
        );
    });
}

#[gpui::test]
async fn test_simple_request(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    let prediction_task = ep_store.update(cx, |ep_store, cx| {
        ep_store.request_prediction(&project, &buffer, position, Default::default(), cx)
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();

    // TODO Put back when we have a structured request again
    // assert_eq!(
    //     request.excerpt_path.as_ref(),
    //     Path::new(path!("root/foo.md"))
    // );
    // assert_eq!(
    //     request.cursor_point,
    //     Point {
    //         line: Line(1),
    //         column: 3
    //     }
    // );

    respond_tx
        .send(model_response(
            &request,
            indoc! { r"
                --- a/root/foo.md
                +++ b/root/foo.md
                @@ ... @@
                 Hello!
                -How
                +How are you?
                 Bye
            "},
        ))
        .unwrap();

    let prediction = prediction_task.await.unwrap().unwrap().prediction.unwrap();

    assert_eq!(prediction.edits.len(), 1);
    assert_eq!(
        prediction.edits[0].0.to_point(&snapshot).start,
        language::Point::new(1, 3)
    );
    assert_eq!(prediction.edits[0].1.as_ref(), " are you?");
}

#[gpui::test]
async fn test_request_events(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md": "Hello!\n\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);
    });

    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(7..7, "How")], None, cx);
    });

    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    let prediction_task = ep_store.update(cx, |ep_store, cx| {
        ep_store.request_prediction(&project, &buffer, position, Default::default(), cx)
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();

    let prompt = prompt_from_request(&request);
    assert!(
        prompt.contains(indoc! {"
        --- a/root/foo.md
        +++ b/root/foo.md
        @@ -1,3 +1,3 @@
         Hello!
        -
        +How
         Bye
    "}),
        "{prompt}"
    );

    respond_tx
        .send(model_response(
            &request,
            indoc! {r#"
                --- a/root/foo.md
                +++ b/root/foo.md
                @@ ... @@
                 Hello!
                -How
                +How are you?
                 Bye
        "#},
        ))
        .unwrap();

    let prediction = prediction_task.await.unwrap().unwrap().prediction.unwrap();

    assert_eq!(prediction.edits.len(), 1);
    assert_eq!(prediction.edits[0].1.as_ref(), " are you?");
}

#[gpui::test]
async fn test_edit_history_getter_pause_splits_last_event(cx: &mut TestAppContext) {
    let (ep_store, _requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md": "Hello!\n\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);
    });

    // First burst: insert "How"
    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(7..7, "How")], None, cx);
    });

    // Simulate a pause longer than the grouping threshold (e.g. 500ms).
    cx.executor().advance_clock(LAST_CHANGE_GROUPING_TIME * 2);
    cx.run_until_parked();

    // Second burst: append " are you?" immediately after "How" on the same line.
    //
    // Keeping both bursts on the same line ensures the existing line-span coalescing logic
    // groups them into a single `LastEvent`, allowing the pause-split getter to return two diffs.
    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(10..10, " are you?")], None, cx);
    });

    // A second edit shortly after the first post-pause edit ensures the last edit timestamp is
    // advanced after the pause boundary is recorded, making pause-splitting deterministic.
    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(19..19, "!")], None, cx);
    });

    // With time-based splitting, there are two distinct events.
    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(events.len(), 2);

    let first_total_edit_range = buffer.read_with(cx, |buffer, _| {
        events[0].total_edit_range.to_point(&buffer.snapshot())
    });
    assert_eq!(first_total_edit_range, Point::new(1, 0)..Point::new(1, 3));

    let zeta_prompt::Event::BufferChange { diff, .. } = events[0].event.as_ref();
    assert_eq!(
        diff.as_str(),
        indoc! {"
            @@ -1,3 +1,3 @@
             Hello!
            -
            +How
             Bye
        "}
    );

    let second_total_edit_range = buffer.read_with(cx, |buffer, _| {
        events[1].total_edit_range.to_point(&buffer.snapshot())
    });
    assert_eq!(second_total_edit_range, Point::new(1, 3)..Point::new(1, 13));

    let zeta_prompt::Event::BufferChange { diff, .. } = events[1].event.as_ref();
    assert_eq!(
        diff.as_str(),
        indoc! {"
            @@ -1,3 +1,3 @@
             Hello!
            -How
            +How are you?!
             Bye
        "}
    );
}

#[gpui::test]
async fn test_predicted_edits_are_separated_in_edit_history(cx: &mut TestAppContext) {
    let (ep_store, _requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());

    // Create a file with 30 lines to test line-based coalescing
    let content = (1..=30)
        .map(|i| format!("Line {}\n", i))
        .collect::<String>();
    fs.insert_tree(
        "/root",
        json!({
            "foo.md": content
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);
    });

    // First edit: multi-line edit spanning rows 10-12 (replacing lines 11-13)
    buffer.update(cx, |buffer, cx| {
        let start = Point::new(10, 0).to_offset(buffer);
        let end = Point::new(13, 0).to_offset(buffer);
        buffer.edit(vec![(start..end, "Middle A\nMiddle B\n")], None, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events(&events),
        indoc! {"
            @@ -8,9 +8,8 @@
             Line 8
             Line 9
             Line 10
            -Line 11
            -Line 12
            -Line 13
            +Middle A
            +Middle B
             Line 14
             Line 15
             Line 16
        "},
        "After first edit"
    );

    // Second edit: insert ABOVE the first edit's range (row 5, within 8 lines of row 10)
    // This tests that coalescing considers the START of the existing range
    buffer.update(cx, |buffer, cx| {
        let offset = Point::new(5, 0).to_offset(buffer);
        buffer.edit(vec![(offset..offset, "Above\n")], None, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events(&events),
        indoc! {"
            @@ -3,14 +3,14 @@
             Line 3
             Line 4
             Line 5
            +Above
             Line 6
             Line 7
             Line 8
             Line 9
             Line 10
            -Line 11
            -Line 12
            -Line 13
            +Middle A
            +Middle B
             Line 14
             Line 15
             Line 16
        "},
        "After inserting above (should coalesce)"
    );

    // Third edit: insert BELOW the first edit's range (row 14 in current buffer, within 8 lines of row 12)
    // This tests that coalescing considers the END of the existing range
    buffer.update(cx, |buffer, cx| {
        let offset = Point::new(14, 0).to_offset(buffer);
        buffer.edit(vec![(offset..offset, "Below\n")], None, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events(&events),
        indoc! {"
            @@ -3,15 +3,16 @@
             Line 3
             Line 4
             Line 5
            +Above
             Line 6
             Line 7
             Line 8
             Line 9
             Line 10
            -Line 11
            -Line 12
            -Line 13
            +Middle A
            +Middle B
             Line 14
            +Below
             Line 15
             Line 16
             Line 17
        "},
        "After inserting below (should coalesce)"
    );

    // Fourth edit: insert FAR BELOW (row 25, beyond 8 lines from the current range end ~row 15)
    // This should NOT coalesce - creates a new event
    buffer.update(cx, |buffer, cx| {
        let offset = Point::new(25, 0).to_offset(buffer);
        buffer.edit(vec![(offset..offset, "Far below\n")], None, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events(&events),
        indoc! {"
            @@ -3,15 +3,16 @@
             Line 3
             Line 4
             Line 5
            +Above
             Line 6
             Line 7
             Line 8
             Line 9
             Line 10
            -Line 11
            -Line 12
            -Line 13
            +Middle A
            +Middle B
             Line 14
            +Below
             Line 15
             Line 16
             Line 17

            ---
            @@ -23,6 +23,7 @@
             Line 22
             Line 23
             Line 24
            +Far below
             Line 25
             Line 26
             Line 27
        "},
        "After inserting far below (should NOT coalesce)"
    );
}

fn render_events(events: &[StoredEvent]) -> String {
    events
        .iter()
        .map(|e| {
            let zeta_prompt::Event::BufferChange { diff, .. } = e.event.as_ref();
            diff.as_str()
        })
        .collect::<Vec<_>>()
        .join("\n---\n")
}

fn render_events_with_predicted(events: &[StoredEvent]) -> Vec<String> {
    events
        .iter()
        .map(|e| {
            let zeta_prompt::Event::BufferChange {
                diff, predicted, ..
            } = e.event.as_ref();
            let prefix = if *predicted { "predicted" } else { "manual" };
            format!("{}\n{}", prefix, diff)
        })
        .collect()
}

fn make_collaborator_replica(
    buffer: &Entity<Buffer>,
    cx: &mut TestAppContext,
) -> (Entity<Buffer>, clock::Global) {
    let (state, version) =
        buffer.read_with(cx, |buffer, _cx| (buffer.to_proto(_cx), buffer.version()));
    let collaborator = cx.new(|_cx| {
        Buffer::from_proto(ReplicaId::new(1), Capability::ReadWrite, state, None).unwrap()
    });
    (collaborator, version)
}

async fn apply_collaborator_edit(
    collaborator: &Entity<Buffer>,
    buffer: &Entity<Buffer>,
    since_version: &mut clock::Global,
    edit_range: Range<usize>,
    new_text: &str,
    cx: &mut TestAppContext,
) {
    collaborator.update(cx, |collaborator, cx| {
        collaborator.edit([(edit_range, new_text)], None, cx);
    });

    let serialize_task = collaborator.read_with(cx, |collaborator, cx| {
        collaborator.serialize_ops(Some(since_version.clone()), cx)
    });
    let ops = serialize_task.await;
    *since_version = collaborator.read_with(cx, |collaborator, _cx| collaborator.version());

    buffer.update(cx, |buffer, cx| {
        buffer.apply_ops(
            ops.into_iter()
                .map(|op| language::proto::deserialize_operation(op).unwrap()),
            cx,
        );
    });
}

#[gpui::test]
async fn test_nearby_collaborator_edits_are_kept_in_history(cx: &mut TestAppContext) {
    let (ep_store, _requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.rs": "line 0\nline 1\nline 2\nline 3\nline 4\nline 5\nline 6\nline 7\nline 8\nline 9\nline 10\nline 11\nline 12\nline 13\nline 14\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
            project.set_active_path(Some(path.clone()), cx);
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    let cursor = buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0)));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);
        let _ = ep_store.prediction_at(&buffer, Some(cursor), &project, cx);
    });

    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(0..6, "LOCAL ZERO")], None, cx);
    });

    let (collaborator, mut collaborator_version) = make_collaborator_replica(&buffer, cx);

    let (line_one_start, line_one_len) = collaborator.read_with(cx, |buffer, _cx| {
        (Point::new(1, 0).to_offset(buffer), buffer.line_len(1))
    });

    apply_collaborator_edit(
        &collaborator,
        &buffer,
        &mut collaborator_version,
        line_one_start..line_one_start + line_one_len as usize,
        "REMOTE ONE",
        cx,
    )
    .await;

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });

    assert_eq!(
        render_events_with_predicted(&events),
        vec![indoc! {"
            manual
            @@ -1,5 +1,5 @@
            -line 0
            -line 1
            +LOCAL ZERO
            +REMOTE ONE
             line 2
             line 3
             line 4
        "}]
    );
}

#[gpui::test]
async fn test_distant_collaborator_edits_are_omitted_from_history(cx: &mut TestAppContext) {
    let (ep_store, _requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.rs": (0..1000)
                .map(|i| format!("line {i}\n"))
                .collect::<String>()
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
            project.set_active_path(Some(path.clone()), cx);
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    let cursor = buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0)));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);
        let _ = ep_store.prediction_at(&buffer, Some(cursor), &project, cx);
    });

    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(0..6, "LOCAL ZERO")], None, cx);
    });

    let (collaborator, mut collaborator_version) = make_collaborator_replica(&buffer, cx);

    let far_line_start = buffer.read_with(cx, |buffer, _cx| Point::new(900, 0).to_offset(buffer));

    apply_collaborator_edit(
        &collaborator,
        &buffer,
        &mut collaborator_version,
        far_line_start..far_line_start + 7,
        "REMOTE FAR",
        cx,
    )
    .await;

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });

    assert_eq!(
        render_events_with_predicted(&events),
        vec![indoc! {"
            manual
            @@ -1,4 +1,4 @@
            -line 0
            +LOCAL ZERO
             line 1
             line 2
             line 3
        "}]
    );
}

#[gpui::test]
async fn test_irrelevant_collaborator_edits_in_different_files_are_omitted_from_history(
    cx: &mut TestAppContext,
) {
    let (ep_store, _requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.rs": "line 0\nline 1\nline 2\nline 3\n",
            "bar.rs": "line 0\nline 1\nline 2\nline 3\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let foo_buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
            project.set_active_path(Some(path.clone()), cx);
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let bar_buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/bar.rs"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    let foo_cursor = foo_buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0)));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&foo_buffer, &project, cx);
        ep_store.register_buffer(&bar_buffer, &project, cx);
        let _ = ep_store.prediction_at(&foo_buffer, Some(foo_cursor), &project, cx);
    });

    let (bar_collaborator, mut bar_version) = make_collaborator_replica(&bar_buffer, cx);

    apply_collaborator_edit(
        &bar_collaborator,
        &bar_buffer,
        &mut bar_version,
        0..6,
        "REMOTE BAR",
        cx,
    )
    .await;

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });

    assert!(events.is_empty());
}

#[gpui::test]
async fn test_large_edits_are_omitted_from_history(cx: &mut TestAppContext) {
    let (ep_store, _requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.rs": (0..20)
                .map(|i| format!("line {i}\n"))
                .collect::<String>()
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
            project.set_active_path(Some(path.clone()), cx);
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    let cursor = buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0)));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);
        let _ = ep_store.prediction_at(&buffer, Some(cursor), &project, cx);
    });

    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(0..6, "LOCAL ZERO")], None, cx);
    });

    let (collaborator, mut collaborator_version) = make_collaborator_replica(&buffer, cx);

    let (line_three_start, line_three_len) = collaborator.read_with(cx, |buffer, _cx| {
        (Point::new(3, 0).to_offset(buffer), buffer.line_len(3))
    });
    let large_edit = "X".repeat(EDIT_HISTORY_DIFF_SIZE_LIMIT + 1);

    apply_collaborator_edit(
        &collaborator,
        &buffer,
        &mut collaborator_version,
        line_three_start..line_three_start + line_three_len as usize,
        &large_edit,
        cx,
    )
    .await;

    buffer.update(cx, |buffer, cx| {
        let line_seven_start = Point::new(7, 0).to_offset(buffer);
        let line_seven_end = Point::new(7, 6).to_offset(buffer);
        buffer.edit(
            vec![(line_seven_start..line_seven_end, "LOCAL SEVEN")],
            None,
            cx,
        );
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });

    let rendered_events = render_events_with_predicted(&events);

    assert_eq!(rendered_events.len(), 2);
    assert!(rendered_events[0].contains("+LOCAL ZERO"));
    assert!(!rendered_events[0].contains(&large_edit));
    assert!(rendered_events[1].contains("+LOCAL SEVEN"));
    assert!(!rendered_events[1].contains(&large_edit));
}

#[gpui::test]
async fn test_predicted_flag_coalescing(cx: &mut TestAppContext) {
    let (ep_store, _requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.rs": "line 0\nline 1\nline 2\nline 3\nline 4\nline 5\nline 6\nline 7\nline 8\nline 9\nline 10\nline 11\nline 12\nline 13\nline 14\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);
    });

    // Case 1: Manual edits have `predicted` set to false.
    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(0..6, "LINE ZERO")], None, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });

    assert_eq!(
        render_events_with_predicted(&events),
        vec![indoc! {"
            manual
            @@ -1,4 +1,4 @@
            -line 0
            +LINE ZERO
             line 1
             line 2
             line 3
        "}]
    );

    // Case 2: Multiple successive manual edits near each other are merged into one
    // event with `predicted` set to false.
    buffer.update(cx, |buffer, cx| {
        let offset = Point::new(1, 0).to_offset(buffer);
        let end = Point::new(1, 6).to_offset(buffer);
        buffer.edit(vec![(offset..end, "LINE ONE")], None, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events_with_predicted(&events),
        vec![indoc! {"
            manual
            @@ -1,5 +1,5 @@
            -line 0
            -line 1
            +LINE ZERO
            +LINE ONE
             line 2
             line 3
             line 4
        "}]
    );

    // Case 3: Accepted predictions have `predicted` set to true.
    // Case 5: A manual edit that follows a predicted edit is not merged with the
    // predicted edit, even if it is nearby.
    ep_store.update(cx, |ep_store, cx| {
        buffer.update(cx, |buffer, cx| {
            let offset = Point::new(2, 0).to_offset(buffer);
            let end = Point::new(2, 6).to_offset(buffer);
            buffer.edit(vec![(offset..end, "LINE TWO")], None, cx);
        });
        ep_store.report_changes_for_buffer(&buffer, &project, true, true, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events_with_predicted(&events),
        vec![
            indoc! {"
                manual
                @@ -1,5 +1,5 @@
                -line 0
                -line 1
                +LINE ZERO
                +LINE ONE
                 line 2
                 line 3
                 line 4
            "},
            indoc! {"
                predicted
                @@ -1,6 +1,6 @@
                 LINE ZERO
                 LINE ONE
                -line 2
                +LINE TWO
                 line 3
                 line 4
                 line 5
            "}
        ]
    );

    // Case 4: Multiple successive accepted predictions near each other are merged
    // into one event with `predicted` set to true.
    ep_store.update(cx, |ep_store, cx| {
        buffer.update(cx, |buffer, cx| {
            let offset = Point::new(3, 0).to_offset(buffer);
            let end = Point::new(3, 6).to_offset(buffer);
            buffer.edit(vec![(offset..end, "LINE THREE")], None, cx);
        });
        ep_store.report_changes_for_buffer(&buffer, &project, true, true, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events_with_predicted(&events),
        vec![
            indoc! {"
                manual
                @@ -1,5 +1,5 @@
                -line 0
                -line 1
                +LINE ZERO
                +LINE ONE
                 line 2
                 line 3
                 line 4
            "},
            indoc! {"
                predicted
                @@ -1,7 +1,7 @@
                 LINE ZERO
                 LINE ONE
                -line 2
                -line 3
                +LINE TWO
                +LINE THREE
                 line 4
                 line 5
                 line 6
            "}
        ]
    );

    // Case 5 (continued): A manual edit that follows a predicted edit is not merged
    // with the predicted edit, even if it is nearby.
    buffer.update(cx, |buffer, cx| {
        let offset = Point::new(4, 0).to_offset(buffer);
        let end = Point::new(4, 6).to_offset(buffer);
        buffer.edit(vec![(offset..end, "LINE FOUR")], None, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events_with_predicted(&events),
        vec![
            indoc! {"
                manual
                @@ -1,5 +1,5 @@
                -line 0
                -line 1
                +LINE ZERO
                +LINE ONE
                 line 2
                 line 3
                 line 4
            "},
            indoc! {"
                predicted
                @@ -1,7 +1,7 @@
                 LINE ZERO
                 LINE ONE
                -line 2
                -line 3
                +LINE TWO
                +LINE THREE
                 line 4
                 line 5
                 line 6
            "},
            indoc! {"
                manual
                @@ -2,7 +2,7 @@
                 LINE ONE
                 LINE TWO
                 LINE THREE
                -line 4
                +LINE FOUR
                 line 5
                 line 6
                 line 7
            "}
        ]
    );

    // Case 6: If we then perform a manual edit at a *different* location (more than
    // 8 lines away), then the edits at the prior location can be merged with each
    // other, even if some are predicted and some are not. `predicted` means all
    // constituent edits were predicted.
    buffer.update(cx, |buffer, cx| {
        let offset = Point::new(14, 0).to_offset(buffer);
        let end = Point::new(14, 7).to_offset(buffer);
        buffer.edit(vec![(offset..end, "LINE FOURTEEN")], None, cx);
    });

    let events = ep_store.update(cx, |ep_store, cx| {
        ep_store.edit_history_for_project(&project, cx)
    });
    assert_eq!(
        render_events_with_predicted(&events),
        vec![
            indoc! {"
                manual
                @@ -1,8 +1,8 @@
                -line 0
                -line 1
                -line 2
                -line 3
                -line 4
                +LINE ZERO
                +LINE ONE
                +LINE TWO
                +LINE THREE
                +LINE FOUR
                 line 5
                 line 6
                 line 7
            "},
            indoc! {"
                manual
                @@ -12,4 +12,4 @@
                 line 11
                 line 12
                 line 13
                -line 14
                +LINE FOURTEEN
            "}
        ]
    );
}

#[gpui::test]
async fn test_empty_prediction(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    let mut response = model_response(&request, "");
    response.model_version = Some("zeta2:test-empty".to_string());
    let id = response.request_id.clone();
    respond_tx.send(response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        assert!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .is_none()
        );
    });

    // prediction is reported as rejected
    let (reject_request, _) = requests.reject.next().await.unwrap();

    assert_eq!(
        &reject_request.rejections,
        &[EditPredictionRejection {
            request_id: id,
            reason: EditPredictionRejectReason::Empty,
            was_shown: false,
            model_version: Some("zeta2:test-empty".to_string()),
            e2e_latency_ms: Some(0),
        }]
    );
}

#[gpui::test]
async fn test_interpolated_empty(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();

    buffer.update(cx, |buffer, cx| {
        buffer.set_text("Hello!\nHow are you?\nBye", cx);
    });

    let mut response = model_response(&request, SIMPLE_DIFF);
    response.model_version = Some("zeta2:test-interpolated-empty".to_string());
    let id = response.request_id.clone();
    respond_tx.send(response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        assert!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .is_none()
        );
    });

    // prediction is reported as rejected
    let (reject_request, _) = requests.reject.next().await.unwrap();

    assert_eq!(
        &reject_request.rejections,
        &[EditPredictionRejection {
            request_id: id,
            reason: EditPredictionRejectReason::InterpolatedEmpty,
            was_shown: false,
            model_version: Some("zeta2:test-interpolated-empty".to_string()),
            e2e_latency_ms: Some(0),
        }]
    );
}

const SIMPLE_DIFF: &str = indoc! { r"
    --- a/root/foo.md
    +++ b/root/foo.md
    @@ ... @@
     Hello!
    -How
    +How are you?
     Bye
"};

#[gpui::test]
async fn test_replace_current(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    let first_response = model_response(&request, SIMPLE_DIFF);
    let first_id = first_response.request_id.clone();
    respond_tx.send(first_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            first_id
        );
    });

    // a second request is triggered
    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    let second_response = model_response(&request, SIMPLE_DIFF);
    let second_id = second_response.request_id.clone();
    respond_tx.send(second_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        // second replaces first
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            second_id
        );
    });

    // first is reported as replaced
    let (reject_request, _) = requests.reject.next().await.unwrap();

    assert_eq!(
        &reject_request.rejections,
        &[EditPredictionRejection {
            request_id: first_id,
            reason: EditPredictionRejectReason::Replaced,
            was_shown: false,
            model_version: None,
            e2e_latency_ms: Some(0),
        }]
    );
}

#[gpui::test]
async fn test_current_preferred(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    let first_response = model_response(&request, SIMPLE_DIFF);
    let first_id = first_response.request_id.clone();
    respond_tx.send(first_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            first_id
        );
    });

    // a second request is triggered
    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    // worse than current prediction
    let mut second_response = model_response(
        &request,
        indoc! { r"
            --- a/root/foo.md
            +++ b/root/foo.md
            @@ ... @@
             Hello!
            -How
            +How are
             Bye
        "},
    );
    second_response.model_version = Some("zeta2:test-current-preferred".to_string());
    let second_id = second_response.request_id.clone();
    respond_tx.send(second_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        // first is preferred over second
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            first_id
        );
    });

    // second is reported as rejected
    let (reject_request, _) = requests.reject.next().await.unwrap();

    assert_eq!(
        &reject_request.rejections,
        &[EditPredictionRejection {
            request_id: second_id,
            reason: EditPredictionRejectReason::CurrentPreferred,
            was_shown: false,
            model_version: Some("zeta2:test-current-preferred".to_string()),
            e2e_latency_ms: Some(0),
        }]
    );
}

#[gpui::test]
async fn test_cancel_earlier_pending_requests(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    // start two refresh tasks
    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request1, respond_first) = requests.predict.next().await.unwrap();

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_second) = requests.predict.next().await.unwrap();

    // wait for throttle
    cx.run_until_parked();

    // second responds first
    let second_response = model_response(&request, SIMPLE_DIFF);
    let second_id = second_response.request_id.clone();
    respond_second.send(second_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        // current prediction is second
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            second_id
        );
    });

    let mut first_response = model_response(&request1, SIMPLE_DIFF);
    first_response.model_version = Some("zeta2:test-canceled".to_string());
    let first_id = first_response.request_id.clone();
    respond_first.send(first_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        // current prediction is still second, since first was cancelled
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            second_id
        );
    });

    // first is reported as rejected
    let (reject_request, _) = requests.reject.next().await.unwrap();

    cx.run_until_parked();

    assert_eq!(
        &reject_request.rejections,
        &[EditPredictionRejection {
            request_id: first_id,
            reason: EditPredictionRejectReason::Canceled,
            was_shown: false,
            model_version: Some("zeta2:test-canceled".to_string()),
            e2e_latency_ms: None,
        }]
    );
}

#[gpui::test]
async fn test_cancel_second_on_third_request(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    // start two refresh tasks
    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request1, respond_first) = requests.predict.next().await.unwrap();

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request2, respond_second) = requests.predict.next().await.unwrap();

    // wait for throttle, so requests are sent
    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        // start a third request
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);

        // 2 are pending, so 2nd is cancelled
        assert_eq!(
            ep_store
                .get_or_init_project(&project, cx)
                .cancelled_predictions
                .iter()
                .copied()
                .collect::<Vec<_>>(),
            [1]
        );
    });

    // wait for throttle
    cx.run_until_parked();

    let (request3, respond_third) = requests.predict.next().await.unwrap();

    let first_response = model_response(&request1, SIMPLE_DIFF);
    let first_id = first_response.request_id.clone();
    respond_first.send(first_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        // current prediction is first
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            first_id
        );
    });

    let mut cancelled_response = model_response(&request2, SIMPLE_DIFF);
    cancelled_response.model_version = Some("zeta2:test-canceled-second".to_string());
    let cancelled_id = cancelled_response.request_id.clone();
    respond_second.send(cancelled_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        // current prediction is still first, since second was cancelled
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            first_id
        );
    });

    let third_response = model_response(&request3, SIMPLE_DIFF);
    let third_response_id = third_response.request_id.clone();
    respond_third.send(third_response).unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        // third completes and replaces first
        assert_eq!(
            ep_store
                .prediction_at(&buffer, None, &project, cx)
                .unwrap()
                .id
                .0,
            third_response_id
        );
    });

    // second is reported as rejected
    let (reject_request, _) = requests.reject.next().await.unwrap();

    cx.run_until_parked();

    assert_eq!(
        &reject_request.rejections,
        &[
            EditPredictionRejection {
                request_id: cancelled_id,
                reason: EditPredictionRejectReason::Canceled,
                was_shown: false,
                model_version: Some("zeta2:test-canceled-second".to_string()),
                e2e_latency_ms: None,
            },
            EditPredictionRejection {
                request_id: first_id,
                reason: EditPredictionRejectReason::Replaced,
                was_shown: false,
                model_version: None,
                // 2 throttle waits (for 2nd and 3rd requests) elapsed
                // between this request's start and response.
                e2e_latency_ms: Some(2 * EditPredictionStore::THROTTLE_TIMEOUT.as_millis()),
            }
        ]
    );
}

#[gpui::test]
async fn test_jump_and_edit_throttles_are_independent(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);

    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n",
            "bar.md": "Hola!\nComo\nAdios\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.set_active_path(Some(path.clone()), cx);
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_project(&project, cx);
        ep_store.register_buffer(&buffer, &project, cx);
    });

    // First edit request - no prior edit, so not throttled.
    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });
    let (_edit_request, edit_response_tx) = requests.predict.next().await.unwrap();
    edit_response_tx.send(empty_response()).unwrap();
    cx.run_until_parked();

    let diagnostic = lsp::Diagnostic {
        range: lsp::Range::new(lsp::Position::new(1, 1), lsp::Position::new(1, 5)),
        severity: Some(lsp::DiagnosticSeverity::ERROR),
        message: "Sentence is incomplete".to_string(),
        ..Default::default()
    };

    // First jump request triggered by diagnostic event on buffer - no prior jump, so not throttled (independent from edit).
    project.update(cx, |project, cx| {
        project.lsp_store().update(cx, |lsp_store, cx| {
            lsp_store
                .update_diagnostics(
                    LanguageServerId(0),
                    lsp::PublishDiagnosticsParams {
                        uri: lsp::Uri::from_file_path(path!("/root/bar.md")).unwrap(),
                        diagnostics: vec![diagnostic],
                        version: None,
                    },
                    None,
                    language::DiagnosticSourceKind::Pushed,
                    &[],
                    cx,
                )
                .unwrap();
        });
    });
    let (_jump_request, jump_response_tx) = requests.predict.next().await.unwrap();
    jump_response_tx.send(empty_response()).unwrap();
    cx.run_until_parked();

    // Second edit request - should be throttled by the first edit.
    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });
    assert_no_predict_request_ready(&mut requests.predict);

    // Second jump request - should be throttled by the first jump.
    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_diagnostics(
            project.clone(),
            DiagnosticSearchScope::Global,
            cx,
        );
    });
    assert_no_predict_request_ready(&mut requests.predict);

    // Wait for both throttles to expire.
    cx.background_executor
        .advance_clock(EditPredictionStore::THROTTLE_TIMEOUT);
    cx.background_executor.run_until_parked();
    cx.run_until_parked();

    // Both requests should now go through.
    let (_request_1, response_tx_1) = requests.predict.next().await.unwrap();
    response_tx_1.send(empty_response()).unwrap();
    cx.run_until_parked();

    let (_request_2, response_tx_2) = requests.predict.next().await.unwrap();
    response_tx_2.send(empty_response()).unwrap();
    cx.run_until_parked();
}

#[gpui::test]
async fn test_same_frame_duplicate_requests_deduplicated(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md":  "Hello!\nHow\nBye\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();
    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(1, 3));

    // Enqueue two refresh calls in the same synchronous frame (no yielding).
    // Both `cx.spawn` tasks are created before either executes, so they both
    // capture the same `proceed_count_at_enqueue`. Only the first task should
    // pass the deduplication gate; the second should be skipped.
    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    // Let both spawned tasks run to completion (including any throttle waits).
    cx.run_until_parked();

    // Exactly one prediction request should have been sent.
    let (request, respond_tx) = requests.predict.next().await.unwrap();
    respond_tx
        .send(model_response(&request, SIMPLE_DIFF))
        .unwrap();
    cx.run_until_parked();

    // No second request should be pending.
    assert_no_predict_request_ready(&mut requests.predict);
}

#[gpui::test]
async fn test_rejections_flushing(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);

    ep_store.update(cx, |ep_store, cx| {
        ep_store.reject_prediction(
            EditPredictionId("test-1".into()),
            EditPredictionRejectReason::Discarded,
            false,
            None,
            None,
            cx,
        );
        ep_store.reject_prediction(
            EditPredictionId("test-2".into()),
            EditPredictionRejectReason::Canceled,
            true,
            None,
            None,
            cx,
        );
    });

    cx.executor().advance_clock(REJECT_REQUEST_DEBOUNCE);
    cx.run_until_parked();

    let (reject_request, respond_tx) = requests.reject.next().await.unwrap();
    respond_tx.send(()).unwrap();

    // batched
    assert_eq!(reject_request.rejections.len(), 2);
    assert_eq!(
        reject_request.rejections[0],
        EditPredictionRejection {
            request_id: "test-1".to_string(),
            reason: EditPredictionRejectReason::Discarded,
            was_shown: false,
            model_version: None,
            e2e_latency_ms: None
        }
    );
    assert_eq!(
        reject_request.rejections[1],
        EditPredictionRejection {
            request_id: "test-2".to_string(),
            reason: EditPredictionRejectReason::Canceled,
            was_shown: true,
            model_version: None,
            e2e_latency_ms: None
        }
    );

    // Reaching batch size limit sends without debounce
    ep_store.update(cx, |ep_store, cx| {
        for i in 0..70 {
            ep_store.reject_prediction(
                EditPredictionId(format!("batch-{}", i).into()),
                EditPredictionRejectReason::Discarded,
                false,
                None,
                None,
                cx,
            );
        }
    });

    // First MAX/2 items are sent immediately
    cx.run_until_parked();
    let (reject_request, respond_tx) = requests.reject.next().await.unwrap();
    respond_tx.send(()).unwrap();

    assert_eq!(reject_request.rejections.len(), 50);
    assert_eq!(reject_request.rejections[0].request_id, "batch-0");
    assert_eq!(reject_request.rejections[49].request_id, "batch-49");

    // Remaining items are debounced with the next batch
    cx.executor().advance_clock(Duration::from_secs(15));
    cx.run_until_parked();

    let (reject_request, respond_tx) = requests.reject.next().await.unwrap();
    respond_tx.send(()).unwrap();

    assert_eq!(reject_request.rejections.len(), 20);
    assert_eq!(reject_request.rejections[0].request_id, "batch-50");
    assert_eq!(reject_request.rejections[19].request_id, "batch-69");

    // Request failure
    ep_store.update(cx, |ep_store, cx| {
        ep_store.reject_prediction(
            EditPredictionId("retry-1".into()),
            EditPredictionRejectReason::Discarded,
            false,
            None,
            None,
            cx,
        );
    });

    cx.executor().advance_clock(REJECT_REQUEST_DEBOUNCE);
    cx.run_until_parked();

    let (reject_request, _respond_tx) = requests.reject.next().await.unwrap();
    assert_eq!(reject_request.rejections.len(), 1);
    assert_eq!(reject_request.rejections[0].request_id, "retry-1");
    // Simulate failure
    drop(_respond_tx);

    // Add another rejection
    ep_store.update(cx, |ep_store, cx| {
        ep_store.reject_prediction(
            EditPredictionId("retry-2".into()),
            EditPredictionRejectReason::Discarded,
            false,
            None,
            None,
            cx,
        );
    });

    cx.executor().advance_clock(REJECT_REQUEST_DEBOUNCE);
    cx.run_until_parked();

    // Retry should include both the failed item and the new one
    let (reject_request, respond_tx) = requests.reject.next().await.unwrap();
    respond_tx.send(()).unwrap();

    assert_eq!(reject_request.rejections.len(), 2);
    assert_eq!(reject_request.rejections[0].request_id, "retry-1");
    assert_eq!(reject_request.rejections[1].request_id, "retry-2");
}

#[gpui::test]
fn test_active_buffer_diagnostics_fetching(cx: &mut TestAppContext) {
    let diagnostic_marker: TextRangeMarker = ('«', '»').into();
    let search_range_marker: TextRangeMarker = ('[', ']').into();

    let (text, mut ranges) = marked_text_ranges_by(
        indoc! {r#"
            fn alpha() {
                let «first_value» = 1;
            }

            [fn beta() {
                let «second_value» = 2;
                let third_value = second_value + missing_symbol;
            }ˇ]

            fn gamma() {
                let «fourth_value» = missing_other_symbol;
            }
        "#},
        vec![diagnostic_marker.clone(), search_range_marker.clone()],
    );

    let diagnostic_ranges = ranges.remove(&diagnostic_marker).unwrap_or_default();
    let search_ranges = ranges.remove(&search_range_marker).unwrap_or_default();

    let buffer = cx.new(|cx| Buffer::local(&text, cx));

    buffer.update(cx, |buffer, cx| {
        let snapshot = buffer.snapshot();
        let diagnostics = DiagnosticSet::new(
            diagnostic_ranges
                .iter()
                .enumerate()
                .map(|(index, range)| DiagnosticEntry {
                    range: snapshot.offset_to_point_utf16(range.start)
                        ..snapshot.offset_to_point_utf16(range.end),
                    diagnostic: Diagnostic {
                        severity: match index {
                            0 => DiagnosticSeverity::WARNING,
                            1 => DiagnosticSeverity::ERROR,
                            _ => DiagnosticSeverity::HINT,
                        },
                        message: match index {
                            0 => "first warning".to_string(),
                            1 => "second error".to_string(),
                            _ => "third hint".to_string(),
                        },
                        group_id: index + 1,
                        is_primary: true,
                        source_kind: language::DiagnosticSourceKind::Pushed,
                        ..Diagnostic::default()
                    },
                }),
            &snapshot,
        );
        buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx);
    });

    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let search_range = snapshot.offset_to_point(search_ranges[0].start)
        ..snapshot.offset_to_point(search_ranges[0].end);

    let active_buffer_diagnostics = zeta::active_buffer_diagnostics(&snapshot, search_range, 5, 0);

    assert_eq!(
        active_buffer_diagnostics,
        vec![zeta_prompt::ActiveBufferDiagnostic {
            severity: Some(1),
            message: "second error".to_string(),
            snippet: "    let second_value = 2;".to_string(),
            snippet_buffer_row_range: 5..5,
            diagnostic_range_in_snippet: 8..20,
        }]
    );

    let active_buffer_diagnostics =
        zeta::active_buffer_diagnostics(&snapshot, Point::new(0, 0)..snapshot.max_point(), 5, 100);
    assert_eq!(
        active_buffer_diagnostics,
        vec![
            zeta_prompt::ActiveBufferDiagnostic {
                severity: Some(1),
                message: "second error".to_string(),
                snippet: String::new(),
                snippet_buffer_row_range: 5..5,
                diagnostic_range_in_snippet: 0..0,
            },
            zeta_prompt::ActiveBufferDiagnostic {
                severity: Some(2),
                message: "first warning".to_string(),
                snippet: String::new(),
                snippet_buffer_row_range: 1..1,
                diagnostic_range_in_snippet: 0..0,
            },
            zeta_prompt::ActiveBufferDiagnostic {
                severity: Some(4),
                message: "third hint".to_string(),
                snippet: String::new(),
                snippet_buffer_row_range: 10..10,
                diagnostic_range_in_snippet: 0..0,
            },
        ]
    );

    let buffer = cx.new(|cx| {
        Buffer::local(
            indoc! {"
                one
                two
                three
                four
                five
            "},
            cx,
        )
    });

    buffer.update(cx, |buffer, cx| {
        let snapshot = buffer.snapshot();
        let diagnostics = DiagnosticSet::new(
            vec![
                DiagnosticEntry {
                    range: text::PointUtf16::new(0, 0)..text::PointUtf16::new(0, 3),
                    diagnostic: Diagnostic {
                        severity: DiagnosticSeverity::ERROR,
                        message: "row zero".to_string(),
                        group_id: 1,
                        is_primary: true,
                        source_kind: language::DiagnosticSourceKind::Pushed,
                        ..Diagnostic::default()
                    },
                },
                DiagnosticEntry {
                    range: text::PointUtf16::new(2, 0)..text::PointUtf16::new(2, 5),
                    diagnostic: Diagnostic {
                        severity: DiagnosticSeverity::WARNING,
                        message: "row two".to_string(),
                        group_id: 2,
                        is_primary: true,
                        source_kind: language::DiagnosticSourceKind::Pushed,
                        ..Diagnostic::default()
                    },
                },
                DiagnosticEntry {
                    range: text::PointUtf16::new(4, 0)..text::PointUtf16::new(4, 4),
                    diagnostic: Diagnostic {
                        severity: DiagnosticSeverity::INFORMATION,
                        message: "row four".to_string(),
                        group_id: 3,
                        is_primary: true,
                        source_kind: language::DiagnosticSourceKind::Pushed,
                        ..Diagnostic::default()
                    },
                },
            ],
            &snapshot,
        );
        buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx);
    });

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

    let active_buffer_diagnostics =
        zeta::active_buffer_diagnostics(&snapshot, Point::new(2, 0)..Point::new(4, 0), 3, 0);

    assert_eq!(
        active_buffer_diagnostics
            .iter()
            .map(|diagnostic| (
                diagnostic.severity,
                diagnostic.message.clone(),
                diagnostic.snippet.clone(),
                diagnostic.snippet_buffer_row_range.clone(),
                diagnostic.diagnostic_range_in_snippet.clone(),
            ))
            .collect::<Vec<_>>(),
        vec![
            (
                Some(2),
                "row two".to_string(),
                "three".to_string(),
                2..2,
                0..5,
            ),
            (
                Some(3),
                "row four".to_string(),
                "five".to_string(),
                4..4,
                0..4,
            ),
        ]
    );
}

#[gpui::test]
fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) {
    let text = (0..25)
        .map(|row| format!("line {row}\n"))
        .collect::<String>();
    let buffer = cx.new(|cx| Buffer::local(&text, cx));

    buffer.update(cx, |buffer, cx| {
        let snapshot = buffer.snapshot();
        let diagnostics = DiagnosticSet::new(
            (0..25)
                .map(|row| DiagnosticEntry {
                    range: text::PointUtf16::new(row, 0)..text::PointUtf16::new(row, 4),
                    diagnostic: Diagnostic {
                        severity: DiagnosticSeverity::ERROR,
                        message: format!("row {row}"),
                        group_id: row as usize,
                        is_primary: true,
                        source_kind: language::DiagnosticSourceKind::Pushed,
                        ..Diagnostic::default()
                    },
                })
                .collect::<Vec<_>>(),
            &snapshot,
        );
        buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx);
    });

    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let active_buffer_diagnostics =
        zeta::active_buffer_diagnostics(&snapshot, Point::new(0, 0)..Point::new(25, 0), 12, 0);

    assert_eq!(active_buffer_diagnostics.len(), 20);
    assert!(
        active_buffer_diagnostics
            .iter()
            .any(|diagnostic| diagnostic.message == "row 12")
    );
    assert!(
        active_buffer_diagnostics
            .iter()
            .all(|diagnostic| diagnostic.message != "row 0" && diagnostic.message != "row 24")
    );

    let text = (0..300)
        .map(|row| format!("line {row} has some diagnostic context\n"))
        .collect::<String>();
    let buffer = cx.new(|cx| Buffer::local(&text, cx));

    buffer.update(cx, |buffer, cx| {
        let snapshot = buffer.snapshot();
        let diagnostics = DiagnosticSet::new(
            vec![DiagnosticEntry {
                range: text::PointUtf16::new(150, 0)..text::PointUtf16::new(150, 4),
                diagnostic: Diagnostic {
                    severity: DiagnosticSeverity::ERROR,
                    message: "long snippet".to_string(),
                    group_id: 1,
                    is_primary: true,
                    source_kind: language::DiagnosticSourceKind::Pushed,
                    ..Diagnostic::default()
                },
            }],
            &snapshot,
        );
        buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx);
    });

    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let active_buffer_diagnostics = zeta::active_buffer_diagnostics(
        &snapshot,
        Point::new(100, 0)..Point::new(200, 0),
        150,
        2000,
    );

    assert_eq!(active_buffer_diagnostics.len(), 1);
    assert!(active_buffer_diagnostics[0].snippet.len() <= 512 * 3 + 2);
    assert!(active_buffer_diagnostics[0].snippet.len() < text.len());
}

// Generate a model response that would apply the given diff to the active file.
fn model_response(request: &PredictEditsV3Request, diff_to_apply: &str) -> PredictEditsV3Response {
    let editable_range =
        zeta_prompt::excerpt_range_for_format(Default::default(), &request.input.excerpt_ranges).1;
    let excerpt = request.input.cursor_excerpt[editable_range.clone()].to_string();
    let new_excerpt = apply_diff_to_string(diff_to_apply, &excerpt).unwrap();

    PredictEditsV3Response {
        request_id: Uuid::new_v4().to_string(),
        editable_range,
        output: new_excerpt,
        cursor_offset: None,
        model_version: None,
    }
}

fn empty_response() -> PredictEditsV3Response {
    PredictEditsV3Response {
        request_id: Uuid::new_v4().to_string(),
        editable_range: 0..0,
        output: String::new(),
        cursor_offset: None,
        model_version: None,
    }
}

fn prompt_from_request(request: &PredictEditsV3Request) -> String {
    zeta_prompt::format_zeta_prompt(&request.input, zeta_prompt::ZetaFormat::default())
        .expect("default zeta prompt formatting should succeed in edit prediction tests")
}

fn assert_no_predict_request_ready(
    requests: &mut mpsc::UnboundedReceiver<(
        PredictEditsV3Request,
        oneshot::Sender<PredictEditsV3Response>,
    )>,
) {
    if requests.next().now_or_never().flatten().is_some() {
        panic!("Unexpected prediction request while throttled.");
    }
}

struct RequestChannels {
    predict: mpsc::UnboundedReceiver<(
        PredictEditsV3Request,
        oneshot::Sender<PredictEditsV3Response>,
    )>,
    reject: mpsc::UnboundedReceiver<(RejectEditPredictionsBody, oneshot::Sender<()>)>,
    settled: mpsc::UnboundedReceiver<SubmitEditPredictionSettledBody>,
}

fn init_test_with_fake_client(
    cx: &mut TestAppContext,
) -> (Entity<EditPredictionStore>, RequestChannels) {
    init_test_with_fake_client_and_legacy_data_collection(cx, None)
}

fn init_test_with_fake_client_and_legacy_data_collection(
    cx: &mut TestAppContext,
    legacy_data_collection_choice: Option<&str>,
) -> (Entity<EditPredictionStore>, RequestChannels) {
    cx.update(move |cx| {
        cx.set_global(AppDatabase::test_new());
        let settings_store = SettingsStore::test(cx);
        cx.set_global(settings_store);
        zlog::init_test();

        if let Some(legacy_data_collection_choice) = legacy_data_collection_choice {
            KeyValueStore::global(cx)
                .write_kvp(
                    ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
                    legacy_data_collection_choice.to_string(),
                )
                .now_or_never()
                .expect("legacy data collection write should complete immediately")
                .expect("legacy data collection write should succeed");
        }

        let (predict_req_tx, predict_req_rx) = mpsc::unbounded();
        let (reject_req_tx, reject_req_rx) = mpsc::unbounded();
        let (settled_req_tx, settled_req_rx) = mpsc::unbounded();

        let http_client = FakeHttpClient::create({
            move |req| {
                let uri = req.uri().path().to_string();
                let content_encoding = req
                    .headers()
                    .get("Content-Encoding")
                    .and_then(|value| value.to_str().ok())
                    .map(str::to_owned);
                let mut body = req.into_body();
                let predict_req_tx = predict_req_tx.clone();
                let reject_req_tx = reject_req_tx.clone();
                let settled_req_tx = settled_req_tx.clone();
                async move {
                    let resp = match uri.as_str() {
                        "/client/llm_tokens" => serde_json::to_string(&json!({
                            "token": "test"
                        }))
                        .unwrap(),
                        "/predict_edits/v3" => {
                            let mut buf = Vec::new();
                            body.read_to_end(&mut buf).await.ok();
                            let decompressed = zstd::decode_all(&buf[..]).unwrap();
                            let req = serde_json::from_slice(&decompressed).unwrap();

                            let (res_tx, res_rx) = oneshot::channel();
                            predict_req_tx.unbounded_send((req, res_tx)).unwrap();
                            serde_json::to_string(&res_rx.await?).unwrap()
                        }
                        "/predict_edits/reject" => {
                            let mut buf = Vec::new();
                            body.read_to_end(&mut buf).await.ok();
                            let req = serde_json::from_slice(&buf).unwrap();

                            let (res_tx, res_rx) = oneshot::channel();
                            reject_req_tx.unbounded_send((req, res_tx)).unwrap();
                            serde_json::to_string(&res_rx.await?).unwrap()
                        }
                        "/predict_edits/settled" => {
                            let mut buf = Vec::new();
                            body.read_to_end(&mut buf).await.ok();
                            let body = if content_encoding.as_deref() == Some("zstd") {
                                zstd::decode_all(&buf[..]).unwrap()
                            } else {
                                buf
                            };
                            let req = serde_json::from_slice(&body).unwrap();
                            settled_req_tx.unbounded_send(req).unwrap();
                            serde_json::to_string(&SubmitEditPredictionSettledResponse {}).unwrap()
                        }
                        _ => {
                            panic!("Unexpected path: {}", uri)
                        }
                    };

                    Ok(Response::builder().body(resp.into()).unwrap())
                }
            }
        });

        let client = client::Client::new(Arc::new(FakeSystemClock::new()), http_client, cx);
        client.cloud_client().set_credentials(1, "test".into());

        let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
        language_model::init(cx);
        RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
        let ep_store = EditPredictionStore::global(&client, &user_store, cx);

        (
            ep_store,
            RequestChannels {
                predict: predict_req_rx,
                reject: reject_req_rx,
                settled: settled_req_rx,
            },
        )
    })
}

#[gpui::test]
async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
    let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
    let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
        to_completion_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
    });

    let edit_preview = cx
        .read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
        .await;

    let prediction = EditPrediction {
        edits,
        cursor_position: None,
        editable_range: None,
        edit_preview,
        buffer: buffer.clone(),
        snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
        id: EditPredictionId("the-id".into()),
        inputs: ZetaPromptInput {
            events: Default::default(),
            related_files: Default::default(),
            active_buffer_diagnostics: vec![],
            cursor_path: Path::new("").into(),
            cursor_excerpt: "".into(),
            cursor_offset_in_excerpt: 0,
            excerpt_start_row: None,
            excerpt_ranges: Default::default(),
            syntax_ranges: None,
            in_open_source_repo: false,
            can_collect_data: false,
            repo_url: None,
        },
        model_version: None,
    };

    cx.update(|cx| {
        assert_eq!(
            from_completion_edits(
                &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                &buffer,
                cx
            ),
            vec![(2..5, "REM".into()), (9..11, "".into())]
        );

        buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
        assert_eq!(
            from_completion_edits(
                &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                &buffer,
                cx
            ),
            vec![(2..2, "REM".into()), (6..8, "".into())]
        );

        buffer.update(cx, |buffer, cx| buffer.undo(cx));
        assert_eq!(
            from_completion_edits(
                &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                &buffer,
                cx
            ),
            vec![(2..5, "REM".into()), (9..11, "".into())]
        );

        buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
        assert_eq!(
            from_completion_edits(
                &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                &buffer,
                cx
            ),
            vec![(3..3, "EM".into()), (7..9, "".into())]
        );

        buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
        assert_eq!(
            from_completion_edits(
                &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                &buffer,
                cx
            ),
            vec![(4..4, "M".into()), (8..10, "".into())]
        );

        buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
        assert_eq!(
            from_completion_edits(
                &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                &buffer,
                cx
            ),
            vec![(9..11, "".into())]
        );

        buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
        assert_eq!(
            from_completion_edits(
                &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                &buffer,
                cx
            ),
            vec![(4..4, "M".into()), (8..10, "".into())]
        );

        buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
        assert_eq!(
            from_completion_edits(
                &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
                &buffer,
                cx
            ),
            vec![(4..4, "M".into())]
        );

        buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
        assert_eq!(prediction.interpolate(&buffer.read(cx).snapshot()), None);
    })
}

#[gpui::test]
async fn test_clean_up_diff(cx: &mut TestAppContext) {
    init_test(cx);

    assert_eq!(
        apply_edit_prediction(
            indoc! {"
                    fn main() {
                        let word_1 = \"lorem\";
                        let range = word.len()..word.len();
                    }
                "},
            indoc! {"
                    fn main() {
                        let word_1 = \"lorem\";
                        let range = word_1.len()..word_1.len();
                    }
                "},
            cx,
        )
        .await,
        indoc! {"
                fn main() {
                    let word_1 = \"lorem\";
                    let range = word_1.len()..word_1.len();
                }
            "},
    );

    assert_eq!(
        apply_edit_prediction(
            indoc! {"
                    fn main() {
                        let story = \"the quick\"
                    }
                "},
            indoc! {"
                    fn main() {
                        let story = \"the quick brown fox jumps over the lazy dog\";
                    }
                "},
            cx,
        )
        .await,
        indoc! {"
                fn main() {
                    let story = \"the quick brown fox jumps over the lazy dog\";
                }
            "},
    );
}

#[gpui::test]
async fn test_edit_prediction_end_of_buffer(cx: &mut TestAppContext) {
    init_test(cx);

    let buffer_content = "lorem\n";
    let completion_response = "lorem\nipsum\n";

    assert_eq!(
        apply_edit_prediction(buffer_content, completion_response, cx).await,
        "lorem\nipsum\n"
    );
}

#[gpui::test]
async fn test_edit_prediction_no_spurious_trailing_newline(cx: &mut TestAppContext) {
    // Test that zeta2's newline normalization logic doesn't insert spurious newlines.
    // When the buffer ends without a trailing newline, but the model returns output
    // with a trailing newline, zeta2 should normalize both sides before diffing
    // so no spurious newline is inserted.
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());

    // Single line buffer with no trailing newline
    fs.insert_tree(
        "/root",
        json!({
            "foo.txt": "hello"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project
                .find_project_path(path!("root/foo.txt"), cx)
                .unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(0, 5));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();

    // Model returns output WITH a trailing newline, even though the buffer doesn't have one.
    // Zeta2 should normalize both sides before diffing, so no spurious newline is inserted.
    let excerpt_length = request.input.cursor_excerpt.len();
    let response = PredictEditsV3Response {
        request_id: Uuid::new_v4().to_string(),
        output: "hello world\n".to_string(),
        editable_range: 0..excerpt_length,
        model_version: None,
        cursor_offset: None,
    };
    respond_tx.send(response).unwrap();

    cx.run_until_parked();

    // The prediction should insert " world" without adding a newline
    ep_store.update(cx, |ep_store, cx| {
        let prediction = ep_store
            .prediction_at(&buffer, None, &project, cx)
            .expect("should have prediction");
        let edits: Vec<_> = prediction
            .edits
            .iter()
            .map(|(range, text)| {
                let snapshot = buffer.read(cx).snapshot();
                (range.to_offset(&snapshot), text.clone())
            })
            .collect();
        assert_eq!(edits, vec![(5..5, " world".into())]);
    });
}

#[gpui::test]
async fn test_v3_prediction_strips_cursor_marker_from_edit_text(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());

    fs.insert_tree(
        "/root",
        json!({
            "foo.txt": "hello"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project
                .find_project_path(path!("root/foo.txt"), cx)
                .unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let position = snapshot.anchor_before(language::Point::new(0, 5));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx);
    });

    let (request, respond_tx) = requests.predict.next().await.unwrap();
    let excerpt_length = request.input.cursor_excerpt.len();
    respond_tx
        .send(PredictEditsV3Response {
            request_id: Uuid::new_v4().to_string(),
            output: "hello world".to_string(),
            editable_range: 0..excerpt_length,
            model_version: None,
            cursor_offset: Some(5),
        })
        .unwrap();

    cx.run_until_parked();

    ep_store.update(cx, |ep_store, cx| {
        let prediction = ep_store
            .prediction_at(&buffer, None, &project, cx)
            .expect("should have prediction");
        let snapshot = buffer.read(cx).snapshot();
        let edits: Vec<_> = prediction
            .edits
            .iter()
            .map(|(range, text)| (range.to_offset(&snapshot), text.clone()))
            .collect();

        assert_eq!(edits, vec![(5..5, " world".into())]);
    });
}

fn init_test(cx: &mut TestAppContext) {
    cx.update(|cx| {
        cx.set_global(AppDatabase::test_new());
        let settings_store = SettingsStore::test(cx);
        cx.set_global(settings_store);
    });
}

async fn apply_edit_prediction(
    buffer_content: &str,
    completion_response: &str,
    cx: &mut TestAppContext,
) -> String {
    let fs = project::FakeFs::new(cx.executor());
    let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
    let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
    let (ep_store, response) = make_test_ep_store(&project, cx).await;
    *response.lock() = completion_response.to_string();
    let edit_prediction = run_edit_prediction(&buffer, &project, &ep_store, cx).await;
    buffer.update(cx, |buffer, cx| {
        buffer.edit(edit_prediction.edits.iter().cloned(), None, cx)
    });
    buffer.read_with(cx, |buffer, _| buffer.text())
}

async fn run_edit_prediction(
    buffer: &Entity<Buffer>,
    project: &Entity<Project>,
    ep_store: &Entity<EditPredictionStore>,
    cx: &mut TestAppContext,
) -> EditPrediction {
    let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(buffer, &project, cx)
    });
    cx.background_executor.run_until_parked();
    let prediction_task = ep_store.update(cx, |ep_store, cx| {
        ep_store.request_prediction(&project, buffer, cursor, Default::default(), cx)
    });
    prediction_task.await.unwrap().unwrap().prediction.unwrap()
}

async fn make_test_ep_store(
    project: &Entity<Project>,
    cx: &mut TestAppContext,
) -> (Entity<EditPredictionStore>, Arc<Mutex<String>>) {
    let default_response = "hello world\n".to_string();
    let completion_response: Arc<Mutex<String>> = Arc::new(Mutex::new(default_response));
    let http_client = FakeHttpClient::create({
        let completion_response = completion_response.clone();
        let mut next_request_id = 0;
        move |req| {
            let completion_response = completion_response.clone();
            let method = req.method().clone();
            let uri = req.uri().path().to_string();
            let mut body = req.into_body();
            async move {
                match (method, uri.as_str()) {
                    (Method::POST, "/client/llm_tokens") => Ok(http_client::Response::builder()
                        .status(200)
                        .body(
                            serde_json::to_string(&CreateLlmTokenResponse {
                                token: LlmToken("the-llm-token".to_string()),
                            })
                            .unwrap()
                            .into(),
                        )
                        .unwrap()),
                    (Method::POST, "/predict_edits/v3") => {
                        let mut buf = Vec::new();
                        body.read_to_end(&mut buf).await.ok();
                        let decompressed = zstd::decode_all(&buf[..]).unwrap();
                        let req: PredictEditsV3Request =
                            serde_json::from_slice(&decompressed).unwrap();

                        next_request_id += 1;
                        Ok(http_client::Response::builder()
                            .status(200)
                            .body(
                                serde_json::to_string(&PredictEditsV3Response {
                                    request_id: format!("request-{next_request_id}"),
                                    editable_range: 0..req.input.cursor_excerpt.len(),
                                    output: completion_response.lock().clone(),
                                    model_version: None,
                                    cursor_offset: None,
                                })
                                .unwrap()
                                .into(),
                            )
                            .unwrap())
                    }
                    _ => Ok(http_client::Response::builder()
                        .status(404)
                        .body("Not Found".to_string().into())
                        .unwrap()),
                }
            }
        }
    });

    let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
    let user_store = cx.update(|cx| cx.new(|cx| client::UserStore::new(client.clone(), cx)));
    cx.update(|cx| {
        RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
    });
    let _server = FakeServer::for_client(42, &client, cx).await;

    let ep_store = cx.new(|cx| {
        let mut ep_store = EditPredictionStore::new(client, project.read(cx).user_store(), cx);
        ep_store.set_edit_prediction_model(EditPredictionModel::Zeta);

        let worktrees = project.read(cx).worktrees(cx).collect::<Vec<_>>();
        for worktree in worktrees {
            let worktree_id = worktree.read(cx).id();
            ep_store
                .get_or_init_project(project, cx)
                .license_detection_watchers
                .entry(worktree_id)
                .or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(&worktree, cx)));
        }

        ep_store
    });

    (ep_store, completion_response)
}

fn to_completion_edits(
    iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
    buffer: &Entity<Buffer>,
    cx: &App,
) -> Vec<(Range<Anchor>, Arc<str>)> {
    let buffer = buffer.read(cx);
    iterator
        .into_iter()
        .map(|(range, text)| {
            (
                buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
                text,
            )
        })
        .collect()
}

fn from_completion_edits(
    editor_edits: &[(Range<Anchor>, Arc<str>)],
    buffer: &Entity<Buffer>,
    cx: &App,
) -> Vec<(Range<usize>, Arc<str>)> {
    let buffer = buffer.read(cx);
    editor_edits
        .iter()
        .map(|(range, text)| {
            (
                range.start.to_offset(buffer)..range.end.to_offset(buffer),
                text.clone(),
            )
        })
        .collect()
}

#[gpui::test]
async fn test_unauthenticated_without_custom_url_blocks_prediction_impl(cx: &mut TestAppContext) {
    init_test(cx);

    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/project",
        serde_json::json!({
            "main.rs": "fn main() {\n    \n}\n"
        }),
    )
    .await;

    let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;

    let request_count = Arc::new(std::sync::atomic::AtomicUsize::default());
    let http_client = FakeHttpClient::create({
        let request_count = request_count.clone();
        move |_req| {
            request_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            async move {
                Ok(gpui::http_client::Response::builder()
                    .status(401)
                    .body("Unauthorized".into())
                    .unwrap())
            }
        }
    });

    let client =
        cx.update(|cx| client::Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
    let user_store = cx.update(|cx| cx.new(|cx| client::UserStore::new(client.clone(), cx)));
    cx.update(|cx| {
        RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
    });

    let ep_store = cx.new(|cx| EditPredictionStore::new(client, project.read(cx).user_store(), cx));

    let buffer = project
        .update(cx, |project, cx| {
            let path = project
                .find_project_path(path!("/project/main.rs"), cx)
                .unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 4)));
    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx)
    });
    cx.background_executor.run_until_parked();

    let completion_task = ep_store.update(cx, |ep_store, cx| {
        ep_store.set_edit_prediction_model(EditPredictionModel::Zeta);
        ep_store.request_prediction(&project, &buffer, cursor, Default::default(), cx)
    });

    assert!(completion_task.await.unwrap().is_none());
    assert_eq!(request_count.load(std::sync::atomic::Ordering::SeqCst), 0);
}

#[gpui::test]
async fn test_diagnostic_jump_excludes_collaborator_regions(cx: &mut TestAppContext) {
    fn set_collaborator_cursor(buffer: &Entity<Buffer>, row: u32, cx: &mut TestAppContext) {
        let collab_replica = clock::ReplicaId::new(10);
        let anchor = buffer.read_with(cx, |buffer, _| {
            buffer.snapshot().anchor_before(Point::new(row, 0))
        });
        let selections: Arc<[Selection<Anchor>]> = Arc::new([Selection {
            id: 1,
            start: anchor,
            end: anchor,
            reversed: false,
            goal: SelectionGoal::None,
        }]);
        buffer.update(cx, |buffer, cx| {
            buffer.apply_ops(
                [Operation::UpdateSelections {
                    selections,
                    lamport_timestamp: clock::Lamport {
                        replica_id: collab_replica,
                        value: 1,
                    },
                    line_mode: false,
                    cursor_shape: CursorShape::Bar,
                }],
                cx,
            );
        });
    }

    fn publish_diagnostics(
        uri_path: &'static str,
        rows: &[u32],
        project: &Entity<Project>,
        cx: &mut TestAppContext,
    ) {
        let diagnostics: Vec<_> = rows
            .iter()
            .map(|&row| lsp::Diagnostic {
                range: lsp::Range::new(lsp::Position::new(row, 0), lsp::Position::new(row, 5)),
                severity: Some(lsp::DiagnosticSeverity::ERROR),
                message: format!("error at row {row}"),
                ..Default::default()
            })
            .collect();
        project.update(cx, |project, cx| {
            project.lsp_store().update(cx, |lsp_store, cx| {
                lsp_store
                    .update_diagnostics(
                        LanguageServerId(0),
                        lsp::PublishDiagnosticsParams {
                            uri: lsp::Uri::from_file_path(uri_path).expect("invalid uri"),
                            diagnostics,
                            version: None,
                        },
                        None,
                        language::DiagnosticSourceKind::Pushed,
                        &[],
                        cx,
                    )
                    .expect("failed to update diagnostics");
            });
        });
    }

    init_test(cx);

    let mut lines = String::new();
    for i in 0..60 {
        lines.push_str(&format!("line {i}\n"));
    }

    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "active.txt": lines,
            "collab_file.txt": "error here\nsecond line\n",
            "free_file.txt": "another error\nsecond line\n",
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let active_buffer = project
        .update(cx, |project, cx| {
            let path = project
                .find_project_path(path!("/root/active.txt"), cx)
                .expect("active.txt not found");
            project.set_active_path(Some(path.clone()), cx);
            project.open_buffer(path, cx)
        })
        .await
        .expect("failed to open active buffer");

    set_collaborator_cursor(&active_buffer, 5, cx);

    publish_diagnostics(path!("/root/active.txt"), &[3, 25, 50], &project, cx);

    cx.run_until_parked();

    let cursor_point = Point::new(25, 0);
    let empty_search_range: Range<Point> = Default::default();

    let snapshot = active_buffer.read_with(cx, |buffer, _| buffer.snapshot());
    let result = EditPredictionStore::next_diagnostic_location(
        active_buffer.clone(),
        &snapshot,
        empty_search_range.clone(),
        cursor_point,
        &project,
        &mut cx.to_async(),
    )
    .await
    .expect("next_diagnostic_location failed");

    let (result_buffer, result_anchor) = result.expect("expected a diagnostic location");
    assert_eq!(result_buffer.entity_id(), active_buffer.entity_id());
    let result_row = result_buffer.read_with(cx, |buffer, _| {
        result_anchor.to_point(&buffer.snapshot()).row
    });
    assert_ne!(
        result_row, 3,
        "row 3 is near collaborator (row 5) but far from local cursor (row 25), should be excluded"
    );
    assert!(
        result_row == 25 || result_row == 50,
        "expected row 25 or 50, got {result_row}"
    );

    let snapshot_near = active_buffer.read_with(cx, |buffer, _| buffer.snapshot());
    let near_cursor_point = Point::new(4, 0);
    let result_near = EditPredictionStore::next_diagnostic_location(
        active_buffer.clone(),
        &snapshot_near,
        empty_search_range.clone(),
        near_cursor_point,
        &project,
        &mut cx.to_async(),
    )
    .await
    .expect("next_diagnostic_location failed");

    let (_, near_anchor) = result_near.expect("expected a diagnostic location when both are near");
    let near_row =
        active_buffer.read_with(cx, |buffer, _| near_anchor.to_point(&buffer.snapshot()).row);
    assert_eq!(
        near_row, 3,
        "row 3 should be included when local cursor (row 4) is also near the collaborator"
    );

    let snapshot_far = active_buffer.read_with(cx, |buffer, _| buffer.snapshot());
    let far_cursor_point = Point::new(50, 0);
    let result_far = EditPredictionStore::next_diagnostic_location(
        active_buffer.clone(),
        &snapshot_far,
        empty_search_range.clone(),
        far_cursor_point,
        &project,
        &mut cx.to_async(),
    )
    .await
    .expect("next_diagnostic_location failed");

    let (_, far_anchor) = result_far.expect("expected a diagnostic location");
    let far_row =
        active_buffer.read_with(cx, |buffer, _| far_anchor.to_point(&buffer.snapshot()).row);
    assert_eq!(
        far_row, 50,
        "row 50 is near local cursor (row 50) and far from collaborator, should be picked"
    );

    publish_diagnostics(path!("/root/collab_file.txt"), &[0], &project, cx);
    publish_diagnostics(path!("/root/free_file.txt"), &[0], &project, cx);
    cx.run_until_parked();

    let collab_buffer = project
        .update(cx, |project, cx| {
            let path = project
                .find_project_path(path!("/root/collab_file.txt"), cx)
                .expect("collab_file.txt not found");
            project.open_buffer(path, cx)
        })
        .await
        .expect("failed to open collab buffer");

    set_collaborator_cursor(&collab_buffer, 0, cx);
    cx.run_until_parked();

    let no_same_file_search_range = Point::new(0, 0)..Point::new(59, 0);
    let snapshot_cross = active_buffer.read_with(cx, |buffer, _| buffer.snapshot());
    let result_cross = EditPredictionStore::next_diagnostic_location(
        active_buffer.clone(),
        &snapshot_cross,
        no_same_file_search_range,
        Point::new(0, 0),
        &project,
        &mut cx.to_async(),
    )
    .await
    .expect("cross-file next_diagnostic_location failed");

    let (cross_buffer, _) = result_cross.expect("expected a cross-file diagnostic location");
    let cross_path = cross_buffer.read_with(cx, |buffer, cx| {
        buffer
            .file()
            .expect("buffer should have a file")
            .full_path(cx)
    });
    assert_eq!(
        cross_path,
        Path::new(path!("root/free_file.txt")),
        "should skip collab_file.txt (has collaborator) and pick free_file.txt"
    );
}

#[gpui::test]
async fn test_edit_prediction_settled(cx: &mut TestAppContext) {
    let (ep_store, _requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());

    // Buffer with two clearly separated regions:
    //   Region A = lines 0-9   (offsets 0..50)
    //   Region B = lines 20-29 (offsets 105..155)
    // A big gap in between so edits in one region never overlap the other.
    let mut content = String::new();
    for i in 0..30 {
        content.push_str(&format!("line {i:02}\n"));
    }

    fs.insert_tree(
        "/root",
        json!({
            "foo.md": content.clone()
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;

    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    type SettledEventRecord = (EditPredictionId, String);
    let settled_events: Arc<Mutex<Vec<SettledEventRecord>>> = Arc::new(Mutex::new(Vec::new()));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);

        let settled_events = settled_events.clone();
        ep_store.settled_event_callback = Some(Box::new(move |id, text| {
            settled_events.lock().push((id, text));
        }));
    });

    // --- Phase 1: edit in region A and enqueue prediction A ---

    buffer.update(cx, |buffer, cx| {
        // Edit at the start of line 0.
        buffer.edit(vec![(0..0, "ADDED ")], None, cx);
    });
    cx.run_until_parked();

    let snapshot_a = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let empty_edits: Arc<[(Range<Anchor>, Arc<str>)]> = Vec::new().into();
    let edit_preview_a = buffer
        .read_with(cx, |buffer, cx| {
            buffer.preview_edits(empty_edits.clone(), cx)
        })
        .await;

    // Region A: first 10 lines of the buffer.
    let editable_region_a = 0..snapshot_a.point_to_offset(Point::new(10, 0));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.enqueue_settled_prediction(
            EditPredictionId("prediction-a".into()),
            &project,
            &buffer,
            &snapshot_a,
            editable_region_a.clone(),
            &edit_preview_a,
            None,
            None,
            Duration::from_secs(0),
            cx,
        );
    });

    // --- Phase 2: repeatedly edit in region A to keep it unsettled ---

    // Let the worker process the channel message before we start advancing.
    cx.run_until_parked();

    for region_a_edit_offset in (5..).take(3) {
        // Edit inside region A (not at the boundary) so `last_edit_at` is
        // updated before the worker's next wake.
        buffer.update(cx, |buffer, cx| {
            buffer.edit(
                vec![(region_a_edit_offset..region_a_edit_offset, "x")],
                None,
                cx,
            );
        });
        cx.run_until_parked();

        cx.executor()
            .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE / 2);
        cx.run_until_parked();
        assert!(
            settled_events.lock().is_empty(),
            "no settled events should fire while region A is still being edited"
        );
    }

    // Still nothing settled.
    assert!(settled_events.lock().is_empty());

    // --- Phase 3: edit in distinct region B, enqueue prediction B ---
    // Advance a small amount so B's quiescence window starts later than A's,
    // but not so much that A settles (A's last edit was at the start of
    // iteration 3, and it needs a full Q to settle).
    cx.executor()
        .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE / 4);
    cx.run_until_parked();
    assert!(settled_events.lock().is_empty());

    let snapshot_b = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let line_20_offset = snapshot_b.point_to_offset(Point::new(20, 0));

    buffer.update(cx, |buffer, cx| {
        buffer.edit(vec![(line_20_offset..line_20_offset, "NEW ")], None, cx);
    });
    cx.run_until_parked();

    let snapshot_b2 = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let edit_preview_b = buffer
        .read_with(cx, |buffer, cx| buffer.preview_edits(empty_edits, cx))
        .await;
    let editable_region_b = line_20_offset..snapshot_b2.point_to_offset(Point::new(25, 0));

    ep_store.update(cx, |ep_store, cx| {
        ep_store.enqueue_settled_prediction(
            EditPredictionId("prediction-b".into()),
            &project,
            &buffer,
            &snapshot_b2,
            editable_region_b.clone(),
            &edit_preview_b,
            None,
            None,
            Duration::from_secs(0),
            cx,
        );
    });

    cx.run_until_parked();
    assert!(
        settled_events.lock().is_empty(),
        "neither prediction should have settled yet"
    );

    // --- Phase 4: let enough time pass for region A to settle ---
    // A's last edit was at T_a (during the last loop iteration). The worker is
    // sleeping until T_a + Q. We advance just enough to reach that wake time
    // (Q/4 since we already advanced Q/4 in phase 3 on top of the loop's
    // 3*Q/2). At that point A has been quiet for Q and settles, but B was
    // enqueued only Q/4 ago and stays pending.
    cx.executor()
        .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE / 4);
    cx.run_until_parked();

    {
        let events = settled_events.lock().clone();
        assert_eq!(
            events.len(),
            1,
            "prediction and capture_sample for A should have settled, got: {events:?}"
        );
        assert_eq!(events[0].0, EditPredictionId("prediction-a".into()));
    }

    // --- Phase 5: let more time pass for region B to settle ---
    // B's last edit was Q/4 before A settled. The worker rescheduled to
    // B's last_edit_at + Q, which is 3Q/4 from now.
    cx.executor()
        .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE * 3 / 4);
    cx.run_until_parked();

    {
        let events = settled_events.lock().clone();
        assert_eq!(
            events.len(),
            2,
            "both prediction and capture_sample settled events should be emitted for each request, got: {events:?}"
        );
        assert_eq!(events[1].0, EditPredictionId("prediction-b".into()));
    }
}

#[gpui::test]
async fn test_edit_prediction_settled_omits_body_when_data_collection_is_disabled(
    cx: &mut TestAppContext,
) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
    let fs = FakeFs::new(cx.executor());
    fs.insert_tree(
        "/root",
        json!({
            "foo.md": "sensitive source\n"
        }),
    )
    .await;
    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
    let buffer = project
        .update(cx, |project, cx| {
            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
            project.open_buffer(path, cx)
        })
        .await
        .unwrap();

    ep_store.update(cx, |ep_store, cx| {
        ep_store.register_buffer(&buffer, &project, cx);
    });

    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
    let edits: Arc<[(Range<Anchor>, Arc<str>)]> =
        cx.update(|cx| to_completion_edits([(0..9, "replacement".into())], &buffer, cx).into());
    let edit_preview = buffer
        .read_with(cx, |buffer, cx| buffer.preview_edits(edits, cx))
        .await;

    ep_store.update(cx, |ep_store, cx| {
        ep_store.enqueue_settled_prediction(
            EditPredictionId("prediction-private".into()),
            &project,
            &buffer,
            &snapshot,
            0..snapshot.len(),
            &edit_preview,
            Some(ExampleSpec {
                name: "test example".to_string(),
                repository_url: "https://example.com/repo".to_string(),
                revision: "rev".to_string(),
                tags: Vec::new(),
                reasoning: None,
                uncommitted_diff: String::new(),
                recently_opened_files: Vec::new(),
                recently_viewed_files: Vec::new(),
                uncommitted_diff_contains_edit_history: false,
                cursor_path: Path::new("foo.md").into(),
                cursor_position: "0".to_string(),
                edit_history: "sensitive edit history".to_string(),
                expected_patches: vec!["sensitive patch".to_string()],
                rejected_patch: None,
                telemetry: None,
                human_feedback: Vec::new(),
                rating: None,
            }),
            Some("test-model".to_string()),
            Duration::from_millis(42),
            cx,
        );
    });

    cx.run_until_parked();
    cx.executor()
        .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
    cx.run_until_parked();

    let settled_request = requests
        .settled
        .next()
        .await
        .expect("settled request should be sent");
    assert!(!settled_request.can_collect_data);
    assert_eq!(settled_request.settled_editable_region, None);
    assert_eq!(settled_request.example, None);
}

#[gpui::test]
fn test_buffer_path_with_id_fallback_for_untitled_buffers(cx: &mut TestAppContext) {
    let buffer_1 = cx.new(|cx| Buffer::local("one", cx));
    let buffer_2 = cx.new(|cx| Buffer::local("two", cx));

    let snapshot_1 = buffer_1.read_with(cx, |buffer, _| buffer.text_snapshot());
    let snapshot_2 = buffer_2.read_with(cx, |buffer, _| buffer.text_snapshot());

    let path_1 = cx.read(|cx| buffer_path_with_id_fallback(None, &snapshot_1, cx));
    let path_2 = cx.read(|cx| buffer_path_with_id_fallback(None, &snapshot_2, cx));

    assert_eq!(
        path_1.as_ref(),
        Path::new(&format!("untitled-{}", snapshot_1.remote_id()))
    );
    assert_eq!(
        path_2.as_ref(),
        Path::new(&format!("untitled-{}", snapshot_2.remote_id()))
    );
    assert_ne!(path_1.as_ref(), path_2.as_ref());
}

#[gpui::test]
async fn test_data_collection_disabled_by_default(cx: &mut TestAppContext) {
    let (ep_store, _channels) = init_test_with_fake_client(cx);

    cx.update(|cx| {
        assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
    });
}

#[gpui::test]
async fn test_data_collection_enabled_via_legacy_kv_store(cx: &mut TestAppContext) {
    let (ep_store, _channels) =
        init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));

    cx.update(|cx| {
        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
    });
}

#[gpui::test]
async fn test_data_collection_default_uses_cached_legacy_value(cx: &mut TestAppContext) {
    let (ep_store, _channels) =
        init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));

    cx.update(|cx| {
        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
    });

    cx.update(|cx| KeyValueStore::global(cx))
        .delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
        .await
        .unwrap();

    cx.update(|cx| {
        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
    });
}

#[gpui::test]
async fn test_data_collection_setting_overrides_kv_store(cx: &mut TestAppContext) {
    let (ep_store, _channels) =
        init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));

    // An explicit false in settings.json wins over the KV store.
    cx.update_global::<SettingsStore, _>(|settings, cx| {
        settings.update_user_settings(cx, |content| {
            content
                .project
                .all_languages
                .edit_predictions
                .get_or_insert_default()
                .allow_data_collection = Some(EditPredictionDataCollectionChoice::No);
        });
    });

    cx.update(|cx| {
        assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
    });
}

#[gpui::test]
async fn test_data_collection_enabled_via_setting(cx: &mut TestAppContext) {
    let (ep_store, _channels) = init_test_with_fake_client(cx);

    cx.update_global::<SettingsStore, _>(|settings, cx| {
        settings.update_user_settings(cx, |content| {
            content
                .project
                .all_languages
                .edit_predictions
                .get_or_insert_default()
                .allow_data_collection = Some(EditPredictionDataCollectionChoice::Yes);
        });
    });

    cx.update(|cx| {
        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
    });
}

#[gpui::test]
async fn test_data_collection_always_enabled_for_staff(cx: &mut TestAppContext) {
    let (ep_store, _channels) = init_test_with_fake_client(cx);

    cx.update(|cx| {
        cx.set_staff(true);
        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
    });
}

#[gpui::test]
async fn test_data_collection_disabled_by_organization_configuration(cx: &mut TestAppContext) {
    let (ep_store, _channels) = init_test_with_fake_client(cx);

    cx.update_global::<SettingsStore, _>(|settings, cx| {
        settings.update_user_settings(cx, |content| {
            content
                .project
                .all_languages
                .edit_predictions
                .get_or_insert_default()
                .allow_data_collection = Some(EditPredictionDataCollectionChoice::Yes);
        });
    });

    let user_store = cx.update(|cx| ep_store.read(cx).user_store.clone());
    cx.update(|cx| {
        user_store.update(cx, |user_store, cx| {
            user_store.set_current_organization_configuration_for_test(
                Arc::new(Organization {
                    id: OrganizationId("org-1".into()),
                    name: "Org 1".into(),
                    is_personal: false,
                }),
                OrganizationConfiguration {
                    is_zed_model_provider_enabled: true,
                    is_agent_thread_feedback_enabled: true,
                    is_collaboration_enabled: true,
                    edit_prediction: OrganizationEditPredictionConfiguration {
                        is_enabled: true,
                        is_feedback_enabled: false,
                    },
                },
                cx,
            );
        });

        assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
    });
}

// When a user had data collection enabled via the legacy KV store (with no explicit
// setting in settings.json), toggle_data_collection must read the *resolved* state
// (true) and write Some(false).
#[gpui::test]
async fn test_toggle_data_collection_from_kv_enabled_state(cx: &mut TestAppContext) {
    let (ep_store, _channels) =
        init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));

    cx.update(|cx| {
        assert!(
            ep_store.read(cx).is_data_collection_enabled(cx),
            "data collection should be enabled via KV store before toggle"
        );
    });

    // Simulate what toggle_data_collection does: capture the resolved current
    // state, then write its inverse.
    let is_currently_enabled = cx.update(|cx| ep_store.read(cx).is_data_collection_enabled(cx));
    cx.update_global::<SettingsStore, _>(|settings, cx| {
        settings.update_user_settings(cx, |content| {
            content
                .project
                .all_languages
                .edit_predictions
                .get_or_insert_default()
                .allow_data_collection = Some(if is_currently_enabled {
                EditPredictionDataCollectionChoice::No
            } else {
                EditPredictionDataCollectionChoice::Yes
            });
        });
    });

    cx.update(|cx| {
        assert!(
            !ep_store.read(cx).is_data_collection_enabled(cx),
            "data collection should be disabled after toggling off from KV-enabled state"
        );
    });
}

#[gpui::test]
async fn test_upsell_shown_by_default(cx: &mut TestAppContext) {
    init_test(cx);
    let kvp = cx.update(|cx| KeyValueStore::global(cx));
    kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
        .await
        .ok();
    kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.ok();

    cx.update(|cx| assert!(should_show_upsell_modal(cx)));
}

#[gpui::test]
async fn test_upsell_dismissed_when_data_collection_choice_in_kv_store(cx: &mut TestAppContext) {
    init_test(cx);

    // Any value for the data collection key means the old upsell was already
    // shown, regardless of whether data collection was accepted or declined.
    for value in &["true", "false"] {
        cx.update(|cx| KeyValueStore::global(cx))
            .write_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into(), value.to_string())
            .await
            .unwrap();

        cx.update(|cx| {
            assert!(
                !should_show_upsell_modal(cx),
                "upsell should be suppressed when data collection choice is '{value}'"
            );
        });
    }

    cx.update(|cx| KeyValueStore::global(cx))
        .delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
        .await
        .unwrap();
}

#[gpui::test]
async fn test_upsell_dismissed_when_dismissed_key_set(cx: &mut TestAppContext) {
    init_test(cx);
    let kvp = cx.update(|cx| KeyValueStore::global(cx));
    kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
        .await
        .ok();
    kvp.write_kvp(ZedPredictUpsell::KEY.into(), "1".into())
        .await
        .unwrap();

    cx.update(|cx| assert!(!should_show_upsell_modal(cx)));

    kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.unwrap();
}

#[gpui::test]
async fn test_upsell_dismissed_via_dismissable_api(cx: &mut TestAppContext) {
    init_test(cx);
    let kvp = cx.update(|cx| KeyValueStore::global(cx));
    kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
        .await
        .ok();
    kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.ok();

    cx.update(|cx| {
        assert!(should_show_upsell_modal(cx));
        ZedPredictUpsell::set_dismissed(true, cx);
    });
    cx.run_until_parked();

    cx.update(|cx| assert!(!should_show_upsell_modal(cx)));

    kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.unwrap();
}

#[ctor::ctor(unsafe)]
fn init_logger() {
    zlog::init_test();
}
