use crate::{
    example::{Example, ExamplePromptInputs},
    headless::EpAppState,
    load_project::run_load_project,
    progress::{ExampleProgress, InfoStyle, Step, StepProgress},
};
use anyhow::Context as _;
use collections::HashSet;
use edit_prediction::{DebugEvent, EditPredictionStore};
use futures::{FutureExt as _, StreamExt as _, channel::mpsc};
use gpui::{AsyncApp, Entity};
use language::Buffer;
use project::Project;
use std::sync::Arc;
use std::time::Duration;

pub async fn run_context_retrieval(
    example: &mut Example,
    app_state: Arc<EpAppState>,
    example_progress: &ExampleProgress,
    mut cx: AsyncApp,
) -> anyhow::Result<()> {
    if example
        .prompt_inputs
        .as_ref()
        .is_some_and(|inputs| inputs.related_files.is_some())
    {
        return Ok(());
    }

    if let Some(captured) = &example.spec.captured_prompt_input {
        let step_progress = example_progress.start(Step::Context);
        step_progress.set_substatus("using captured prompt input");

        let edit_history: Vec<Arc<zeta_prompt::Event>> = captured
            .events
            .iter()
            .map(|e| Arc::new(e.to_event()))
            .collect();

        let related_files: Vec<zeta_prompt::RelatedFile> = captured
            .related_files
            .iter()
            .map(|rf| rf.to_related_file())
            .collect();

        example.prompt_inputs = Some(ExamplePromptInputs {
            content: captured.cursor_file_content.clone(),
            cursor_row: captured.cursor_row,
            cursor_column: captured.cursor_column,
            cursor_offset: captured.cursor_offset,
            excerpt_start_row: captured.excerpt_start_row,
            edit_history,
            related_files: Some(related_files),
        });

        return Ok(());
    }

    run_load_project(example, app_state.clone(), example_progress, cx.clone()).await?;

    let step_progress: Arc<StepProgress> = example_progress.start(Step::Context).into();

    let state = example.state.as_ref().unwrap();
    let project = state.project.clone();

    let _lsp_handle = project.update(&mut cx, |project, cx| {
        project.register_buffer_with_language_servers(&state.buffer, cx)
    });
    wait_for_language_servers_to_start(&project, &state.buffer, &step_progress, &mut cx).await?;

    let ep_store = cx
        .update(|cx| EditPredictionStore::try_global(cx))
        .context("EditPredictionStore not initialized")?;

    let mut events = ep_store.update(&mut cx, |store, cx| {
        store.register_buffer(&state.buffer, &project, cx);
        store.refresh_context(&project, &state.buffer, state.cursor_position, cx);
        store.debug_info(&project, cx)
    });

    while let Some(event) = events.next().await {
        match event {
            DebugEvent::ContextRetrievalFinished(_) => {
                break;
            }
            _ => {}
        }
    }

    let context_files =
        ep_store.update(&mut cx, |store, cx| store.context_for_project(&project, cx));

    let excerpt_count: usize = context_files.iter().map(|f| f.excerpts.len()).sum();
    step_progress.set_info(format!("{} excerpts", excerpt_count), InfoStyle::Normal);

    if let Some(prompt_inputs) = example.prompt_inputs.as_mut() {
        prompt_inputs.related_files = Some(context_files);
    }
    Ok(())
}

async fn wait_for_language_servers_to_start(
    project: &Entity<Project>,
    buffer: &Entity<Buffer>,
    step_progress: &Arc<StepProgress>,
    cx: &mut AsyncApp,
) -> anyhow::Result<()> {
    let lsp_store = project.read_with(cx, |project, _| project.lsp_store());

    let (language_server_ids, mut starting_language_server_ids) =
        buffer.update(cx, |buffer, cx| {
            lsp_store.update(cx, |lsp_store, cx| {
                let ids = lsp_store.language_servers_for_local_buffer(buffer, cx);
                let starting_ids = ids
                    .iter()
                    .copied()
                    .filter(|id| !lsp_store.language_server_statuses.contains_key(&id))
                    .collect::<HashSet<_>>();
                (ids, starting_ids)
            })
        });

    step_progress.set_substatus(format!("waiting for {} LSPs", language_server_ids.len()));

    let timeout_duration = if starting_language_server_ids.is_empty() {
        Duration::from_secs(30)
    } else {
        Duration::from_secs(60 * 5)
    };

    let timeout = cx.background_executor().timer(timeout_duration).shared();

    let (mut tx, mut rx) = mpsc::channel(language_server_ids.len());
    let added_subscription = cx.subscribe(project, {
        let step_progress = step_progress.clone();
        move |_, event, _| match event {
            project::Event::LanguageServerAdded(language_server_id, name, _) => {
                step_progress.set_substatus(format!("LSP started: {}", name));
                tx.try_send(*language_server_id).ok();
            }
            _ => {}
        }
    });

    while !starting_language_server_ids.is_empty() {
        futures::select! {
            language_server_id = rx.next() => {
                if let Some(id) = language_server_id {
                    starting_language_server_ids.remove(&id);
                }
            },
            _ = timeout.clone().fuse() => {
                return Err(anyhow::anyhow!("LSP wait timed out after {} minutes", timeout_duration.as_secs() / 60));
            }
        }
    }

    drop(added_subscription);

    let (mut tx, mut rx) = mpsc::channel(language_server_ids.len());
    let subscriptions = [
        cx.subscribe(&lsp_store, {
            let step_progress = step_progress.clone();
            move |_, event, _| {
                if let project::LspStoreEvent::LanguageServerUpdate {
                    message:
                        client::proto::update_language_server::Variant::WorkProgress(
                            client::proto::LspWorkProgress {
                                message: Some(message),
                                ..
                            },
                        ),
                    ..
                } = event
                {
                    step_progress.set_substatus(message.clone());
                }
            }
        }),
        cx.subscribe(project, {
            let step_progress = step_progress.clone();
            let lsp_store = lsp_store.clone();
            move |_, event, cx| match event {
                project::Event::DiskBasedDiagnosticsFinished { language_server_id } => {
                    let lsp_store = lsp_store.read(cx);
                    let name = lsp_store
                        .language_server_adapter_for_id(*language_server_id)
                        .unwrap()
                        .name();
                    step_progress.set_substatus(format!("LSP idle: {}", name));
                    tx.try_send(*language_server_id).ok();
                }
                _ => {}
            }
        }),
    ];

    project
        .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))
        .await?;

    let mut pending_language_server_ids = lsp_store.read_with(cx, |lsp_store, _| {
        language_server_ids
            .iter()
            .copied()
            .filter(|id| {
                lsp_store
                    .language_server_statuses
                    .get(id)
                    .is_some_and(|status| status.has_pending_diagnostic_updates)
            })
            .collect::<HashSet<_>>()
    });
    while !pending_language_server_ids.is_empty() {
        futures::select! {
            language_server_id = rx.next() => {
                if let Some(id) = language_server_id {
                    pending_language_server_ids.remove(&id);
                }
            },
            _ = timeout.clone().fuse() => {
                return Err(anyhow::anyhow!("LSP wait timed out after {} minutes", timeout_duration.as_secs() / 60));
            }
        }
    }

    drop(subscriptions);
    step_progress.clear_substatus();
    Ok(())
}