//! Input queue restore and thread-input snapshot behavior for `ChatWidget`.

use std::collections::HashSet;

use crate::bottom_pane::ComposerDraftSnapshot;
use crate::bottom_pane::KillBufferSnapshot;

use super::user_messages::remap_colliding_paste_placeholders;
use super::*;

impl ChatWidget {
    pub(crate) fn take_kill_buffer_snapshot(&mut self) -> KillBufferSnapshot {
        self.bottom_pane.take_kill_buffer_snapshot()
    }

    pub(crate) fn restore_kill_buffer_snapshot(&mut self, snapshot: KillBufferSnapshot) {
        self.bottom_pane.restore_kill_buffer_snapshot(snapshot);
    }

    /// Restore the exact draft entered before the fully initialized composer became available.
    pub(crate) fn restore_startup_draft(&mut self, draft: ComposerDraftSnapshot) {
        self.bottom_pane
            .inherit_startup_sparkle(draft.sparkle_draft);
        let existing_draft = self.bottom_pane.composer_draft_snapshot();
        let existing_cursor = existing_draft.cursor;
        let existing_message = UserMessage {
            text: existing_draft.text,
            text_elements: existing_draft.text_elements,
            local_images: existing_draft.local_images,
            remote_image_urls: existing_draft.remote_image_urls,
            mention_bindings: existing_draft.mention_bindings,
        };
        let existing_has_content =
            !self.bottom_pane.composer_is_empty() || !existing_draft.pending_pastes.is_empty();

        let startup_message = UserMessage {
            text: draft.text,
            text_elements: draft.text_elements,
            local_images: draft.local_images,
            remote_image_urls: draft.remote_image_urls,
            mention_bindings: draft.mention_bindings,
        };
        let startup_has_content = !startup_message.text.is_empty()
            || !startup_message.local_images.is_empty()
            || !startup_message.remote_image_urls.is_empty()
            || !draft.pending_pastes.is_empty();

        let cursor = if existing_has_content && startup_has_content {
            let mut used_paste_placeholders = HashSet::new();
            let (existing_message, mut pending_pastes) = remap_colliding_paste_placeholders(
                existing_message,
                existing_draft.pending_pastes,
                &mut used_paste_placeholders,
            );
            let startup_offset = existing_message.text.len().saturating_add(1);
            let preceding_text_element = startup_message
                .text_elements
                .iter()
                .enumerate()
                .take_while(|(_, element)| element.byte_range.end <= draft.cursor)
                .last()
                .map(|(index, element)| (index, element.byte_range.end));
            let (startup_message, startup_pending_pastes) = remap_colliding_paste_placeholders(
                startup_message,
                draft.pending_pastes,
                &mut used_paste_placeholders,
            );
            let cursor_adjustment = preceding_text_element.map_or(0, |(index, original_end)| {
                startup_message.text_elements[index]
                    .byte_range
                    .end
                    .saturating_sub(original_end)
            });
            pending_pastes.extend(startup_pending_pastes);
            self.restore_composer_state(Self::composer_state_from_user_message(
                merge_user_messages(vec![existing_message, startup_message]),
                pending_pastes,
            ));
            startup_offset
                .saturating_add(draft.cursor)
                .saturating_add(cursor_adjustment)
        } else if existing_has_content || !startup_has_content {
            existing_cursor
        } else {
            self.restore_composer_state(Self::composer_state_from_user_message(
                startup_message,
                draft.pending_pastes,
            ));
            draft.cursor
        };
        self.bottom_pane.set_composer_cursor(cursor);
        self.bottom_pane.restore_startup_composer_state(
            draft.last_composer_activity_at,
            draft.startup_local_history,
        );
        if startup_has_content && self.local_settings.tui.vim_mode_default {
            self.bottom_pane.enable_vim_in_insert_mode();
        }
        crate::startup_recovery::handed_off(|| self.bottom_pane.composer_recovery_snapshot());
    }

    /// Includes protected prompts deferred by streaming or the approval idle timer.
    pub(crate) fn has_pending_protected_request(&self) -> bool {
        self.bottom_pane.has_pending_approval() || self.interrupts.has_pending_prompt()
    }

    /// Transfer startup input only after protected views and required sandbox setup finish.
    pub(crate) fn restore_startup_draft_when_ready(
        &mut self,
        pending_draft: &mut Option<ComposerDraftSnapshot>,
    ) {
        if let Some(draft) = pending_draft.as_ref() {
            self.bottom_pane
                .inherit_startup_sparkle(draft.sparkle_draft);
        }
        if self.has_active_view()
            || self
                .bottom_pane
                .questions
                .as_ref()
                .is_some_and(|q| q.expanded)
            || self.has_pending_protected_request()
            || !self.bottom_pane.composer_input_enabled()
        {
            return;
        }
        #[cfg(any(target_os = "windows", test))]
        if self.windows_sandbox_host == crate::app::WindowsSandboxHost::Local
            && (self.windows_sandbox_local_server
                && self.windows_sandbox_config.requirements.is_none()
                || self.elevated_windows_sandbox_setup_required())
        {
            return;
        }
        if let Some(draft) = pending_draft.take() {
            self.restore_startup_draft(draft);
        }
    }

    pub(crate) fn set_initial_user_message_submit_suppressed(&mut self, suppressed: bool) {
        self.suppress_initial_user_message_submit = suppressed;
    }

    pub(crate) fn submit_initial_user_message_if_pending(&mut self) {
        if self.suppress_initial_user_message_submit || self.input_queue.rate_limit_recovery_pending
        {
            return;
        }
        #[cfg(any(target_os = "windows", test))]
        if self.windows_sandbox_local_server
            && self.windows_sandbox_host != crate::app::WindowsSandboxHost::Remote
            && self.windows_sandbox_config.requirements.is_none()
        {
            return;
        }
        #[cfg(any(target_os = "windows", test))]
        if self.windows_sandbox_host == crate::app::WindowsSandboxHost::Local
            && self.elevated_windows_sandbox_setup_required()
        {
            return;
        }
        #[cfg(any(target_os = "windows", test))]
        if matches!(
            self.windows_sandbox_host,
            crate::app::WindowsSandboxHost::Mixed | crate::app::WindowsSandboxHost::Unknown
        ) && self.elevated_windows_sandbox_setup_required()
        {
            if let Some(user_message) = self.initial_user_message.take() {
                self.restore_user_message_to_composer(user_message);
            }
            return;
        }
        if self.blocks_direct_input {
            if let Some(user_message) = self.initial_user_message.take() {
                self.restore_user_message_to_composer(user_message);
            }
            return;
        }
        if let Some(user_message) = self.initial_user_message.take() {
            self.submit_user_message(user_message);
        }
    }

    pub(super) fn pop_next_queued_user_message(
        &mut self,
    ) -> Option<(QueuedUserMessage, UserMessageHistoryRecord)> {
        if self.input_queue.rejected_steers_queue.is_empty() {
            self.input_queue
                .queued_user_messages
                .pop_front()
                .map(|user_message| {
                    let history_record = self
                        .input_queue
                        .queued_user_message_history_records
                        .pop_front()
                        .unwrap_or(UserMessageHistoryRecord::UserMessageText);
                    (user_message, history_record)
                })
        } else {
            // Reply envelopes must remain separate messages so other clients can read them.
            let count = self
                .input_queue
                .rejected_steers_queue
                .iter()
                .position(|message| crate::async_question_reply::parse(&message.text).is_some())
                .map(|index| index.max(/*other*/ 1))
                .unwrap_or(self.input_queue.rejected_steers_queue.len());
            let rejected_messages = self
                .input_queue
                .rejected_steers_queue
                .drain(..count)
                .collect::<Vec<_>>();
            let sources = self
                .input_queue
                .rejected_steer_sources
                .drain(..count.min(self.input_queue.rejected_steer_sources.len()))
                .collect::<Vec<_>>();
            let source = if !rejected_messages.is_empty()
                && sources.len() == rejected_messages.len()
                && sources
                    .iter()
                    .all(|source| *source == UserMessageSource::QuestionAnswer)
            {
                UserMessageSource::QuestionAnswer
            } else {
                UserMessageSource::Prompt
            };
            let mut history_records = self
                .input_queue
                .rejected_steer_history_records
                .drain(..count.min(self.input_queue.rejected_steer_history_records.len()))
                .collect::<Vec<_>>();
            history_records.resize(
                rejected_messages.len(),
                UserMessageHistoryRecord::UserMessageText,
            );
            let (message, history_record) = merge_user_messages_with_history_record(
                rejected_messages
                    .into_iter()
                    .zip(history_records)
                    .collect::<Vec<_>>(),
            );
            Some((
                QueuedUserMessage {
                    source,
                    ..QueuedUserMessage::from(message)
                },
                history_record,
            ))
        }
    }

    pub(super) fn pop_latest_queued_composer_state(&mut self) -> Option<ThreadComposerState> {
        if let Some(user_message) = self.input_queue.queued_user_messages.pop_back() {
            self.input_queue.recovered_queue &= self.input_queue.has_queued_follow_up_messages()
                || !self.input_queue.pending_steers.is_empty();
            let history_record = self
                .input_queue
                .queued_user_message_history_records
                .pop_back()
                .unwrap_or(UserMessageHistoryRecord::UserMessageText);
            let QueuedUserMessage {
                user_message,
                pending_pastes,
                ..
            } = user_message;
            Some(Self::composer_state_from_user_message(
                user_message_for_restore(user_message, &history_record),
                pending_pastes,
            ))
        } else {
            let user_message = self.input_queue.rejected_steers_queue.pop_back()?;
            self.input_queue.rejected_steer_sources.pop_back();
            self.input_queue.recovered_queue &= self.input_queue.has_queued_follow_up_messages()
                || !self.input_queue.pending_steers.is_empty();
            let history_record = self
                .input_queue
                .rejected_steer_history_records
                .pop_back()
                .unwrap_or(UserMessageHistoryRecord::UserMessageText);
            Some(Self::composer_state_from_user_message(
                user_message_for_restore(user_message, &history_record),
                Vec::new(),
            ))
        }
    }

    pub(crate) fn enqueue_rejected_steer(&mut self) -> bool {
        let Some(pending_steer) = self.input_queue.pending_steers.pop_front() else {
            tracing::warn!(
                "received active-turn-not-steerable error without a matching pending steer"
            );
            return false;
        };
        self.input_queue
            .rejected_steers_queue
            .push_back(pending_steer.user_message);
        self.input_queue
            .rejected_steer_sources
            .push_back(pending_steer.source);
        self.input_queue
            .rejected_steer_history_records
            .push_back(pending_steer.history_record);
        self.refresh_pending_input_preview();
        true
    }

    /// Handle a turn aborted due to user interrupt (Esc), budget exhaustion,
    /// or review completion.
    /// When there are queued user messages, restore them into the composer
    /// separated by newlines rather than auto-submitting the next one.
    pub(super) fn on_interrupted_turn(&mut self, reason: TurnAbortReason) {
        self.requeue_image_submission();
        // Finalize, log a gentle prompt, and clear running state.
        self.finalize_turn();
        let send_pending_steers_immediately =
            self.input_queue.submit_pending_steers_after_interrupt;
        self.input_queue.submit_pending_steers_after_interrupt = false;
        if self.interrupted_turn_notice_mode != InterruptedTurnNoticeMode::Suppress {
            if send_pending_steers_immediately {
                self.add_to_history(history_cell::new_info_event(
                    "Model interrupted to submit steer instructions.".to_owned(),
                    /*hint*/ None,
                ));
            } else {
                self.add_to_history(match reason {
                    TurnAbortReason::BudgetLimited => history_cell::new_error_event(
                        "Goal budget reached - the turn was stopped.".to_string(),
                    ),
                    TurnAbortReason::Interrupted => PlainHistoryCell::new(vec![
                        Line::from(
                            "■ Conversation interrupted - use /feedback if something went wrong",
                        )
                        .style(crate::style::secondary_text_style()),
                    ]),
                });
            }
        }

        // The server has already discarded pending input by the time the
        // interrupted turn reaches the UI, so any unacknowledged steers still
        // tracked here must be restored locally instead of waiting for a later commit.
        if send_pending_steers_immediately {
            let pending_steers = self
                .input_queue
                .pending_steers
                .drain(..)
                .collect::<Vec<_>>();
            if pending_steers
                .iter()
                .any(|pending| pending.source == UserMessageSource::QuestionAnswer)
            {
                // Keep answers intact when an interrupt retries uncommitted input.
                for pending in pending_steers {
                    self.input_queue
                        .rejected_steers_queue
                        .push_back(pending.user_message);
                    self.input_queue
                        .rejected_steer_sources
                        .push_back(pending.source);
                    self.input_queue
                        .rejected_steer_history_records
                        .push_back(pending.history_record);
                }
                if let Some((message, history_record)) = self.pop_next_queued_user_message() {
                    let source = message.source;
                    self.submit_user_message_with_history_and_shell_escape_policy(
                        message.into_user_message(),
                        history_record,
                        ShellEscapePolicy::Allow,
                        source,
                    );
                }
            } else if !pending_steers.is_empty() {
                let (user_message, history_record) = merge_user_messages_with_history_record(
                    pending_steers
                        .into_iter()
                        .map(|pending| (pending.user_message, pending.history_record))
                        .collect(),
                );
                self.submit_user_message_with_history_and_shell_escape_policy(
                    user_message,
                    history_record,
                    ShellEscapePolicy::Allow,
                    UserMessageSource::Prompt,
                );
            } else if let Some(combined) = self.drain_pending_messages_for_restore() {
                self.restore_composer_state(combined);
            }
        } else if let Some(combined) = self.drain_pending_messages_for_restore() {
            self.restore_composer_state(combined);
        }
        self.refresh_pending_input_preview();
        self.request_redraw();
    }

    /// Merge pending steers, queued drafts, and the current composer state into a single message.
    ///
    /// Each pending message numbers attachments from `[Image #1]` relative to its own remote
    /// images. When we concatenate multiple messages after interrupt, we must renumber local-image
    /// placeholders in a stable order and rebase text element byte ranges so the restored composer
    /// state stays aligned with the merged attachment list. Returns `None` when there is nothing to
    /// restore.
    fn drain_pending_messages_for_restore(&mut self) -> Option<ThreadComposerState> {
        if self.input_queue.pending_steers.is_empty() && !self.has_queued_follow_up_messages() {
            return None;
        }

        let composer = self.bottom_pane.composer_draft_snapshot();
        let composer_pending_pastes = composer.pending_pastes;
        let existing_message = UserMessage {
            text: composer.text,
            text_elements: composer.text_elements,
            local_images: composer.local_images,
            remote_image_urls: composer.remote_image_urls,
            mention_bindings: composer.mention_bindings,
        };

        let rejected_messages = self
            .input_queue
            .rejected_steers_queue
            .drain(..)
            .collect::<Vec<_>>();
        self.input_queue.rejected_steer_sources.clear();
        let mut rejected_history_records = self
            .input_queue
            .rejected_steer_history_records
            .drain(..)
            .collect::<Vec<_>>();
        rejected_history_records.resize(
            rejected_messages.len(),
            UserMessageHistoryRecord::UserMessageText,
        );
        let mut to_merge: Vec<UserMessage> = rejected_messages
            .into_iter()
            .zip(rejected_history_records.iter())
            .map(|(message, history_record)| user_message_for_restore(message, history_record))
            .collect();
        to_merge.extend(
            self.input_queue
                .pending_steers
                .drain(..)
                .map(|steer| user_message_for_restore(steer.user_message, &steer.history_record)),
        );
        let queued_messages = self
            .input_queue
            .queued_user_messages
            .drain(..)
            .collect::<Vec<_>>();
        let mut queued_history_records = self
            .input_queue
            .queued_user_message_history_records
            .drain(..)
            .collect::<Vec<_>>();
        queued_history_records.resize(
            queued_messages.len(),
            UserMessageHistoryRecord::UserMessageText,
        );
        let mut pending_pastes = Vec::new();
        let mut used_paste_placeholders = HashSet::new();
        for (message, history_record) in queued_messages
            .into_iter()
            .zip(queued_history_records.iter())
        {
            let (message, message_pastes) = remap_colliding_paste_placeholders(
                user_message_for_restore(message.user_message, history_record),
                message.pending_pastes,
                &mut used_paste_placeholders,
            );
            pending_pastes.extend(message_pastes);
            to_merge.push(message);
        }
        let has_existing_message = !existing_message.text.is_empty()
            || !existing_message.local_images.is_empty()
            || !existing_message.remote_image_urls.is_empty();
        if has_existing_message {
            let (existing_message, composer_pending_pastes) = remap_colliding_paste_placeholders(
                existing_message,
                composer_pending_pastes,
                &mut used_paste_placeholders,
            );
            to_merge.push(existing_message);
            pending_pastes.extend(composer_pending_pastes);
        }

        Some(Self::composer_state_from_user_message(
            merge_user_messages(to_merge),
            pending_pastes,
        ))
    }

    pub(crate) fn restore_user_message_to_composer(&mut self, user_message: UserMessage) {
        let draft = self.bottom_pane.composer_draft_snapshot();
        let pending_pastes = draft.pending_pastes;
        let draft_message = UserMessage {
            text: draft.text,
            text_elements: draft.text_elements,
            local_images: draft.local_images,
            remote_image_urls: draft.remote_image_urls,
            mention_bindings: draft.mention_bindings,
        };
        let mut messages = vec![user_message];
        if !draft_message.text.is_empty()
            || !draft_message.local_images.is_empty()
            || !draft_message.remote_image_urls.is_empty()
        {
            messages.push(draft_message);
        }
        self.restore_composer_state(Self::composer_state_from_user_message(
            merge_user_messages(messages),
            pending_pastes,
        ));
    }

    pub(super) fn restore_composer_state(&mut self, composer: ThreadComposerState) {
        let ThreadComposerState {
            text,
            local_images,
            remote_image_urls,
            text_elements,
            mention_bindings,
            pending_pastes,
        } = composer;
        let local_image_paths = local_images.into_iter().map(|img| img.path).collect();
        self.set_remote_image_urls(remote_image_urls);
        self.bottom_pane.set_composer_text_with_mention_bindings(
            text,
            text_elements,
            local_image_paths,
            mention_bindings,
        );
        self.bottom_pane.set_composer_pending_pastes(pending_pastes);
    }

    fn composer_state_from_user_message(
        user_message: UserMessage,
        pending_pastes: Vec<(String, String)>,
    ) -> ThreadComposerState {
        let UserMessage {
            text,
            local_images,
            remote_image_urls,
            text_elements,
            mention_bindings,
        } = user_message;
        ThreadComposerState {
            text,
            local_images,
            remote_image_urls,
            text_elements,
            mention_bindings,
            pending_pastes,
        }
    }

    pub(crate) fn capture_thread_input_state(&mut self) -> Option<ThreadInputState> {
        self.cancel_image_submission();
        let draft = self.bottom_pane.composer_draft_snapshot();
        let composer = ThreadComposerState {
            text: draft.text,
            text_elements: draft.text_elements,
            local_images: draft.local_images,
            remote_image_urls: draft.remote_image_urls,
            mention_bindings: draft.mention_bindings,
            pending_pastes: draft.pending_pastes,
        };
        Some(ThreadInputState {
            questions: self
                .bottom_pane
                .questions
                .as_deref_mut()
                .map(crate::bottom_pane::AsyncQuestions::capture),
            pending_thread_settings: None,
            composer: composer.has_content().then_some(composer),
            safety_buffering_prompt: self.safety_buffering_prompt.clone(),
            safety_buffering_source: self.safety_buffering_source,
            pending_steers: self.input_queue.pending_steers.clone(),
            rejected_steers_queue: self.input_queue.rejected_steers_queue.clone(),
            rejected_steer_sources: self.input_queue.rejected_steer_sources.clone(),
            rejected_steer_history_records: self.input_queue.rejected_steer_history_records.clone(),
            queued_user_messages: self.input_queue.queued_user_messages.clone(),
            queued_user_message_history_records: self
                .input_queue
                .queued_user_message_history_records
                .clone(),
            recovered_queue: self.input_queue.recovered_queue,
            user_turn_pending_start: self.input_queue.user_turn_pending_start,
            submit_pending_steers_after_interrupt: self
                .input_queue
                .submit_pending_steers_after_interrupt,
            current_collaboration_mode: self.current_collaboration_mode.clone(),
            active_collaboration_mask: self.active_collaboration_mask.clone(),
            plan_mode_reasoning_effort: self.config.plan_mode_reasoning_effort.clone(),
            task_running: self.bottom_pane.is_task_running(),
            agent_turn_running: self.turn_lifecycle.agent_turn_running,
        })
    }

    pub(crate) fn restore_thread_input_state(
        &mut self,
        input_state: Option<ThreadInputState>,
        restore_mode: ThreadInputStateRestoreMode,
    ) {
        let preserve_in_flight_turn = restore_mode.preserve_in_flight_turn;
        let restored_task_running =
            preserve_in_flight_turn && input_state.as_ref().is_some_and(|state| state.task_running);
        if let Some(input_state) = input_state {
            self.bottom_pane.restore_questions(input_state.questions);
            self.input_queue.recovered_queue = input_state.recovered_queue;
            self.current_collaboration_mode = input_state.current_collaboration_mode;
            self.active_collaboration_mask = input_state.active_collaboration_mask;
            self.config.plan_mode_reasoning_effort = input_state.plan_mode_reasoning_effort;
            self.safety_buffering_prompt = input_state.safety_buffering_prompt;
            self.safety_buffering_source = input_state.safety_buffering_source;
            self.turn_lifecycle.restore_running(
                preserve_in_flight_turn && input_state.agent_turn_running,
                Instant::now(),
            );
            self.input_queue.user_turn_pending_start =
                preserve_in_flight_turn && input_state.user_turn_pending_start;
            self.input_queue.submit_pending_steers_after_interrupt =
                preserve_in_flight_turn && input_state.submit_pending_steers_after_interrupt;
            self.update_collaboration_mode_indicator();
            self.refresh_model_dependent_surfaces();
            self.restore_composer_state(input_state.composer.unwrap_or_default());
            let pending_steers = input_state.pending_steers;
            let mut queued_user_messages = input_state.queued_user_messages;
            let mut queued_user_message_history_records =
                input_state.queued_user_message_history_records;
            if preserve_in_flight_turn {
                self.input_queue.pending_steers = pending_steers;
            } else {
                self.input_queue.pending_steers.clear();
                for pending in pending_steers.into_iter().rev() {
                    queued_user_messages.push_front(QueuedUserMessage {
                        source: pending.source,
                        ..QueuedUserMessage::from(pending.user_message)
                    });
                    queued_user_message_history_records.push_front(pending.history_record);
                }
            }
            self.input_queue.rejected_steers_queue = input_state.rejected_steers_queue;
            self.input_queue.rejected_steer_sources = input_state.rejected_steer_sources;
            self.input_queue.rejected_steer_sources.resize(
                self.input_queue.rejected_steers_queue.len(),
                UserMessageSource::Prompt,
            );
            self.input_queue.rejected_steer_history_records =
                input_state.rejected_steer_history_records;
            self.input_queue.rejected_steer_history_records.resize(
                self.input_queue.rejected_steers_queue.len(),
                UserMessageHistoryRecord::UserMessageText,
            );
            self.input_queue.queued_user_messages = queued_user_messages;
            self.input_queue.queued_user_message_history_records =
                queued_user_message_history_records;
            self.input_queue.queued_user_message_history_records.resize(
                self.input_queue.queued_user_messages.len(),
                UserMessageHistoryRecord::UserMessageText,
            );
        } else {
            self.turn_lifecycle
                .restore_running(/*running*/ false, Instant::now());
            self.safety_buffering_prompt = None;
            self.safety_buffering_source = UserMessageSource::Prompt;
            self.input_queue.clear();
            self.restore_composer_state(Default::default());
        }
        self.input_queue.recovered_queue &= self.input_queue.has_queued_follow_up_messages()
            || !self.input_queue.pending_steers.is_empty();
        let effort = self.effective_reasoning_effort();
        self.bottom_pane
            .set_active_reasoning_effort_baseline(effort.as_ref());
        self.turn_lifecycle
            .restore_running(self.turn_lifecycle.agent_turn_running, Instant::now());
        self.update_task_running_state();
        if restored_task_running && !self.bottom_pane.is_task_running() {
            self.bottom_pane.set_task_running(/*running*/ true);
            self.refresh_status_surfaces();
        }
        self.refresh_pending_input_preview();
        self.request_redraw();
    }

    pub(crate) fn set_queue_autosend_suppressed(&mut self, suppressed: bool) {
        self.input_queue.suppress_queue_autosend = suppressed;
    }
}