use std::collections::HashMap;
use std::sync::Arc;
#[cfg(feature = "local_fs")]
use std::time::Duration;

use chrono::Utc;
use futures::stream::AbortHandle;
use galaxy_completer::completer::{
    self, expand_command_aliases, AliasExpansionResult, CompleterOptions,
    CompletionsFallbackStrategy, MatchStrategy,
};
use galaxy_completer::meta::Spanned;
use galaxy_completer::parsers::hir::{Command, Expression, FlagType};
use galaxy_completer::parsers::ParsedExpression;
use galaxy_core::features::FeatureFlag;
#[cfg(feature = "local_fs")]
use galaxyui::r#async::FutureExt;
use galaxyui::{AppContext, Entity, ModelContext, ModelHandle, SingletonEntity};
use itertools::Itertools;
#[cfg_attr(not(feature = "local_fs"), allow(unused_imports))]
use parking_lot::{FairMutex, Mutex};

use super::generate_ai_input_suggestions::{
    create_generate_ai_input_suggestions_request, get_context_messages,
    GenerateAIInputSuggestionsRequest, GenerateAIInputSuggestionsResponseV2, NextCommandContext,
};
use crate::ai::block_context::BlockContext;
use crate::ai_assistant::execution_context::WarpAiExecutionContext;
use crate::completer::SessionContext;
#[cfg(feature = "local_fs")]
use crate::persistence::{database_file_path_for_scope, establish_ro_connection, PersistenceScope};
use crate::report_error;
use crate::server::server_api::{AIApiError, ServerApi};
use crate::settings::AISettings;
use crate::terminal::event::UserBlockCompleted;
use crate::terminal::input::{CompleterData, IntelligentAutosuggestionResult};
use crate::terminal::model::session::Sessions;
use crate::terminal::{History, HistoryEntry, TerminalModel};
use crate::workspaces::user_workspaces::UserWorkspaces;

cfg_if::cfg_if! {
    if #[cfg(feature = "local_fs")] {
        use diesel::SqliteConnection;
        use std::path::PathBuf;
        use galaxy_completer::parsers::hir::ArgType;
    }
}

#[cfg_attr(not(feature = "local_fs"), allow(dead_code))]
const MAX_NUM_SIMILAR_HISTORY_CONTEXT: usize = 25;

/// The number of additional preceding commands for each HistoryContext
/// included in the LLM request.
#[cfg_attr(not(feature = "local_fs"), allow(dead_code))]
const NUM_ADDITIONAL_PREV_COMMAND_CONTEXT_LLM: usize = 2;

#[cfg(feature = "local_fs")]
const ARG_GENERATOR_VALIDATION_TIMEOUT: Duration = Duration::from_millis(150);

pub fn is_next_command_enabled(app: &galaxyui::AppContext) -> bool {
    AISettings::as_ref(app).is_intelligent_autosuggestions_enabled(app)
        && UserWorkspaces::as_ref(app).is_next_command_enabled()
}

/// Information about an autosuggestion that would have been made if purely based off history.
/// If there was no history, history_command_prediction would be an empty string.
#[derive(Clone, Default, PartialEq, Debug)]
pub struct HistoryBasedAutosuggestionState {
    pub history_command_prediction: String,
    pub history_command_prediction_likelihood: f64,
    pub total_history_count: usize,
}

#[derive(Clone, Default, PartialEq)]
pub enum NextCommandSuggestionState {
    #[default]
    None,
    Cycling,
    Ready {
        request: Box<GenerateAIInputSuggestionsRequest>,
        response: GenerateAIInputSuggestionsResponseV2,
        /// How long the request took to complete, in milliseconds.
        request_duration_ms: i64,
        /// If true, we made a call to an LLM to generate this.
        /// Otherwise it came from history.
        is_from_ai: bool,
        /// If true, this suggestion came from the user explicitly cycling (is not the initial suggestion).
        is_from_cycle: bool,
        history_based_autosuggestion_state: HistoryBasedAutosuggestionState,
    },
}

impl NextCommandSuggestionState {
    pub fn is_ready(&self) -> bool {
        matches!(self, NextCommandSuggestionState::Ready { .. })
    }

    pub fn is_cycling(&self) -> bool {
        matches!(self, NextCommandSuggestionState::Cycling)
    }

    pub fn command_suggestion(&self) -> Option<&str> {
        match &self {
            // The server only returns one command suggestion as the most likely action.
            NextCommandSuggestionState::Ready { response, .. } => {
                let command = &response.most_likely_action;
                // If AI accidentally returned JSON instead of a plain string for the most likely action, don't use it.
                if command.starts_with('{') {
                    return None;
                }
                Some(command)
            }
            _ => None,
        }
    }
}

/// Struct storing the result of the zero-state next command suggestion,
/// used for telemetry purposes.
#[derive(Clone)]
pub struct ZeroStateSuggestionInfo {
    pub request: Box<GenerateAIInputSuggestionsRequest>,
    pub response: GenerateAIInputSuggestionsResponseV2,
    /// How long the request took to complete, in milliseconds.
    pub request_duration_ms: i64,
    /// If true, we made a call to an LLM to generate this.
    /// Otherwise it came from history.
    pub is_from_ai: bool,
    pub history_based_autosuggestion_state: HistoryBasedAutosuggestionState,
}

pub struct NextCommandModel {
    sessions: ModelHandle<Sessions>,
    model: Arc<FairMutex<TerminalModel>>,
    server_api: Arc<ServerApi>,
    #[cfg(feature = "local_fs")]
    conn: Option<Arc<Mutex<SqliteConnection>>>,

    next_command_state: NextCommandSuggestionState,
    /// Context used to generate the zero-state suggestion.
    /// We reuse this but apply filtering on history_contexts as the user edits the input.
    cached_zerostate_next_command_context: Option<NextCommandContext>,
    zerostate_suggestion_info: Option<ZeroStateSuggestionInfo>,
    next_command_abort_handle: Option<AbortHandle>,
}

impl Entity for NextCommandModel {
    type Event = NextCommandModelEvent;
}

pub enum NextCommandModelEvent {
    NextCommandSuggestionReady,
}

impl NextCommandModel {
    pub fn new(
        sessions: ModelHandle<Sessions>,
        model: Arc<FairMutex<TerminalModel>>,
        server_api: Arc<ServerApi>,
    ) -> Self {
        #[cfg(feature = "local_fs")]
        let conn = database_file_path_for_scope(&PersistenceScope::App)
            .to_str()
            .and_then(|db_url| {
                establish_ro_connection(db_url)
                    .ok()
                    .map(|conn| Arc::new(Mutex::new(conn)))
            });
        Self {
            sessions,
            model,
            server_api,
            #[cfg(feature = "local_fs")]
            conn,
            next_command_state: NextCommandSuggestionState::None,
            cached_zerostate_next_command_context: None,
            zerostate_suggestion_info: None,
            next_command_abort_handle: None,
        }
    }

    /// Returns snippets of command history (HistoryContext) that are similar to the completed_block.
    /// Each HistoryContext contains some sequential commands run in the same session,
    /// where the last element of HistoryContext.previous_commands is the same as completed_block.
    /// Returns None if there was a connection issue, and Some(empty vec)
    /// if there is no similar historical context.
    #[cfg(feature = "local_fs")]
    pub fn get_similar_history_context(
        conn: &mut SqliteConnection,
        completed_block: &UserBlockCompleted,
        num_additional_preceding_commands: usize,
    ) -> Vec<crate::ai::predict::generate_ai_input_suggestions::HistoryContext> {
        // The number of commands from history affects how quickly we "learn" new patterns, the lower the faster.
        let Ok(same_commands_from_history) =
            crate::persistence::commands::get_same_commands_from_history(
                conn,
                completed_block,
                MAX_NUM_SIMILAR_HISTORY_CONTEXT,
            )
        else {
            return vec![];
        };
        // Iterate from oldest to newest
        same_commands_from_history
            .into_iter()
            .rev()
            .filter_map(|command| {
                let next_command =
                    crate::persistence::commands::get_next_command(conn, &command).ok()?;
                if num_additional_preceding_commands == 0 {
                    return Some(
                        crate::ai::predict::generate_ai_input_suggestions::HistoryContext {
                            previous_commands: vec![command],
                            next_command,
                        },
                    );
                }
                // We know next_command comes after command.
                // Get some more commands that came before command so there's additional context before next_command.
                let mut previous_commands = crate::persistence::commands::get_previous_commands(
                    conn,
                    &command,
                    num_additional_preceding_commands,
                )
                .ok()?;
                previous_commands.push(command);
                Some(
                    crate::ai::predict::generate_ai_input_suggestions::HistoryContext {
                        previous_commands,
                        next_command,
                    },
                )
            })
            .collect()
    }

    pub fn get_state(&self) -> &NextCommandSuggestionState {
        &self.next_command_state
    }

    pub fn get_zero_state_suggestion_info(&self) -> Option<&ZeroStateSuggestionInfo> {
        self.zerostate_suggestion_info.as_ref()
    }

    pub fn clear_state(&mut self) {
        self.next_command_state = NextCommandSuggestionState::None;
        self.cached_zerostate_next_command_context = None;
        self.zerostate_suggestion_info = None;
        self.abort_inflight_request();
    }

    pub fn abort_inflight_request(&mut self) {
        if let Some(abort_handle) = self.next_command_abort_handle.take() {
            abort_handle.abort();
        }
    }

    pub fn cycle_next_command_suggestion(&mut self, _ctx: &mut ModelContext<Self>) {
        // TODO(roland): reconsider down arrow UX for next command or remove completely
    }

    #[cfg_attr(not(feature = "local_fs"), allow(unused_variables))]
    fn get_next_command_context(
        terminal_model: Arc<FairMutex<TerminalModel>>,
        #[cfg(feature = "local_fs")] conn: Option<Arc<Mutex<SqliteConnection>>>,
        ai_execution_context: WarpAiExecutionContext,
        block_completed: &UserBlockCompleted,
    ) -> NextCommandContext {
        #[cfg_attr(not(feature = "local_fs"), allow(unused_mut))]
        let mut history_contexts = vec![];
        let context_messages = get_context_messages(terminal_model.clone(), 5, 100, 200);
        #[cfg(feature = "local_fs")]
        if let Some(conn) = conn {
            let mut conn = conn.lock();
            history_contexts = Self::get_similar_history_context(
                &mut conn,
                block_completed,
                NUM_ADDITIONAL_PREV_COMMAND_CONTEXT_LLM,
            );
        }
        NextCommandContext {
            history_contexts,
            ai_execution_context,
            context_messages,
        }
    }

    /// Generates a zero-state next command suggestion immediately after a block completes.
    pub fn generate_next_command_suggestion(
        &mut self,
        block_completed: UserBlockCompleted,
        context: WarpAiExecutionContext,
        completer_data: CompleterData,
        block_context: Option<Box<BlockContext>>,
        previous_result: Option<IntelligentAutosuggestionResult>,
        ctx: &mut ModelContext<Self>,
    ) {
        // Clear the cached next command context so we don't use stale data.
        self.cached_zerostate_next_command_context = None;
        self.zerostate_suggestion_info = None;
        self.generate_next_command_suggestion_with_prefix(
            None,
            block_completed,
            context,
            completer_data,
            block_context,
            previous_result,
            ctx,
        );
    }

    /// Returns the most recent command with a matching prefix run in the user's current working directory.
    /// If no such command exists, returns the most recent command anywhere with a matching prefix.
    pub fn get_reverse_chronological_potential_autosuggestions(
        prefix: &str,
        completer_data: &CompleterData,
        app: &AppContext,
    ) -> Option<Vec<HistoryEntry>> {
        let session_id = completer_data.active_block_session_id()?;
        let history_entries = History::as_ref(app).commands(session_id)?;
        let working_dir = completer_data
            .active_block_metadata
            .as_ref()
            .and_then(|block_metadata| block_metadata.current_working_directory());
        Some(find_potential_autosuggestions_from_history(
            history_entries.into_iter(),
            prefix,
            working_dir,
        ))
    }

    /// Generates a next command suggestion with a prefix in the input.
    #[cfg_attr(not(feature = "local_fs"), allow(unused_variables))]
    #[expect(clippy::too_many_arguments)]
    pub fn generate_next_command_suggestion_with_prefix(
        &mut self,
        prefix: Option<String>,
        block_completed: UserBlockCompleted,
        context: WarpAiExecutionContext,
        completer_data: CompleterData,
        block_context: Option<Box<BlockContext>>,
        previous_result: Option<IntelligentAutosuggestionResult>,
        ctx: &mut ModelContext<Self>,
    ) {
        let server_api = self.server_api.clone();
        let terminal_model = self.model.clone();
        let cached_next_command_context = self.cached_zerostate_next_command_context.clone();

        let completion_context = completer_data.completion_session_context(ctx);
        // This is only needed if we have a prefix.
        let reverse_chronological_potential_autosuggestions = if let Some(prefix) = &prefix {
            Self::get_reverse_chronological_potential_autosuggestions(prefix, &completer_data, ctx)
        } else {
            None
        };

        #[cfg(feature = "local_fs")]
        let conn = self.conn.clone();
        self.next_command_state = NextCommandSuggestionState::None;
        self.abort_inflight_request();
        let session_env_vars = completer_data
            .active_block_session_id()
            .and_then(|session_id| {
                self.sessions.read(ctx, |sessions, _| {
                    sessions.get_env_vars_for_session(session_id)
                })
            });
        self.next_command_abort_handle = Some(
            ctx.spawn(
                async move {
                    let mut history_based_autosuggestion_state =
                        HistoryBasedAutosuggestionState::default();
                    let start_ts_ms = Utc::now().timestamp_millis();

                    let mut next_command_context =
                        if let Some(cached_next_command_context) = cached_next_command_context {
                            cached_next_command_context
                        } else {
                            Self::get_next_command_context(
                                terminal_model,
                                #[cfg(feature = "local_fs")]
                                conn,
                                context,
                                &block_completed,
                            )
                        };
                    // Filter history contexts for only cases where the next command matches our prefix.
                    if let Some(prefix) = &prefix {
                        next_command_context.history_contexts = next_command_context
                            .history_contexts
                            .into_iter()
                            .filter(|context| context.next_command.command.starts_with(prefix))
                            .collect_vec();
                    }
                    // First, use rich history to find commands with a matching prefix that were run
                    // in a similar context, taking into account the most recent block run.
                    if !next_command_context.history_contexts.is_empty() {
                        let mut history_next_command_counts = counter::Counter::new();
                        for history_context in &next_command_context.history_contexts {
                            history_next_command_counts[&history_context.next_command.command] += 1;
                        }

                        let mut total_history_count = history_next_command_counts.total::<usize>();
                        let most_likely_next_commands =
                            history_next_command_counts.k_most_common_ordered(5);
                        for (most_likely_next_command, count) in &most_likely_next_commands {
                            if !is_command_valid(most_likely_next_command, completion_context.as_ref(), session_env_vars.as_ref()).await {
                                log::debug!("Discarding most likely next command from rich history that failed validation: `{most_likely_next_command}`");
                                total_history_count -= *count;
                                continue;
                            }
                            let history_command_prediction_likelihood =
                                *count as f64 / total_history_count as f64;
                            history_based_autosuggestion_state = HistoryBasedAutosuggestionState {
                                history_command_prediction: most_likely_next_command.to_owned(),
                                history_command_prediction_likelihood,
                                total_history_count,
                            };

                            // If one command is very likely based on history, skip the LLM call and return it directly.
                            // Use history-based autosuggestion if there are a min number of similar cases in history
                            // AND the same command was run at least skip_llm_confidence_threshold of the time.
                            // Partial autosuggestions with a prefix have more lenient requirements
                            // because latency matters more as the user is typing.
                            let (min_history_count, skip_llm_confidence_threshold) =
                                if prefix.is_some() {
                                    (1, 0.1)
                                } else {
                                    (2, 0.25)
                                };
                            if total_history_count >= min_history_count
                                && history_command_prediction_likelihood
                                    >= skip_llm_confidence_threshold
                            {
                                // We construct the request even though we're not sending it to the server because
                                // it might be used later for cycling next command suggestions.
                                let request = create_generate_ai_input_suggestions_request(
                                    next_command_context.clone(),
                                    prefix,
                                    block_context,
                                    previous_result,
                                );
                                return (
                                    Ok(GenerateAIInputSuggestionsResponseV2 {
                                        commands: vec![most_likely_next_command.to_owned()],
                                        ai_queries: vec![],
                                        most_likely_action: most_likely_next_command.to_owned(),
                                    }),
                                    request,
                                    false,
                                    start_ts_ms,
                                    history_based_autosuggestion_state,
                                    false,
                                    next_command_context,
                                );
                            }
                        }
                    }
                    let request = create_generate_ai_input_suggestions_request(
                        next_command_context.clone(),
                        prefix.clone(),
                        block_context,
                        previous_result,
                    );

                    // For zero-state next command suggestions, return the result immediately.
                    let Some(prefix) = prefix else {
                        return (
                            server_api.generate_ai_input_suggestions(&request).await,
                            request,
                            true,
                            start_ts_ms,
                            history_based_autosuggestion_state,
                            false,
                            next_command_context,
                        );
                    };

                    // At this point we know we're generating a partial suggestion with a prefix.
                    // First, return the most recent command with a matching prefix run in the same pwd
                    // (if exists, otherwise just most recent command anywhere with matching prefix).
                    for reverse_chronological_command in reverse_chronological_potential_autosuggestions.unwrap_or_default() {
                        if is_command_valid(&reverse_chronological_command.command, completion_context.as_ref(), session_env_vars.as_ref()).await {
                            return (
                                Ok(GenerateAIInputSuggestionsResponseV2 {
                                    commands: vec![reverse_chronological_command.command.clone()],
                                ai_queries: vec![],
                                most_likely_action: reverse_chronological_command.command,
                            }),
                            request,
                            false,
                            start_ts_ms,
                                history_based_autosuggestion_state,
                                false,
                                next_command_context,
                            );
                        }
                    }

                    // If we have no command anywhere in history with a matching prefix, fallback to the first completer result.
                    if let Some(completion_context) = completion_context {
                        let completion_result = completer::suggestions(
                            &prefix,
                            prefix.len(),
                            session_env_vars.as_ref(),
                            CompleterOptions {
                                match_strategy: MatchStrategy::CaseSensitive,
                                fallback_strategy: CompletionsFallbackStrategy::None,
                                suggest_file_path_completions_only: false,
                                parse_quotes_as_literals: false,
                            },
                            &completion_context,
                        )
                        .await;

                        let autosuggestion = completion_result.and_then(|result| {
                            let replacement_span = result.replacement_span;
                            result.suggestions.into_iter().next().map(|s| {
                                // Reproduce the final buffer text with the autosuggestion since the
                                // completer only gives the replacement span of the suggestion.
                                let result = format!(
                                    "{}{}",
                                    &prefix[0..replacement_span.start()],
                                    s.replacement()
                                );
                                result
                            })
                        });

                        if let Some(autosuggestion) = autosuggestion {
                            if is_command_valid(&autosuggestion, Some(&completion_context), session_env_vars.as_ref()).await {
                                return (
                                    Ok(GenerateAIInputSuggestionsResponseV2 {
                                        commands: vec![autosuggestion.clone()],
                                    ai_queries: vec![],
                                    most_likely_action: autosuggestion,
                                }),
                                request,
                                false,
                                    start_ts_ms,
                                    history_based_autosuggestion_state,
                                    false,
                                    next_command_context,
                                );
                            }
                        }
                    };

                    // Only if we have no commands from history and no completions, use the LLM to generate a partial suggestion.
                    let response = server_api.generate_ai_input_suggestions(&request).await;
                    (
                        response,
                        request,
                        true,
                        start_ts_ms,
                        history_based_autosuggestion_state,
                        false,
                        next_command_context,
                    )
                },
                Self::on_next_command_suggestion_result,
            )
            .abort_handle(),
        );
    }

    fn on_next_command_suggestion_result(
        &mut self,
        result: (
            Result<GenerateAIInputSuggestionsResponseV2, AIApiError>,
            GenerateAIInputSuggestionsRequest,
            bool,
            i64,
            HistoryBasedAutosuggestionState,
            bool,
            NextCommandContext,
        ),
        ctx: &mut ModelContext<Self>,
    ) {
        self.next_command_abort_handle = None;
        let (
            result,
            request,
            is_from_ai,
            start_ts_ms,
            history_based_autosuggestion_state,
            is_from_cycle,
            next_command_context,
        ) = result;
        let end_ts_ms = Utc::now().timestamp_millis();
        let request_duration_ms = end_ts_ms - start_ts_ms;
        if request.prefix.is_none() {
            self.cached_zerostate_next_command_context = Some(next_command_context);
        }
        match result {
            Ok(response) => {
                if let Some(prefix) = &request.prefix {
                    if !response.most_likely_action.starts_with(prefix) {
                        // This is not expected to happen because the server applies its own filtering,
                        // but check just in case.
                        log::warn!(
                            "Next command suggestion `{}` does not start with prefix `{}`.",
                            response.most_likely_action,
                            prefix
                        );
                        return;
                    }
                } else {
                    self.zerostate_suggestion_info = Some(ZeroStateSuggestionInfo {
                        request: Box::new(request.clone()),
                        response: response.clone(),
                        request_duration_ms,
                        is_from_ai,
                        history_based_autosuggestion_state: history_based_autosuggestion_state
                            .clone(),
                    });
                }

                self.next_command_state = NextCommandSuggestionState::Ready {
                    request: Box::new(request),
                    response,
                    request_duration_ms,
                    is_from_ai,
                    is_from_cycle,
                    history_based_autosuggestion_state,
                };
                ctx.emit(NextCommandModelEvent::NextCommandSuggestionReady);
            }
            Err(err) => {
                report_error!(
                    anyhow::anyhow!(err).context("Failed to generate Next Command suggestion")
                );
            }
        };
    }
}

impl SingletonEntity for NextCommandModel {}

/// Validates that the arg is valid given its type (e.g. filepath exists if it's a filepath arg).
/// This uses a file system call, so this function should be called only in background threads.
#[cfg_attr(not(feature = "local_fs"), allow(unused_variables))]
async fn is_arg_valid(
    full_command: &str,
    arg: &Spanned<ParsedExpression>,
    ctx: &SessionContext,
    session_env_vars: Option<&HashMap<String, String>>,
) -> bool {
    let Expression::ValidatableArgument(arg_types_to_validate) = arg.expression() else {
        return true;
    };
    // The expression shouldn't be parsed as a `ValidatableArgument` if the arg types are empty,
    // but we check this just in case.
    if arg_types_to_validate.is_empty() {
        return true;
    }
    cfg_if::cfg_if! {
        if #[cfg(feature = "local_fs")] {
            // If we have arg types to validate, the arg must pass validation for at least one of them.
            // If the argument has one or more generators, validate these last because they're more expensive
            // and we can check all generators together using completions suggestions.
            let mut has_generator_arg_type = false;
            for arg_type in arg_types_to_validate {
                match arg_type {
                    ArgType::File => {
                        let mut path_arg = PathBuf::from(arg.value().as_str());
                        if path_arg.is_relative() {
                            if let Ok(working_dir) = PathBuf::try_from(ctx.current_working_directory.clone()) {
                                path_arg = working_dir.join(path_arg);
                            }
                        }
                        if path_arg.is_file() {
                            return true;
                        }
                    }
                    ArgType::Folder => {
                        let mut path_arg = PathBuf::from(arg.value().as_str());
                        if path_arg.is_relative() {
                            if let Ok(working_dir) = PathBuf::try_from(ctx.current_working_directory.clone()) {
                                path_arg = working_dir.join(path_arg);
                            }
                        }
                        if path_arg.is_dir() {
                            return true;
                        }
                    }
                    ArgType::Generator(_) => {
                        has_generator_arg_type = true;
                    }
                };
            }
            if has_generator_arg_type {
                // We don't have completions implemented for feature flags like --features=with_local_server.
                // If arg is the span of `with_local_server`, attempting to complete on --features= to validate it will return no results.
                // We should only use completions to validate the arg if the previous character is whitespace, until completions handles this case.
                let prev_char = full_command.get(..arg.span.start()).and_then(|s| s.chars().next_back());
                if prev_char.is_some_and(|c| !c.is_whitespace()) {
                    return true;
                }
                // Running completions runs all generators, so we only need to do this once.
                // TODO(roland): this also generates completions from sources other than generators, which are unnecessary.
                // If performance becomes a concern, consider validating against generators sequentially and returning early if valid.
                // We use completions suggestions because it's simpler to implement and read.
                let completions_future = completer::suggestions(
                    full_command,
                    arg.span.start(),
                    session_env_vars,
                    CompleterOptions {
                        match_strategy: MatchStrategy::CaseSensitive,
                        fallback_strategy: CompletionsFallbackStrategy::None,
                        suggest_file_path_completions_only: false,
                        parse_quotes_as_literals: false,
                    },
                    ctx,
                );

                // If the completions call times out, assume the arg is valid.
                // This is necessary because some generators can hang (e.g. kubectl commands if the cluster isn't running).
                let Ok(completion_result) = completions_future.with_timeout(ARG_GENERATOR_VALIDATION_TIMEOUT).await else {
                    log::debug!("Generator validation for arg `{}` in command `{}` timed out - assuming it's valid", arg.value().as_str(), full_command);
                    return true;
                };

                let Some(completion_result) = completion_result else {
                    return true;
                };
                for suggestion in completion_result.suggestions {
                    if suggestion.display() == arg.value().as_str() {
                        return true;
                    }
                }
            }
            // If we didn't pass validation for any of the possible arg types, this arg is invalid.
            log::debug!("arg `{}` in command `{}` failed validation", arg.value().as_str(), full_command);
            false
        } else {
            true
        }
    }
}

/// Validates the command is valid.
/// Currently uses completions specs to check if parsing is successful, and validates
/// that any filepaths args actually exist on disk.
/// This uses a file system call, so this function should be called only in background threads.
pub async fn is_command_valid(
    command: &str,
    ctx: Option<&SessionContext>,
    session_env_vars: Option<&HashMap<String, String>>,
) -> bool {
    if !FeatureFlag::ValidateAutosuggestions.is_enabled() {
        return true;
    }
    let Some(ctx) = ctx else {
        return true;
    };
    let AliasExpansionResult {
        expanded_command_line,
        classified_command,
        ..
    } = expand_command_aliases(command, false, ctx).await;

    let Some(classified_command) = classified_command else {
        return true;
    };

    // We assume the command is valid on parse error because
    // 1. Our completion specs are not always comprehensive (unknown args/options cause parse error)
    // 2. Our parsing logic has some bugs that need to be investigated (INT-816)
    if let Some(err) = &classified_command.error {
        log::debug!(
            "Assuming command `{}` is valid because it failed to parse: {err:?}",
            expanded_command_line,
        );
        return true;
    }
    // If we can't classify the command, it means we don't have completion specs for it.
    // Assume it's valid.
    let Command::Classified(shell_command) = classified_command.command else {
        return true;
    };
    if let Some(positionals) = &shell_command.args.positionals {
        for positional in positionals {
            if !is_arg_valid(&expanded_command_line, positional, ctx, session_env_vars).await {
                return false;
            }
        }
    }
    if let Some(flags) = shell_command.args.flags {
        for flag in flags.iter() {
            if let FlagType::Argument { value } = &flag.flag_type {
                if !is_arg_valid(&expanded_command_line, value, ctx, session_env_vars).await {
                    return false;
                }
            }
        }
    }
    true
}

/// Scans the given history entries in reverse order for commands that start
/// with the buffer text to return as a potential autosuggestion. Prioritizes commands
/// in history that were executed in the user's current working directory,
/// with any command executed in other directories at the end.
fn find_potential_autosuggestions_from_history<'a>(
    history_entries: impl DoubleEndedIterator<Item = &'a HistoryEntry>,
    buffer_text: &str,
    working_dir: Option<&str>,
) -> Vec<HistoryEntry> {
    let mut commands_in_same_dir = vec![];
    let mut commands_in_other_dirs = vec![];
    for entry in history_entries.rev() {
        if !entry.command.starts_with(buffer_text) {
            continue;
        }
        let same_dir = entry
            .pwd
            .as_ref()
            .zip(working_dir)
            .is_some_and(|(pwd, working_dir)| pwd == working_dir);

        if same_dir {
            commands_in_same_dir.push(entry.clone());
        } else {
            commands_in_other_dirs.push(entry.clone());
        }
    }
    commands_in_same_dir.extend(commands_in_other_dirs);
    commands_in_same_dir
}

#[cfg(test)]
#[path = "next_command_model_tests.rs"]
mod tests;