use crate::{    EditPredictionExampleCaptureFeatureFlag, StoredEvent,    cursor_excerpt::editable_and_context_ranges_for_cursor_position, example_spec::ExampleSpec,};use anyhow::Result;use buffer_diff::BufferDiffSnapshot;use collections::HashMap;use feature_flags::FeatureFlagAppExt as _;use gpui::{App, Entity, Task};use language::{Buffer, ToPoint as _};use project::{Project, WorktreeId};use std::{collections::hash_map, fmt::Write as _, path::Path, sync::Arc};use text::BufferSnapshot as TextBufferSnapshot; pub(crate) const DEFAULT_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS: u16 = 10; pub fn capture_example(    project: Entity<Project>,    buffer: Entity<Buffer>,    cursor_anchor: language::Anchor,    mut events: Vec<StoredEvent>,    cx: &mut App,) -> Option<Task<Result<ExampleSpec>>> {    let snapshot = buffer.read(cx).snapshot();    let file = snapshot.file()?;    let worktree_id = file.worktree_id(cx);    let repository = project.read(cx).active_repository(cx)?;    let repository_snapshot = repository.read(cx).snapshot();    let worktree = project.read(cx).worktree_for_id(worktree_id, cx)?;    let cursor_path = worktree.read(cx).root_name().join(file.path());    if worktree.read(cx).abs_path() != repository_snapshot.work_directory_abs_path {        return None;    }     let repository_url = repository_snapshot        .remote_origin_url        .clone()        .or_else(|| repository_snapshot.remote_upstream_url.clone())?;    let revision = repository_snapshot.head_commit.as_ref()?.sha.to_string();     let git_store = project.read(cx).git_store().clone();     Some(cx.spawn(async move |mut cx| {        let snapshots_by_path =            collect_snapshots(&project, &git_store, worktree_id, &events, &mut cx).await?;         events.retain(|stored_event| {            match stored_event.event.as_ref() {                zeta_prompt::Event::BufferChange { path, .. } => {                    if !snapshots_by_path.contains_key(path) {                        return false;                    }                }            }            true        });         let line_comment_prefix = snapshot            .language()            .and_then(|lang| lang.config().line_comments.first())            .map(|s| s.to_string())            .unwrap_or_default();        let (cursor_excerpt, cursor_offset) = cx            .background_executor()            .spawn(async move { compute_cursor_excerpt(&snapshot, cursor_anchor) })            .await;        let uncommitted_diff = cx            .background_executor()            .spawn(async move { compute_uncommitted_diff(snapshots_by_path) })            .await;         let mut edit_history = String::new();        for stored_event in &events {            zeta_prompt::write_event(&mut edit_history, &stored_event.event);            if !edit_history.ends_with('\n') {                edit_history.push('\n');            }        }         let mut spec = ExampleSpec {            name: generate_timestamp_name(),            repository_url,            revision,            tags: Vec::new(),            reasoning: None,            uncommitted_diff,            cursor_path: cursor_path.as_std_path().into(),            cursor_position: String::new(),            edit_history,            expected_patches: Vec::new(),        };        spec.set_cursor_excerpt(&cursor_excerpt, cursor_offset, &line_comment_prefix);        Ok(spec)    }))} fn compute_cursor_excerpt(    snapshot: &language::BufferSnapshot,    cursor_anchor: language::Anchor,) -> (String, usize) {    use text::ToOffset as _;     let cursor_point = cursor_anchor.to_point(snapshot);    let (_editable_range, context_range) =        editable_and_context_ranges_for_cursor_position(cursor_point, snapshot, 100, 50);    let context_start_offset = context_range.start.to_offset(snapshot);    let cursor_offset = cursor_anchor.to_offset(snapshot);    let cursor_offset_in_excerpt = cursor_offset.saturating_sub(context_start_offset);    let excerpt = snapshot.text_for_range(context_range).collect::<String>();    (excerpt, cursor_offset_in_excerpt)} async fn collect_snapshots(    project: &Entity<Project>,    git_store: &Entity<project::git_store::GitStore>,    worktree_id: WorktreeId,    events: &[StoredEvent],    cx: &mut gpui::AsyncApp,) -> Result<HashMap<Arc<Path>, (TextBufferSnapshot, BufferDiffSnapshot)>> {    let mut snapshots_by_path = HashMap::default();    let root_name = project.read_with(cx, |project, cx| {        project            .worktree_for_id(worktree_id, cx)            .unwrap()            .read(cx)            .root_name()            .to_owned()    })?;    for stored_event in events {        let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref();        if let Some((project_path, full_path)) = project.read_with(cx, |project, cx| {            let project_path = project                .find_project_path(path, cx)                .filter(|path| path.worktree_id == worktree_id)?;            let full_path = root_name.join(&project_path.path).as_std_path().into();            Some((project_path, full_path))        })? {            if let hash_map::Entry::Vacant(entry) = snapshots_by_path.entry(full_path) {                let buffer = project                    .update(cx, |project, cx| {                        project.open_buffer(project_path.clone(), cx)                    })?                    .await?;                let diff = git_store                    .update(cx, |git_store, cx| {                        git_store.open_uncommitted_diff(buffer.clone(), cx)                    })?                    .await?;                let diff_snapshot = diff.update(cx, |diff, cx| diff.snapshot(cx))?;                entry.insert((stored_event.old_snapshot.clone(), diff_snapshot));            }        }    }    Ok(snapshots_by_path)} fn compute_uncommitted_diff(    snapshots_by_path: HashMap<Arc<Path>, (TextBufferSnapshot, BufferDiffSnapshot)>,) -> String {    let mut uncommitted_diff = String::new();    for (full_path, (before_text, diff_snapshot)) in snapshots_by_path {        if let Some(head_text) = &diff_snapshot.base_text_string() {            let file_diff = language::unified_diff(head_text, &before_text.text());            if !file_diff.is_empty() {                let path_str = full_path.to_string_lossy();                writeln!(uncommitted_diff, "--- a/{path_str}").ok();                writeln!(uncommitted_diff, "+++ b/{path_str}").ok();                uncommitted_diff.push_str(&file_diff);                if !uncommitted_diff.ends_with('\n') {                    uncommitted_diff.push('\n');                }            }        }    }    uncommitted_diff} fn generate_timestamp_name() -> String {    let format = time::format_description::parse("[year]-[month]-[day] [hour]:[minute]:[second]");    match format {        Ok(format) => {            let now = time::OffsetDateTime::now_local()                .unwrap_or_else(|_| time::OffsetDateTime::now_utc());            now.format(&format)                .unwrap_or_else(|_| "unknown-time".to_string())        }        Err(_) => "unknown-time".to_string(),    }} pub(crate) fn should_sample_edit_prediction_example_capture(cx: &App) -> bool {    let capture_rate = language::language_settings::all_language_settings(None, cx)        .edit_predictions        .example_capture_rate        .unwrap_or(DEFAULT_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS);    cx.has_flag::<EditPredictionExampleCaptureFeatureFlag>()        && rand::random::<u16>() % 10_000 < capture_rate} #[cfg(test)]mod tests {    use super::*;    use crate::EditPredictionStore;    use client::{Client, UserStore};    use clock::FakeSystemClock;    use gpui::{AppContext as _, TestAppContext, http_client::FakeHttpClient};    use indoc::indoc;    use language::{Anchor, Point};    use project::{FakeFs, Project};    use serde_json::json;    use settings::SettingsStore;    use std::path::Path;     #[gpui::test]    async fn test_capture_example(cx: &mut TestAppContext) {        init_test(cx);        let fs = FakeFs::new(cx.executor());         let committed_contents = indoc! {"            fn main() {                one();                two();                three();                four();                five();                six();                seven();                eight();                nine();            }        "};         let disk_contents = indoc! {"            fn main() {                // comment 1                one();                two();                three();                four();                five();                six();                seven();                eight();                // comment 2                nine();            }        "};         fs.insert_tree(            "/project",            json!({                ".git": {},                "src": {                    "main.rs": disk_contents,                }            }),        )        .await;         fs.set_head_for_repo(            Path::new("/project/.git"),            &[("src/main.rs", committed_contents.to_string())],            "abc123def456",        );        fs.set_remote_for_repo(            Path::new("/project/.git"),            "origin",            "https://github.com/test/repo.git",        );         let project = Project::test(fs.clone(), ["/project".as_ref()], cx).await;         let buffer = project            .update(cx, |project, cx| {                project.open_local_buffer("/project/src/main.rs", cx)            })            .await            .unwrap();         let ep_store = cx.read(|cx| EditPredictionStore::try_global(cx).unwrap());        ep_store.update(cx, |ep_store, cx| {            ep_store.register_buffer(&buffer, &project, cx)        });        cx.run_until_parked();         buffer.update(cx, |buffer, cx| {            let point = Point::new(6, 0);            buffer.edit([(point..point, "    // comment 3\n")], None, cx);            let point = Point::new(4, 0);            buffer.edit([(point..point, "    // comment 4\n")], None, cx);             pretty_assertions::assert_eq!(                buffer.text(),                indoc! {"                    fn main() {                        // comment 1                        one();                        two();                        // comment 4                        three();                        four();                        // comment 3                        five();                        six();                        seven();                        eight();                        // comment 2                        nine();                    }                "}            );        });        cx.run_until_parked();         let events = ep_store.update(cx, |store, cx| {            store.edit_history_for_project_with_pause_split_last_event(&project, cx)        });        let mut example = cx            .update(|cx| {                capture_example(project.clone(), buffer.clone(), Anchor::MIN, events, cx).unwrap()            })            .await            .unwrap();        example.name = "test".to_string();         pretty_assertions::assert_eq!(            example,            ExampleSpec {                name: "test".to_string(),                repository_url: "https://github.com/test/repo.git".to_string(),                revision: "abc123def456".to_string(),                tags: Vec::new(),                reasoning: None,                uncommitted_diff: indoc! {"                    --- a/project/src/main.rs                    +++ b/project/src/main.rs                    @@ -1,4 +1,5 @@                     fn main() {                    +    // comment 1                         one();                         two();                         three();                    @@ -7,5 +8,6 @@                         six();                         seven();                         eight();                    +    // comment 2                         nine();                     }                "}                .to_string(),                cursor_path: Path::new("project/src/main.rs").into(),                cursor_position: indoc! {"                    fn main() {                    ^[CURSOR_POSITION]                        // comment 1                        one();                        two();                        // comment 4                        three();                        four();                        // comment 3                        five();                        six();                        seven();                        eight();                        // comment 2                        nine();                    }                "}                .to_string(),                edit_history: indoc! {"                    --- a/project/src/main.rs                    +++ b/project/src/main.rs                    @@ -2,8 +2,10 @@                         // comment 1                         one();                         two();                    +    // comment 4                         three();                         four();                    +    // comment 3                         five();                         six();                         seven();                "}                .to_string(),                expected_patches: Vec::new()            }        );    }     fn init_test(cx: &mut TestAppContext) {        cx.update(|cx| {            let settings_store = SettingsStore::test(cx);            cx.set_global(settings_store);            zlog::init_test();            let http_client = FakeHttpClient::with_404_response();            let client = Client::new(Arc::new(FakeSystemClock::new()), http_client, cx);            language_model::init(client.clone(), cx);            let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));            EditPredictionStore::global(&client, &user_store, cx);        })    }}