// *****************************************************************************
// Copyright (C) 2024 EclipseSource GmbH.
//
// This program and the accompanying materials are made available under the
// terms of the Eclipse Public License v. 2.0 which is available at
// http://www.eclipse.org/legal/epl-2.0.
//
// This Source Code may also be made available under the following Secondary
// Licenses when the conditions for such availability set forth in the Eclipse
// Public License v. 2.0 are satisfied: GNU General Public License, version 2
// with the GNU Classpath Exception which is available at
// https://www.gnu.org/software/classpath/license.html.
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************
/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
// Partially copied from https://github.com/microsoft/vscode/blob/a2cab7255c0df424027be05d58e1b7b941f4ea60/src/vs/workbench/contrib/chat/common/chatModel.ts

import {
    AIVariableResolutionRequest,
    CompactionMessage,
    CompactionSettings,
    GenericCapabilitySelections,
    LanguageModelMessage,
    ReasoningSettings,
    ResolvedAIContextVariable,
    ResolvedAIVariable,
    ServerToolUseMessage,
    TextMessage,
    ThinkingMessage,
    ToolCallResult,
    ToolRequest,
    ToolResultMessage,
    ToolUseMessage
} from '@theia/ai-core';
import { ArrayUtils, CancellationToken, CancellationTokenSource, Command, Disposable, DisposableCollection, Emitter, Event, generateUuid, URI } from '@theia/core';
import { MarkdownString, MarkdownStringImpl } from '@theia/core/lib/common/markdown-rendering';
import { Position } from '@theia/core/shared/vscode-languageserver-protocol';
import { ChangeSet, ChangeSetElement, ChangeSetImpl, ChatUpdateChangeSetEvent } from './change-set';
import { ChatAgentLocation } from './chat-agents';
import {
    SerializedChatModel,
    SerializableChatRequestData,
    SerializableChatResponseContentData,
    SerializableChatResponseData,
    SerializableHierarchy,
    SerializableHierarchyBranch,
    SerializableHierarchyBranchItem,
    SerializableChangeSetElement,
    SerializableParsedRequest,
    SerializableParsedRequestPart
} from './chat-model-serialization';
import {
    ParsedChatRequest,
    ParsedChatRequestTextPart,
    ParsedChatRequestVariablePart,
    ParsedChatRequestFunctionPart,
    ParsedChatRequestAgentPart,
    ParsedChatRequestPart
} from './parsed-chat-request';
import debounce = require('@theia/core/shared/lodash.debounce');
export { ChangeSet, ChangeSetElement, ChangeSetImpl };

/**********************
 * INTERFACES AND TYPE GUARDS
 **********************/

export type ChatChangeEvent =
    | ChatAddRequestEvent
    | ChatAddResponseEvent
    | ChatAddVariableEvent
    | ChatRemoveVariableEvent
    | ChatSetVariablesEvent
    | ChatRemoveRequestEvent
    | ChatSuggestionsChangedEvent
    | ChatUpdateChangeSetEvent
    | ChatEditRequestEvent
    | ChatEditCancelEvent
    | ChatEditSubmitEvent
    | ChatResponseChangedEvent
    | ChatChangeHierarchyBranchEvent
    | ChatInteractionNeededEvent
    | ChatSessionStatusChangedEvent
    | ChatSettingsChangedEvent;

export interface ChatAddRequestEvent {
    kind: 'addRequest';
    request: ChatRequestModel;
}

export interface ChatEditRequestEvent {
    kind: 'enableEdit';
    request: EditableChatRequestModel;
    branch: ChatHierarchyBranch<ChatRequestModel>;
}

export interface ChatEditCancelEvent {
    kind: 'cancelEdit';
    request: EditableChatRequestModel;
    branch: ChatHierarchyBranch<ChatRequestModel>;
}

export interface ChatEditSubmitEvent {
    kind: 'submitEdit';
    request: EditableChatRequestModel;
    branch: ChatHierarchyBranch<ChatRequestModel>;
    newRequest: ChatRequest;
}

export interface ChatChangeHierarchyBranchEvent {
    kind: 'changeHierarchyBranch';
    branch: ChatHierarchyBranch<ChatRequestModel>;
}

export interface ChatAddResponseEvent {
    kind: 'addResponse';
    response: ChatResponseModel;
}

export interface ChatAddVariableEvent {
    kind: 'addVariable';
}

export interface ChatRemoveVariableEvent {
    kind: 'removeVariable';
}

export interface ChatSetVariablesEvent {
    kind: 'setVariables';
}

export interface ChatSuggestionsChangedEvent {
    kind: 'suggestionsChanged';
    suggestions: ChatSuggestion[];
}

export interface ChatResponseChangedEvent {
    kind: 'responseChanged';
}

export interface ChatInteractionNeededEvent {
    kind: 'interactionNeeded';
    contentPart: InteractiveContent & ChatResponseContent;
}

export interface ChatSessionStatusChangedEvent {
    kind: 'statusChanged';
    status: ChatSessionStatus;
}

export interface ChatSettingsChangedEvent {
    kind: 'settingsChanged';
    settings: ChatSessionSettings;
}

export namespace ChatChangeEvent {
    export function isChangeSetEvent(event: ChatChangeEvent): event is ChatUpdateChangeSetEvent {
        return event.kind === 'updateChangeSet';
    }
    export function isInteractionNeededEvent(event: ChatChangeEvent): event is ChatInteractionNeededEvent {
        return event.kind === 'interactionNeeded';
    }
    export function isStatusChangedEvent(event: ChatChangeEvent): event is ChatSessionStatusChangedEvent {
        return event.kind === 'statusChanged';
    }
}

export type ChatRequestRemovalReason = 'removal' | 'resend' | 'adoption';

export interface ChatRemoveRequestEvent {
    kind: 'removeRequest';
    requestId: string;
    responseId?: string;
    reason: ChatRequestRemovalReason;
}

/**
 * A model that contains information about a chat request that may branch off.
 *
 * The hierarchy of requests is represented by a tree structure.
 * - The root of the tree is the initial request
 * - Within each branch, the requests are stored in a list. Those requests are the alternatives to the original request.
 *   Each of those items can have a next branch, which is the next request in the hierarchy.
 */
export interface ChatRequestHierarchy<TRequest extends ChatRequestModel = ChatRequestModel> extends Disposable {
    readonly branch: ChatHierarchyBranch<TRequest>

    onDidChange: Event<ChangeActiveBranchEvent<TRequest>>;

    append(request: TRequest): ChatHierarchyBranch<TRequest>;
    activeRequests(): TRequest[];
    activeBranches(): ChatHierarchyBranch<TRequest>[];
    findRequest(requestId: string): TRequest | undefined;
    findBranch(requestId: string): ChatHierarchyBranch<TRequest> | undefined;

    notifyChange(event: ChangeActiveBranchEvent<TRequest>): void;
    toSerializable(): SerializableHierarchy;
}

export interface ChangeActiveBranchEvent<TRequest extends ChatRequestModel = ChatRequestModel> {
    branch: ChatHierarchyBranch<TRequest>,
    item: ChatHierarchyBranchItem<TRequest>
}

/**
 * A branch of the chat request hierarchy.
 * It contains a list of items, each representing a request.
 * Those items can have a next branch, which is the next request in the hierarchy.
 */
export interface ChatHierarchyBranch<TRequest extends ChatRequestModel = ChatRequestModel> extends Disposable {
    readonly id: string;
    readonly hierarchy: ChatRequestHierarchy<TRequest>;
    readonly previous?: ChatHierarchyBranch<TRequest>;
    readonly items: ChatHierarchyBranchItem<TRequest>[];
    readonly activeBranchIndex: number;

    next(): ChatHierarchyBranch<TRequest> | undefined;
    get(): TRequest;
    add(request: TRequest): void;
    remove(request: TRequest | string): void;
    /**
     * Create a new branch by inserting it as the next branch of the active item.
     */
    continue(request: TRequest): ChatHierarchyBranch<TRequest>;

    enable(request: TRequest): ChatHierarchyBranchItem<TRequest>;
    enablePrevious(): ChatHierarchyBranchItem<TRequest>;
    enableNext(): ChatHierarchyBranchItem<TRequest>;

    succeedingBranches(): ChatHierarchyBranch<TRequest>[];
}

export interface ChatHierarchyBranchItem<TRequest extends ChatRequestModel = ChatRequestModel> {
    readonly element: TRequest;
    readonly next?: ChatHierarchyBranch<TRequest>;
}

export interface CommonChatSessionSettings {
    /**
     * Language model (or alias) id to use for this session only, overriding the agent's default.
     * New sessions start without an override and use the agent's configured model. Cleared to
     * revert to the default.
     */
    modelId?: string;
    /** Reasoning configuration for this session; applied to reasoning-capable models. */
    reasoning?: ReasoningSettings;
    /** Per-session tool confirmation timeout in seconds. Overrides the global preference when set. */
    confirmationTimeout?: number;
    /** Per-session server-side compaction settings; set values win over the per-provider and global settings. */
    compaction?: CompactionSettings;
}

export interface ChatSessionSettings {
    /**
     * Theia-specific common settings processed by Theia.
     * This key is excluded when passing settings to LLM providers.
     */
    commonSettings?: CommonChatSessionSettings;

    /**
     * Allow arbitrary provider-specific settings.
     * These correspond to the "Advanced Settings (JSON)" in the UI
     * and are passed directly to the LLM.
     */
    [key: string]: unknown;
}

/**
 * Aggregated, observable status of a chat session, derived from the state of the session's last request.
 */
export type ChatSessionStatus =
    /** No request is in progress. The last request, if any, completed successfully or was canceled. */
    | 'idle'
    /** A request is in progress. */
    | 'running'
    /** A tool call requires user confirmation before it can be executed. */
    | 'awaitingApproval'
    /** An approved or confirmation-free tool call is still executing. */
    | 'awaitingToolCall'
    /** The agent is waiting for user input, e.g. an answer to a structured question. */
    | 'awaitingInput'
    /** The last request ended in an error. */
    | 'failed';

export namespace ChatSessionStatus {
    /**
     * All {@link ChatSessionStatus} values, e.g. for schema declarations. Derived from an
     * exhaustive record so that adding a status without listing it here fails to compile.
     */
    export const VALUES: readonly ChatSessionStatus[] = Object.keys({
        idle: true,
        running: true,
        awaitingApproval: true,
        awaitingToolCall: true,
        awaitingInput: true,
        failed: true
    } satisfies Record<ChatSessionStatus, true>) as ChatSessionStatus[];

    /**
     * Whether a request is in progress in this status, including the states waiting on the user or a tool.
     */
    export function isInProgress(status: ChatSessionStatus): boolean {
        return status !== 'idle' && status !== 'failed';
    }

    /**
     * Whether the session is blocked on the user, i.e. a tool approval or another input is required to proceed.
     */
    export function requiresUserAction(status: ChatSessionStatus): boolean {
        return status === 'awaitingApproval' || status === 'awaitingInput';
    }

    /**
     * Derives the session status from the state of the given request (usually the session's last request).
     */
    export function fromRequest(request: ChatRequestModel | undefined): ChatSessionStatus {
        if (!request) {
            return 'idle';
        }
        const response = request.response;
        if (response.isError) {
            return 'failed';
        }
        if (!ChatRequestModel.isInProgress(request)) {
            return 'idle';
        }
        const content = response.response.content;
        if (content.some(part => ToolCallChatResponseContent.is(part) && part.isAwaitingUserConfirmation)) {
            return 'awaitingApproval';
        }
        // Unresolved non-tool-call interactive parts are structured questions (tool calls are
        // "unresolved" while merely executing). The waiting-for-input flag additionally covers
        // inputs without a dedicated content part, e.g. the user-interaction tool's wizard.
        if (response.isWaitingForInput || content.some(part => !ToolCallChatResponseContent.is(part) && InteractiveContent.is(part) && !part.isResolved)) {
            return 'awaitingInput';
        }
        if (content.some(part => ToolCallChatResponseContent.is(part) && !part.finished)) {
            return 'awaitingToolCall';
        }
        return 'running';
    }
}

export interface ChatModel {
    readonly onDidChange: Event<ChatChangeEvent>;
    readonly id: string;
    readonly location: ChatAgentLocation;
    readonly context: ChatContextManager;
    readonly suggestions: readonly ChatSuggestion[];
    readonly settings?: ChatSessionSettings;
    readonly changeSet: ChangeSet;
    /**
     * Aggregated status of this session, derived from the state of its last request.
     * Changes are announced via {@link onDidChange} with a {@link ChatSessionStatusChangedEvent}.
     */
    readonly status: ChatSessionStatus;
    /** ID of the root session in the delegation chain. For delegated sessions, this points to the topmost session where task contexts are stored. */
    rootSessionId?: string;
    /** ID of the immediate parent session that delegated this one. Undefined for top-level sessions. */
    parentSessionId?: string;
    getRequests(): ChatRequestModel[];
    getBranches(): ChatHierarchyBranch<ChatRequestModel>[];
    isEmpty(): boolean;
    toSerializable(): SerializedChatModel;
}

export interface ChatSuggestionCallback {
    kind: 'callback',
    callback: () => unknown;
    content: string | MarkdownString;
}
export namespace ChatSuggestionCallback {
    export function is(candidate: ChatSuggestion): candidate is ChatSuggestionCallback {
        return typeof candidate === 'object' && 'callback' in candidate;
    }
    export function containsCallbackLink(candidate: ChatSuggestion): candidate is ChatSuggestionCallback {
        if (!is(candidate)) { return false; }
        const text = typeof candidate.content === 'string' ? candidate.content : candidate.content.value;
        return text.includes('](_callback)');
    }
}

export type ChatSuggestion = | string | MarkdownString | ChatSuggestionCallback;

export interface ChatContextManager {
    onDidChange: Event<ChatAddVariableEvent | ChatRemoveVariableEvent | ChatSetVariablesEvent>;
    getVariables(): readonly AIVariableResolutionRequest[]
    addVariables(...variables: AIVariableResolutionRequest[]): void;
    deleteVariables(...indices: number[]): void;
    clear(): void;
}

export interface ChangeSetDecoration {
    readonly priority?: number;
    readonly additionalInfoSuffixIcon?: string[];
}

export interface ChatRequest {
    readonly text: string;
    readonly displayText?: string;
    /**
     * If the request has been triggered in the context of
     * an existing request, this id will be set to the id of the
     * referenced request.
     */
    readonly referencedRequestId?: string;
    readonly variables?: readonly AIVariableResolutionRequest[];
    readonly modeId?: string;
    /**
     * Capability overrides for this request.
     * Maps capability fragment IDs to enabled/disabled state.
     * Only includes capabilities that differ from their default value.
     */
    readonly capabilityOverrides?: Record<string, boolean>;

    /**
     * Generic capability selections for this request.
     * Contains user-selected skills, functions, MCP tools, etc.
     * from the capabilities panel dropdowns.
     */
    readonly genericCapabilitySelections?: GenericCapabilitySelections;

    /**
     * Server tool selections for this request, keyed by model vendor. Only the entry matching
     * the actually selected model's vendor is applied when sending the request.
     */
    readonly serverToolSelections?: Record<string, string[]>;
}

export interface ChatContext {
    variables: ResolvedAIContextVariable[];
}

export interface ChatRequestModel {
    readonly id: string;
    readonly session: ChatModel;
    readonly request: ChatRequest;
    readonly response: ChatResponseModel;
    readonly message: ParsedChatRequest;
    readonly context: ChatContext;
    readonly agentId?: string;
    readonly data?: { [key: string]: unknown };
    toSerializable(): SerializableChatRequestData;
}

export namespace ChatRequestModel {
    export function is(request: unknown): request is ChatRequestModel {
        return !!(
            request &&
            typeof request === 'object' &&
            'id' in request &&
            typeof (request as { id: unknown }).id === 'string' &&
            'session' in request &&
            'request' in request &&
            'response' in request &&
            'message' in request
        );
    }
    export function isInProgress(request: ChatRequestModel | undefined): boolean {
        if (!request) {
            return false;
        }
        const response = request.response;
        return !(
            response.isComplete ||
            response.isCanceled ||
            response.isError
        );
    }
}

export interface EditableChatRequestModel extends ChatRequestModel {
    readonly isEditing: boolean;
    editContextManager: ChatContextManagerImpl;
    enableEdit(): void;
    cancelEdit(): void;
    submitEdit(newRequest: ChatRequest): void;
}

export namespace EditableChatRequestModel {
    export function is(request: unknown): request is EditableChatRequestModel {
        return !!(
            ChatRequestModel.is(request) &&
            'enableEdit' in request &&
            'cancelEdit' in request &&
            'submitEdit' in request
        );
    }

    export function isEditing(request: unknown): request is EditableChatRequestModel {
        return is(request) && request.isEditing;
    }
}

export interface ChatProgressMessage {
    kind: 'progressMessage';
    id: string;
    status: 'inProgress' | 'completed' | 'failed';
    show: 'untilFirstContent' | 'whileIncomplete' | 'forever';
    content: string;
}

/**
 * Interface for ChatResponseContent parts that require user interaction.
 * Content parts that implement this interface can be tracked by the delegation
 * renderer without content-type-specific checks.
 */
export interface InteractiveContent {
    /** Stable identifier for deduplication in pending interaction tracking. */
    readonly interactionId: string | undefined;
    /** Whether the interaction has been resolved (e.g., confirmed/denied, option selected). */
    readonly isResolved: boolean;
    /** Resolves when the interaction is resolved. Used for cleanup in delegation chains. */
    readonly whenResolved: Promise<void>;
    /**
     * Whether the interaction currently requires user action. Unlike {@link isResolved},
     * this reflects the momentary state: e.g. a tool call whose confirmation was granted
     * is not yet resolved (the tool is still executing) but no longer awaits interaction.
     * When `undefined`, consumers should fall back to `!isResolved`.
     */
    readonly isAwaitingInteraction?: boolean;
}

export namespace InteractiveContent {
    export function is(content: unknown): content is InteractiveContent {
        return typeof content === 'object' && !!content
            && 'interactionId' in content
            && 'isResolved' in content
            && 'whenResolved' in content;
    }
}

export interface ChatResponseContent {
    kind: string;
    /**
     * Represents the content as a string. Returns `undefined` if the content
     * is purely informational and/or visual and should not be included in the overall
     * representation of the response.
     */
    asString?(): string | undefined;
    asDisplayString?(): string | undefined;
    merge?(nextChatResponseContent: ChatResponseContent): boolean;
    toLanguageModelMessage?(): LanguageModelMessage | LanguageModelMessage[];
    toSerializable?(): SerializableChatResponseContentData;
}

export namespace ChatResponseContent {
    export function is(obj: unknown): obj is ChatResponseContent {
        return !!(
            obj &&
            typeof obj === 'object' &&
            'kind' in obj &&
            typeof (obj as { kind: unknown }).kind === 'string'
        );
    }
    export function hasAsString(
        obj: ChatResponseContent
    ): obj is Required<Pick<ChatResponseContent, 'asString'>> & ChatResponseContent {
        return typeof obj.asString === 'function';
    }
    export function hasDisplayString(
        obj: ChatResponseContent
    ): obj is Required<Pick<ChatResponseContent, 'asDisplayString'>> & ChatResponseContent {
        return typeof obj.asDisplayString === 'function';
    }
    export function hasMerge(
        obj: ChatResponseContent
    ): obj is Required<Pick<ChatResponseContent, 'merge'>> & ChatResponseContent {
        return typeof obj.merge === 'function';
    }
    export function hasToLanguageModelMessage(
        obj: ChatResponseContent
    ): obj is Required<Pick<ChatResponseContent, 'toLanguageModelMessage'>> & ChatResponseContent {
        return typeof obj.toLanguageModelMessage === 'function';
    }
}

/**
 * Data interfaces for chat response content serialization.
 * These define the structure of the data property in SerializableChatResponseContentData.
 */

export interface TextContentData {
    content: string;
}

export interface ThinkingContentData {
    content: string;
    signature: string;
}

export interface CompactionContentData {
    provider: string;
    data: unknown;
    summary?: string;
}

export interface MarkdownContentData {
    content: string;
}

export interface InformationalContentData {
    content: string;
}

export interface CodeContentData {
    code: string;
    language?: string;
    location?: Location;
}

export interface ToolCallContentData {
    id?: string;
    name?: string;
    arguments?: string;
    result?: ToolCallResult;
    data?: Record<string, string>;
}

export interface ServerToolCallContentData {
    id?: string;
    name?: string;
    arguments?: string;
    result?: ToolCallResult;
    data?: Record<string, string>;
}

export interface CommandContentData {
    commandId?: string;
    commandLabel?: string;
    arguments?: unknown[];
}

export interface HorizontalLayoutContentData {
    content: SerializableChatResponseContentData[];
}

export interface ProgressContentData {
    message: string;
}

export interface ErrorContentData {
    message: string;
    stack?: string;
}

/**
 * Restored questions display the question, options, and any previously selected answer,
 * but do not allow new selections.
 */
export interface QuestionContentData {
    question: string;
    header?: string;
    options: { text: string; value?: string; description?: string }[];
    multiSelect?: boolean;
    selectedOption?: { text: string; value?: string };
    selectedOptions?: { text: string; value?: string }[];
}

export interface TextChatResponseContent
    extends Required<ChatResponseContent> {
    kind: 'text';
    content: string;
}

export interface ErrorChatResponseContent extends ChatResponseContent {
    kind: 'error';
    error: Error;
}

export interface MarkdownChatResponseContent
    extends Required<ChatResponseContent> {
    kind: 'markdownContent';
    content: MarkdownString;
}

export interface CodeChatResponseContent
    extends ChatResponseContent {
    kind: 'code';
    code: string;
    language?: string;
    location?: Location;
}

export interface HorizontalLayoutChatResponseContent extends ChatResponseContent {
    kind: 'horizontal';
    content: ChatResponseContent[];
}

export interface ToolCallChatResponseContent extends Required<ChatResponseContent>, InteractiveContent {
    kind: 'toolCall';
    id?: string;
    name?: string;
    arguments?: string;
    finished: boolean;
    result?: ToolCallResult;
    /** Timeout in seconds for confirmation dialogs. 0 means no timeout. */
    confirmationTimeout?: number;
    confirmed: Promise<boolean>;
    /** Resolves when the tool call requires user confirmation (show Allow/Deny UI). */
    needsUserConfirmation: Promise<void>;
    /**
     * Whether the tool call is currently awaiting the user's confirmation decision, i.e.
     * {@link requestUserConfirmation} was called and the user has neither confirmed nor denied it yet.
     */
    readonly isAwaitingUserConfirmation: boolean;
    whenFinished: Promise<void>;
    /**
     * Provider-specific metadata about the tool call that the language model needs back on
     * subsequent turns (e.g. Google's `thoughtSignature`, cache pointers). The full record is
     * passed to the model via {@link toLanguageModelMessage} as `ToolUseMessage.data`.
     */
    data?: Record<string, string>;
    /**
     * Fires when the serialized form of this tool call changes outside of the agent's stream
     * (e.g. after a renderer persists intermediate state via {@link updateResult}). The
     * parent {@link ChatResponse} forwards this event so chat-session auto-save picks up
     * the change.
     */
    readonly onDidChange: Event<void>;
    confirm(): void;
    deny(reason?: string): void;
    cancelConfirmation(reason?: unknown): void;
    /** Signal that this tool call needs user confirmation. Resolves the needsUserConfirmation promise. */
    requestUserConfirmation(): void;
    /**
     * Whether the tool execution is currently blocked waiting for user input provided
     * through the tool's own UI (e.g. an interactive wizard). Cleared when the tool
     * call finishes or {@link userInputHandled} is called.
     */
    readonly isAwaitingUserInput: boolean;
    /** Signal that the tool execution is blocked waiting for user input. */
    requestUserInput(): void;
    /** Signal that the tool execution is no longer waiting for user input. */
    userInputHandled(): void;
    /**
     * Update the tool call's result without marking it finished. Use this to persist
     * intermediate state for long-running tools (e.g. the user-interaction wizard) so
     * progress survives chat-session reloads. On restore the tool call is always marked
     * finished (no live handler exists anymore), so a partial result must be
     * self-describing; consumers should not rely on `finished` to tell partial from final.
     */
    updateResult(result: ToolCallResult): void;
    /**
     * Mark the tool call as completed with the given result.
     *
     * This is used to update the UI immediately when a tool finishes execution,
     * without waiting for all parallel tool calls to complete. The language model
     * batches tool results (via Promise.all) before yielding them to the stream,
     * so without this early completion signal, the UI wouldn't update until all
     * tools finish. The values set here will be overwritten by merge() when the
     * language model eventually yields the results, but they should be identical.
     */
    complete(result: ToolCallResult): void;
}

/**
 * A tool the provider executed on its own infrastructure (a server tool). Unlike
 * {@link ToolCallChatResponseContent}, it has no confirmation members because it is
 * auto-approved/server-executed; the invocation and its result are surfaced together.
 */
export interface ServerToolCallChatResponseContent extends Required<ChatResponseContent> {
    kind: 'serverToolCall';
    id?: string;
    name?: string;
    arguments?: string;
    finished: boolean;
    result?: ToolCallResult;
    /** Provider-specific metadata needed to faithfully reconstruct the server tool on replay. */
    data?: Record<string, string>;
}

export namespace ServerToolCallChatResponseContent {
    export function is(obj: unknown): obj is ServerToolCallChatResponseContent {
        return ChatResponseContent.is(obj) && obj.kind === 'serverToolCall';
    }
}

export interface ThinkingChatResponseContent
    extends Required<ChatResponseContent> {
    kind: 'thinking';
    content: string;
    signature: string;
}

export interface ProgressChatResponseContent
    extends Required<ChatResponseContent> {
    kind: 'progress';
    message: string;
}

export interface Location {
    uri: URI;
    position: Position;
}
export namespace Location {
    export function is(obj: unknown): obj is Location {
        return !!obj && typeof obj === 'object' &&
            'uri' in obj && (obj as { uri: unknown }).uri instanceof URI &&
            'position' in obj && Position.is((obj as { position: unknown }).position);
    }
}

export interface CustomCallback {
    label: string;
    callback: () => Promise<void>;
}

/**
 * A command chat response content represents a command that is offered to the user for execution.
 * It either refers to an already registered Theia command or provides a custom callback.
 * If both are given, the custom callback will be preferred.
 */
export interface CommandChatResponseContent extends ChatResponseContent {
    kind: 'command';
    command?: Command;
    customCallback?: CustomCallback;
    arguments?: unknown[];
}

/**
 * An informational chat response content represents a message that is purely informational and should not be included in the overall representation of the response.
 */
export interface InformationalChatResponseContent extends ChatResponseContent {
    kind: 'informational';
    content: MarkdownString;
}

export namespace TextChatResponseContent {
    export function is(obj: unknown): obj is TextChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'text' &&
            'content' in obj &&
            typeof (obj as { content: unknown }).content === 'string'
        );
    }
}

export namespace MarkdownChatResponseContent {
    export function is(obj: unknown): obj is MarkdownChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'markdownContent' &&
            'content' in obj &&
            MarkdownString.is((obj as { content: unknown }).content)
        );
    }
}

export namespace InformationalChatResponseContent {
    export function is(obj: unknown): obj is InformationalChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'informational' &&
            'content' in obj &&
            MarkdownString.is((obj as { content: unknown }).content)
        );
    }
}

export namespace CommandChatResponseContent {
    export function is(obj: unknown): obj is CommandChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'command' &&
            'command' in obj &&
            Command.is((obj as { command: unknown }).command)
        );
    }
}

export namespace CodeChatResponseContent {
    export function is(obj: unknown): obj is CodeChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'code' &&
            'code' in obj &&
            typeof (obj as { code: unknown }).code === 'string'
        );
    }
}

export namespace HorizontalLayoutChatResponseContent {
    export function is(
        obj: unknown
    ): obj is HorizontalLayoutChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'horizontal' &&
            'content' in obj &&
            Array.isArray((obj as { content: unknown }).content) &&
            (obj as { content: unknown[] }).content.every(
                ChatResponseContent.is
            )
        );
    }
}

export namespace ToolCallChatResponseContent {
    export function is(obj: unknown): obj is ToolCallChatResponseContent {
        return ChatResponseContent.is(obj) && obj.kind === 'toolCall';
    }

    export interface DenialResult {
        denied: true;
        /** User-provided reason for the denial */
        reason?: string;
    }

    export function isDenialResult(result: unknown): result is DenialResult {
        return typeof result === 'object' && !!result &&
            'denied' in result && (result as DenialResult).denied === true;
    }

    /**
     * Checks if a tool call result contains an error.
     * Supports both ToolCallContent with ToolCallErrorResult items and legacy simple error format.
     */
    export function isErrorResult(result: unknown): boolean {
        if (!result || typeof result !== 'object') {
            return false;
        }
        if ('content' in result && Array.isArray((result as { content: unknown[] }).content)) {
            return (result as { content: Array<{ type?: string }> }).content.some(item => item.type === 'error');
        }
        return 'error' in result && (result as { error: boolean }).error === true;
    }

    /**
     * Extracts the error message from a tool call result.
     * Supports both ToolCallContent with ToolCallErrorResult items and legacy simple error format.
     */
    export function getErrorMessage(result: unknown): string | undefined {
        if (!result || typeof result !== 'object') {
            return undefined;
        }
        if ('content' in result && Array.isArray((result as { content: unknown[] }).content)) {
            const errorItem = (result as { content: Array<{ type?: string; data?: string }> }).content.find(item => item.type === 'error');
            return errorItem?.data;
        }
        if ('error' in result && (result as { error: boolean }).error === true) {
            return (result as { message?: string }).message;
        }
        return undefined;
    }

    /**
     * Checks if a tool call result indicates the tool was not available.
     * This happens when the LLM tries to call a tool that wasn't provided in the request.
     */
    export function isNotAvailableResult(result: unknown): boolean {
        if (!result || typeof result !== 'object') {
            return false;
        }
        if ('content' in result && Array.isArray((result as { content: unknown[] }).content)) {
            return (result as { content: Array<{ type?: string; errorKind?: string }> }).content
                .some(item => item.type === 'error' && item.errorKind === 'tool-not-available');
        }
        return false;
    }
}

/**
 * Represents a streaming delta update for tool call arguments.
 * This content type is used during streaming to append argument fragments
 * to an existing ToolCallChatResponseContent rather than replacing the full arguments.
 */
export interface ToolCallArgumentsDeltaContent extends ChatResponseContent {
    kind: 'toolCallArgumentsDelta';
    /** The tool call ID this delta belongs to */
    id: string;
    /** The argument fragment to append */
    delta: string;
}

export namespace ToolCallArgumentsDeltaContent {
    export function is(obj: unknown): obj is ToolCallArgumentsDeltaContent {
        return ChatResponseContent.is(obj) && obj.kind === 'toolCallArgumentsDelta';
    }
}

export namespace ErrorChatResponseContent {
    export function is(obj: unknown): obj is ErrorChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'error' &&
            'error' in obj &&
            obj.error instanceof Error
        );
    }
}

export namespace ThinkingChatResponseContent {
    export function is(obj: unknown): obj is ThinkingChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'thinking' &&
            'content' in obj &&
            typeof obj.content === 'string'
        );
    }
}

export interface CompactionChatResponseContent extends ChatResponseContent {
    kind: 'compaction';
    provider: string;
    data: unknown;
    summary?: string;
}
export namespace CompactionChatResponseContent {
    export function is(obj: unknown): obj is CompactionChatResponseContent {
        return ChatResponseContent.is(obj) && obj.kind === 'compaction';
    }
}

export namespace ProgressChatResponseContent {
    export function is(obj: unknown): obj is ProgressChatResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'progress' &&
            'message' in obj &&
            typeof obj.message === 'string'
        );
    }
}

export type QuestionResponseHandler = (
    selectedOption: { text: string, value?: string },
) => void;

export type MultiSelectQuestionResponseHandler = (
    selectedOptions: { text: string, value?: string }[],
) => void;

export interface QuestionResponseContent extends ChatResponseContent, InteractiveContent {
    kind: 'question';
    question: string;
    header?: string;
    options: { text: string, value?: string, description?: string }[];
    multiSelect?: boolean;
    selectedOption?: { text: string, value?: string };
    selectedOptions?: { text: string, value?: string }[];
    handler?: QuestionResponseHandler | MultiSelectQuestionResponseHandler;
    /** Called when the user dismisses a single-select question without choosing an option. */
    onSkip?: () => void;
    request?: MutableChatRequestModel;
    /**
     * Whether this question is read-only (restored from persistence without handler).
     * When true, the UI should disable option selection.
     */
    readonly isReadOnly: boolean;
}

export namespace QuestionResponseContent {
    export function is(obj: unknown): obj is QuestionResponseContent {
        return (
            ChatResponseContent.is(obj) &&
            obj.kind === 'question' &&
            'question' in obj &&
            typeof (obj as { question: unknown }).question === 'string' &&
            'options' in obj &&
            Array.isArray((obj as { options: unknown }).options) &&
            (obj as { options: unknown[] }).options.every(option =>
                typeof option === 'object' &&
                option && 'text' in option &&
                typeof (option as { text: unknown }).text === 'string' &&
                ('value' in option ? typeof (option as { value: unknown }).value === 'string' || typeof (option as { value: unknown }).value === 'undefined' : true)
            ) &&
            // handler and request are optional (undefined for restored/read-only questions)
            ('handler' in obj ? (obj as { handler: unknown }).handler === undefined || typeof (obj as { handler: unknown }).handler === 'function' : true) &&
            ('request' in obj ? (obj as { request: unknown }).request === undefined || (obj as { request: unknown }).request instanceof MutableChatRequestModel : true)
        );
    }
}

export interface ResponseTokenUsage {
    readonly inputTokens: number;
    readonly outputTokens: number;
    readonly cacheCreationInputTokens?: number;
    readonly cacheReadInputTokens?: number;
}

export interface ChatResponse {
    readonly content: ChatResponseContent[];
    asString(): string;
    asDisplayString(): string;
}

/**
 * The ChatResponseModel wraps the actual ChatResponse with additional information like the current state, progress messages, a unique id etc.
 */
export interface ChatResponseModel {
    /**
     * Use this to be notified for any change in the response model
     */
    readonly onDidChange: Event<void>;
    /**
     * Fires when this response requires user interaction (e.g., tool confirmation, question response).
     * The content part that needs interaction is provided as the event payload.
     */
    readonly onInteractionNeeded: Event<InteractiveContent & ChatResponseContent>;
    /**
     * Content parts announced via {@link onInteractionNeeded} that currently await user
     * interaction (see {@link InteractiveContent.isAwaitingInteraction}). Lets late
     * subscribers (e.g. a remounted delegation renderer) rebuild pending interaction
     * state instead of relying solely on the push event.
     */
    readonly pendingInteractions: ReadonlyArray<InteractiveContent & ChatResponseContent>;
    /**
     * The unique identifier of the response model
     */
    readonly id: string;
    /**
     * The unique identifier of the request model this response is associated with
     */
    readonly requestId: string;
    /**
     * In case there are progress messages, then they will be stored here
     */
    readonly progressMessages: ChatProgressMessage[];
    /**
     * The actual response content
     */
    readonly response: ChatResponse;
    /**
     * Indicates whether this response is complete. No further changes are expected if 'true'.
     */
    readonly isComplete: boolean;
    /**
     * Indicates whether this response is canceled. No further changes are expected if 'true'.
     */
    readonly isCanceled: boolean;
    /**
     * Some agents might need to wait for user input to continue. This flag indicates that.
     */
    readonly isWaitingForInput: boolean;
    /**
     * Indicates whether an error occurred when processing the response. No further changes are expected if 'true'.
     */
    readonly isError: boolean;
    /**
     * The agent who produced the response content, if there is one.
     */
    readonly agentId?: string
    /**
     * An optional error object that caused the response to be in an error state.
     */
    readonly errorObject?: Error;
    /**
     * Some functionality might want to store some data associated with the response.
     * This can be used to store and retrieve such data.
     */
    readonly data: { [key: string]: unknown };
    /**
     * The ID of the prompt variant used to generate this response
     */
    readonly promptVariantId?: string;
    /**
     * Indicates whether the prompt variant was customized/edited
     */
    readonly isPromptVariantEdited?: boolean;
    /**
     * The identifier of the language model that produced this response, if recorded.
     */
    readonly languageModel?: string;
    readonly tokenUsage?: ResponseTokenUsage;
    toSerializable(): SerializableChatResponseData;
}

/**********************
 * Implementations
 **********************/

export class MutableChatModel implements ChatModel, Disposable {
    protected readonly _onDidChangeEmitter = new Emitter<ChatChangeEvent>();
    onDidChange: Event<ChatChangeEvent> = this._onDidChangeEmitter.event;

    protected readonly toDispose = new DisposableCollection();

    protected _hierarchy: ChatRequestHierarchy<MutableChatRequestModel>;
    protected _id: string;
    protected _suggestions: readonly ChatSuggestion[] = [];
    protected readonly _contextManager = new ChatContextManagerImpl();
    protected _changeSet: ChatTreeChangeSet;
    protected _settings: ChatSessionSettings;
    protected _location: ChatAgentLocation;
    protected _status: ChatSessionStatus = 'idle';
    rootSessionId?: string;
    parentSessionId?: string;

    get location(): ChatAgentLocation {
        return this._location;
    }

    constructor(
        locationOrSerializedData: ChatAgentLocation | SerializedChatModel = ChatAgentLocation.Panel
    ) {
        // Check if we're restoring from serialized data
        if (this.isSerializedChatModel(locationOrSerializedData)) {
            this.restoreFromSerializedData(locationOrSerializedData);
        } else {
            // Normal creation path
            this._location = locationOrSerializedData;
            this._id = generateUuid();
            this._hierarchy = new ChatRequestHierarchyImpl<MutableChatRequestModel>();
            this._changeSet = new ChatTreeChangeSet(this._hierarchy);
            this.toDispose.push(this._changeSet);
            this._changeSet.onDidChange(this._onDidChangeEmitter.fire, this._onDidChangeEmitter, this.toDispose);
        }

        this.toDispose.pushAll([
            this._onDidChangeEmitter,
            this._contextManager.onDidChange(this._onDidChangeEmitter.fire, this._onDidChangeEmitter),
            this._hierarchy.onDidChange(event => {
                this._onDidChangeEmitter.fire({
                    kind: 'changeHierarchyBranch',
                    branch: event.branch,
                });
            }),
            // Re-derive the aggregated session status whenever anything in the model changes.
            this.onDidChange(event => {
                if (!ChatChangeEvent.isStatusChangedEvent(event)) {
                    this.updateStatus();
                }
            }),
        ]);
        this._status = this.computeStatus();
    }

    /**
     * Type guard to determine if we're receiving serialized data
     */
    protected isSerializedChatModel(data: ChatAgentLocation | SerializedChatModel): data is SerializedChatModel {
        return typeof data === 'object' && 'sessionId' in data && 'requests' in data && 'responses' in data;
    }

    /**
     * Restore this chat model from serialized data
     *
     * Does not restore response content or changesets.
     * This handled by the chat service using the deserializer registries.
     */
    protected restoreFromSerializedData(data: SerializedChatModel): void {
        this._id = data.sessionId;
        this._location = data.location;

        // First, create all request models and build a map
        const requestMap = new Map<string, MutableChatRequestModel>();
        for (const reqData of data.requests) {
            const respData = data.responses.find(r => r.requestId === reqData.id);
            const requestModel = new MutableChatRequestModel(
                this,
                reqData,
                respData
            );
            requestMap.set(requestModel.id, requestModel);
            requestModel.onDidChange(event => {
                if (!ChatChangeEvent.isChangeSetEvent(event)) {
                    this._onDidChangeEmitter.fire(event);
                }
            }, this, this.toDispose);
        }

        // Restore per-session settings (e.g. the per-session model override) so the chat input can
        // reflect the previous selection.
        if (data.settings) {
            this._settings = data.settings;
        }

        // Restore the hierarchy structure with all alternatives
        this._hierarchy = new ChatRequestHierarchyImpl<MutableChatRequestModel>(data.hierarchy, requestMap);

        // Register all requests with changeset
        this._changeSet = new ChatTreeChangeSet(this._hierarchy);
        this.toDispose.push(this._changeSet);
        this._changeSet.onDidChange(this._onDidChangeEmitter.fire, this._onDidChangeEmitter, this.toDispose);

        for (const requestModel of requestMap.values()) {
            this._changeSet.registerRequest(requestModel);
        }
    }

    get id(): string {
        return this._id;
    }

    get changeSet(): ChangeSet {
        return this._changeSet;
    }

    getBranches(): ChatHierarchyBranch<ChatRequestModel>[] {
        return this._hierarchy.activeBranches();
    }

    getBranch(requestId: string): ChatHierarchyBranch<ChatRequestModel> | undefined {
        return this._hierarchy.findBranch(requestId);
    }

    getRequests(): MutableChatRequestModel[] {
        return this._hierarchy.activeRequests();
    }

    getRequest(id: string): MutableChatRequestModel | undefined {
        return this.getRequests().find(request => request.id === id);
    }

    get suggestions(): readonly ChatSuggestion[] {
        return this._suggestions;
    }

    get status(): ChatSessionStatus {
        return this._status;
    }

    protected computeStatus(): ChatSessionStatus {
        return ChatSessionStatus.fromRequest(this.getRequests().at(-1));
    }

    protected updateStatus(): void {
        const status = this.computeStatus();
        if (status !== this._status) {
            this._status = status;
            this._onDidChangeEmitter.fire({ kind: 'statusChanged', status });
        }
    }

    get context(): ChatContextManager {
        return this._contextManager;
    }

    get settings(): ChatSessionSettings {
        return this._settings;
    }

    setSettings(settings: ChatSessionSettings): void {
        this._settings = settings;
        // Emit a change so listeners (e.g. session auto-save) persist selector-only or dialog-only
        // settings updates that are not accompanied by another model change.
        this._onDidChangeEmitter.fire({ kind: 'settingsChanged', settings });
    }

    addChildModel(child: MutableChatModel): Disposable {
        const disposable = new DisposableCollection();
        disposable.push(child.onDidChange(event => {
            if (ChatChangeEvent.isInteractionNeededEvent(event)) {
                this._onDidChangeEmitter.fire(event);
            }
        }));
        this.toDispose.push(disposable);
        return disposable;
    }

    addRequest(parsedChatRequest: ParsedChatRequest, agentId?: string, context: ChatContext = { variables: [] }): MutableChatRequestModel {
        const add = this.getTargetForRequestAddition(parsedChatRequest);
        const requestModel = new MutableChatRequestModel(
            this,
            parsedChatRequest,
            agentId,
            context
        );
        requestModel.onDidChange(event => {
            if (!ChatChangeEvent.isChangeSetEvent(event)) {
                this._onDidChangeEmitter.fire(event);
            }
        }, this, this.toDispose);

        add(requestModel);
        this._changeSet.registerRequest(requestModel);

        this._onDidChangeEmitter.fire({
            kind: 'addRequest',
            request: requestModel,
        });
        return requestModel;
    }

    protected getTargetForRequestAddition(request: ParsedChatRequest): (addendum: MutableChatRequestModel) => void {
        const requestId = request.request.referencedRequestId;
        const branch = requestId !== undefined && this._hierarchy.findBranch(requestId);
        if (requestId !== undefined && !branch) { throw new Error(`Cannot find branch for requestId: ${requestId}`); }
        return branch ? branch.add.bind(branch) : this._hierarchy.append.bind(this._hierarchy);
    }

    setSuggestions(suggestions: ChatSuggestion[]): void {
        this._suggestions = Object.freeze(suggestions);
        this._onDidChangeEmitter.fire({
            kind: 'suggestionsChanged',
            suggestions
        });
    }

    isEmpty(): boolean {
        return this.getRequests().length === 0;
    }

    toSerializable(): SerializedChatModel {
        const hierarchy = this._hierarchy.toSerializable();

        const allRequests = this.getAllRequests();

        const serializedRequests: SerializableChatRequestData[] = allRequests.map(req => req.toSerializable());
        const serializedResponses: SerializableChatResponseData[] = allRequests
            .filter(req => req.response)
            .map(req => req.response.toSerializable());

        return {
            sessionId: this._id,
            location: this.location,
            hierarchy,
            requests: serializedRequests,
            responses: serializedResponses,
            settings: this._settings
        };
    }

    /**
     * Get all requests from the hierarchy.
     * This is used for operations that need to process all requests, such as serialization.
     */
    getAllRequests(): MutableChatRequestModel[] {
        const allRequests: MutableChatRequestModel[] = [];
        const visited = new Set<string>();

        const collectFromBranch = (branch: ChatHierarchyBranch<MutableChatRequestModel>): void => {
            for (const item of branch.items) {
                // Avoid duplicates
                if (!visited.has(item.element.id)) {
                    visited.add(item.element.id);
                    allRequests.push(item.element);
                }

                // Recursively collect from next branches
                if (item.next) {
                    collectFromBranch(item.next);
                }
            }
        };

        collectFromBranch(this._hierarchy.branch);
        return allRequests;
    }

    dispose(): void {
        this.toDispose.dispose();
    }
}

export class ChatTreeChangeSet implements Omit<ChangeSet, 'onDidChange'> {
    protected readonly onDidChangeEmitter = new Emitter<ChatUpdateChangeSetEvent>();
    get onDidChange(): Event<ChatUpdateChangeSetEvent> {
        return this.onDidChangeEmitter.event;
    }

    protected readonly toDispose = new DisposableCollection();

    constructor(protected readonly hierarchy: ChatRequestHierarchy<MutableChatRequestModel>) {
        hierarchy.onDidChange(this.handleChangeSetChange, this, this.toDispose);
    }

    get title(): string {
        return this.getCurrentChangeSet()?.title ?? '';
    }

    removeElements(...uris: URI[]): boolean {
        return this.getMutableChangeSet().removeElements(...uris);
    }

    addElements(...elements: ChangeSetElement[]): boolean {
        return this.getMutableChangeSet().addElements(...elements);
    }

    setElements(...elements: ChangeSetElement[]): void {
        this.getMutableChangeSet().setElements(...elements);
    }

    setTitle(title: string): void {
        this.getMutableChangeSet().setTitle(title);
    }

    getElementByURI(uri: URI): ChangeSetElement | undefined {
        return this.currentElements.find(candidate => candidate.uri.isEqual(uri));
    }

    protected currentElements: ChangeSetElement[] = [];
    protected handleChangeSetChange = debounce(this.doHandleChangeSetChange.bind(this), 100, { leading: false, trailing: true });
    protected doHandleChangeSetChange(): void {
        const newElements = this.computeChangeSetElements();
        this.handleElementChange(newElements);
        this.currentElements = newElements;
        this.onDidChangeEmitter.fire({ kind: 'updateChangeSet', elements: this.currentElements, title: this.getCurrentChangeSet()?.title });
    }

    getElements(): ChangeSetElement[] {
        return this.currentElements;
    }

    protected computeChangeSetElements(): ChangeSetElement[] {
        const allElements = ChangeSetImpl.combine((function* (requests: MutableChatRequestModel[]): IterableIterator<ChangeSetImpl> {
            for (let i = requests.length - 1; i >= 0; i--) {
                const changeSet = requests[i].changeSet;
                if (changeSet) { yield changeSet; }
            }
        })(this.hierarchy.activeRequests()));
        return ArrayUtils.coalesce(Array.from(allElements.values()));
    }

    protected handleElementChange(newElements: ChangeSetElement[]): void {
        const old = new Set(this.currentElements);
        for (const element of newElements) {
            if (!old.delete(element)) {
                element.onShow?.();
            }
        }
        for (const element of old) {
            element.onHide?.();
        }
    }

    protected toDisposeOnRequestAdded = new DisposableCollection();
    registerRequest(request: MutableChatRequestModel): void {
        request.onDidChange(event => event.kind === 'updateChangeSet' && this.handleChangeSetChange(), this, this.toDispose);
        if (this.localChangeSet) {
            request.changeSet = this.localChangeSet;
            this.localChangeSet = undefined;
        }
        this.toDisposeOnRequestAdded.dispose();
    }

    protected localChangeSet?: ChangeSetImpl;
    protected getMutableChangeSet(): ChangeSetImpl {
        const tipRequest = this.hierarchy.activeRequests().at(-1);
        const existingChangeSet = tipRequest?.changeSet;
        if (existingChangeSet) {
            return existingChangeSet;
        }
        if (this.localChangeSet && tipRequest) {
            throw new Error('Non-empty chat model retained reference to own change set. This is unexpected!');
        }
        if (this.localChangeSet) {
            return this.localChangeSet;
        }
        const newChangeSet = new ChangeSetImpl();
        if (tipRequest) {
            tipRequest.changeSet = newChangeSet;
        } else {
            this.localChangeSet = newChangeSet;
            newChangeSet.onDidChange(this.handleChangeSetChange, this, this.toDisposeOnRequestAdded);
        }
        return newChangeSet;
    }

    protected getCurrentChangeSet(): ChangeSet | undefined {
        const holder = this.getBranchParent(candidate => !!candidate.get().changeSet);
        return holder?.get().changeSet ?? this.localChangeSet;
    }

    /** Returns the lowest node among active nodes that satisfies {@link criterion} */
    getBranchParent(criterion: (branch: ChatHierarchyBranch<MutableChatRequestModel>) => boolean): ChatHierarchyBranch<MutableChatRequestModel> | undefined {
        const branches = this.hierarchy.activeBranches();
        for (let i = branches.length - 1; i >= 0; i--) {
            const branch = branches[i];
            if (criterion?.(branch)) { return branch; }
        }
        return branches.at(0);
    }

    dispose(): void {
        this.toDispose.dispose();
    }
}

export class ChatRequestHierarchyImpl<TRequest extends ChatRequestModel = ChatRequestModel> implements ChatRequestHierarchy<TRequest> {
    protected readonly onDidChangeActiveBranchEmitter = new Emitter<ChangeActiveBranchEvent<TRequest>>();
    readonly onDidChange = this.onDidChangeActiveBranchEmitter.event;

    readonly branch: ChatHierarchyBranch<TRequest>;

    constructor(
        serializedHierarchy?: SerializableHierarchy,
        requestMap?: Map<string, TRequest>) {
        this.branch = new ChatRequestHierarchyBranchImpl<TRequest>(this);
        if (serializedHierarchy && requestMap) {
            this.restoreFromSerialized(serializedHierarchy, requestMap);
        }
    }

    /**
     * Restore the hierarchy from serialized data.
     */
    protected restoreFromSerialized(
        serializedHierarchy: SerializableHierarchy,
        requestMap: Map<string, TRequest>
    ): void {
        // Build a map of branch IDs to restored branch objects
        const branchMap = new Map<string, ChatHierarchyBranch<TRequest>>();

        // Function to restore a branch and its descendants
        const restoreBranch = (branchId: string): ChatHierarchyBranch<TRequest> => {
            // Check if already restored
            if (branchMap.has(branchId)) {
                return branchMap.get(branchId)!;
            }

            const serializedBranch = serializedHierarchy.branches[branchId];
            if (!serializedBranch) {
                throw new Error(`Cannot find serialized branch with id: ${branchId}`);
            }

            // Restore items in this branch
            const items: ChatHierarchyBranchItem<TRequest>[] = serializedBranch.items.map(serializedItem => {
                const request = requestMap.get(serializedItem.requestId);
                if (!request) {
                    throw new Error(`Cannot find request with id: ${serializedItem.requestId}`);
                }

                // Restore next branch if present
                const next = serializedItem.nextBranchId
                    ? restoreBranch(serializedItem.nextBranchId)
                    : undefined;

                return {
                    element: request,
                    next
                };
            });

            // Determine if this is the root branch
            const isRoot = branchId === serializedHierarchy.rootBranchId;

            if (isRoot) {
                // For root branch, we need to replace the existing branch's internals
                // Cast is safe here as we know this.branch is a ChatRequestHierarchyBranchImpl
                const rootBranch = this.branch as ChatRequestHierarchyBranchImpl<TRequest>;
                // Use Object.assign to update the readonly properties
                Object.assign(rootBranch, {
                    id: branchId,
                    items,
                    _activeIndex: serializedBranch.activeBranchIndex
                });
                branchMap.set(branchId, rootBranch);
                return rootBranch;
            } else {
                // For non-root branches, use constructor-based deserialization
                const restoredBranch = new ChatRequestHierarchyBranchImpl<TRequest>(
                    this,
                    undefined, // previous will be set by parent
                    items,
                    serializedBranch.activeBranchIndex,
                    branchId
                );
                branchMap.set(branchId, restoredBranch);
                return restoredBranch;
            }
        };

        // Start restoration from the root branch
        restoreBranch(serializedHierarchy.rootBranchId);
    }

    append(request: TRequest): ChatHierarchyBranch<TRequest> {
        const branches = this.activeBranches();

        if (branches.length === 0) {
            this.branch.add(request);
            return this.branch;
        }

        return branches.at(-1)!.continue(request);
    }

    activeRequests(): TRequest[] {
        return this.activeBranches().map(h => h.get());
    }

    activeBranches(): ChatHierarchyBranch<TRequest>[] {
        return Array.from(this.iterateBranches());
    }

    protected *iterateBranches(): Generator<ChatHierarchyBranch<TRequest>> {
        let current: ChatHierarchyBranch<TRequest> | undefined = this.branch;
        while (current) {
            if (current.items.length > 0) {
                yield current;
                current = current.next();
            } else {
                break;
            }
        }
    }

    findRequest(requestId: string): TRequest | undefined {
        const branch = this.findInBranch(this.branch, requestId);
        return branch?.items.find(item => item.element.id === requestId)?.element;
    }

    findBranch(requestId: string): ChatHierarchyBranch<TRequest> | undefined {
        return this.findInBranch(this.branch, requestId);
    }

    protected findInBranch(branch: ChatHierarchyBranch<TRequest>, requestId: string): ChatHierarchyBranch<TRequest> | undefined {
        for (const item of branch.items) {
            if (item.element.id === requestId) {
                return branch;
            }
        }
        for (const item of branch.items) {
            if (item.next) {
                const found = this.findInBranch(item.next, requestId);
                if (found) {
                    return found;
                }
            }
        }
        return undefined;
    }

    notifyChange(event: ChangeActiveBranchEvent<TRequest>): void {
        this.onDidChangeActiveBranchEmitter.fire(event);
    }

    toSerializable(): SerializableHierarchy {
        const branches: { [branchId: string]: SerializableHierarchyBranch } = {};

        // Recursively serialize all branches starting from the root
        this.serializeBranch(this.branch, branches);

        return {
            rootBranchId: this.branch.id,
            branches
        };
    }

    protected serializeBranch(
        branch: ChatHierarchyBranch<TRequest>,
        branches: { [branchId: string]: SerializableHierarchyBranch }
    ): void {
        const items: SerializableHierarchyBranchItem[] = branch.items.map(item => {
            if (item.next) {
                this.serializeBranch(item.next, branches);
            }
            return {
                requestId: item.element.id,
                nextBranchId: item.next?.id
            };
        });

        branches[branch.id] = {
            id: branch.id,
            items,
            activeBranchIndex: branch.activeBranchIndex
        };
    }

    dispose(): void {
        this.onDidChangeActiveBranchEmitter.dispose();
        this.branch.dispose();
    }
}

export class ChatRequestHierarchyBranchImpl<TRequest extends ChatRequestModel> implements ChatHierarchyBranch<TRequest> {
    readonly id: string;

    constructor(
        readonly hierarchy: ChatRequestHierarchy<TRequest>,
        readonly previous?: ChatHierarchyBranch<TRequest>,
        readonly items: ChatHierarchyBranchItem<TRequest>[] = [],
        protected _activeIndex = -1,
        id?: string
    ) {
        this.id = id ?? generateUuid();
    }

    get activeBranchIndex(): number {
        return this._activeIndex;
    }

    protected set activeBranchIndex(value: number) {
        this._activeIndex = value;
        this.hierarchy.notifyChange({
            branch: this,
            item: this.items[this._activeIndex]
        });
    }

    next(): ChatHierarchyBranch<TRequest> | undefined {
        return this.items[this.activeBranchIndex]?.next;
    }

    get(): TRequest {
        return this.items[this.activeBranchIndex].element;
    }

    add(request: TRequest): void {
        const branch: ChatHierarchyBranchItem<TRequest> = {
            element: request
        };
        this.items.push(branch);
        this.activeBranchIndex = this.items.length - 1;
    }

    remove(request: TRequest | string): void {
        const requestId = typeof request === 'string' ? request : request.id;
        const index = this.items.findIndex(version => version.element.id === requestId);
        if (index !== -1) {
            this.items.splice(index, 1);
            if (this.activeBranchIndex >= index) {
                this.activeBranchIndex--;
            }
        }
    }

    continue(request: TRequest): ChatHierarchyBranch<TRequest> {
        if (this.items.length === 0) {
            this.add(request);
            return this;
        }

        const item = this.items[this.activeBranchIndex];

        if (item) {
            const next = new ChatRequestHierarchyBranchImpl(this.hierarchy, this, [{ element: request }], 0);
            this.items[this.activeBranchIndex] = {
                ...item,
                next
            };
            return next;
        }

        throw new Error(`No current branch to continue from. Active Index: ${this.activeBranchIndex}`);
    }

    enable(request: TRequest): ChatHierarchyBranchItem<TRequest> {
        this.activeBranchIndex = this.items.findIndex(pred => pred.element.id === request.id);
        return this.items[this.activeBranchIndex];
    }

    enablePrevious(): ChatHierarchyBranchItem<TRequest> {
        if (this.activeBranchIndex > 0) {
            this.activeBranchIndex--;
            return this.items[this.activeBranchIndex];
        }
        return this.items[0];
    }

    enableNext(): ChatHierarchyBranchItem<TRequest> {
        if (this.activeBranchIndex < this.items.length - 1) {
            this.activeBranchIndex++;
            return this.items[this.activeBranchIndex];
        }

        return this.items[this.activeBranchIndex];
    }

    succeedingBranches(): ChatHierarchyBranch<TRequest>[] {
        const branches: ChatHierarchyBranch<TRequest>[] = [];

        let current: ChatHierarchyBranch<TRequest> | undefined = this;
        while (current !== undefined) {
            branches.push(current);
            current = current.next();
        }

        return branches;
    }

    dispose(): void {
        if (Disposable.is(this.get())) {
            this.items.forEach(({ element }) => Disposable.is(element) && element.dispose());
        }
        this.items.length = 0;
    }
}

export class ChatContextManagerImpl implements ChatContextManager {
    protected readonly variables = new Array<AIVariableResolutionRequest>();
    protected readonly onDidChangeEmitter = new Emitter<ChatAddVariableEvent | ChatRemoveVariableEvent | ChatSetVariablesEvent>();
    get onDidChange(): Event<ChatAddVariableEvent | ChatRemoveVariableEvent | ChatSetVariablesEvent> {
        return this.onDidChangeEmitter.event;
    }

    constructor(context?: ChatContext) {
        if (context) {
            this.variables.push(...context.variables.map(AIVariableResolutionRequest.fromResolved));
        }
    }

    getVariables(): readonly AIVariableResolutionRequest[] {
        const result = this.variables.slice();
        Object.freeze(result);
        return result;
    }

    addVariables(...variables: AIVariableResolutionRequest[]): void {
        let modified = false;
        variables.forEach(variable => {
            if (this.variables.some(existing => existing.variable.id === variable.variable.id && existing.arg === variable.arg)) {
                return;
            }
            this.variables.push(variable);
            modified = true;
        });
        if (modified) {
            this.onDidChangeEmitter.fire({ kind: 'addVariable' });
        }
    }

    deleteVariables(...indices: number[]): void {
        const toDelete = indices.filter(candidate => candidate <= this.variables.length).sort((left, right) => right - left);
        if (toDelete.length) {
            toDelete.forEach(index => {
                this.variables.splice(index, 1);
            });
            this.onDidChangeEmitter.fire({ kind: 'removeVariable' });
        }
    }

    setVariables(variables: AIVariableResolutionRequest[]): void {
        this.variables.length = 0;
        variables.forEach(variable => {
            if (this.variables.some(existing => existing.variable.id === variable.variable.id && existing.arg === variable.arg)) {
                return;
            }
            this.variables.push(variable);
        });
        this.onDidChangeEmitter.fire({ kind: 'setVariables' });
    }

    clear(): void {
        if (this.variables.length) {
            this.variables.length = 0;
            this.onDidChangeEmitter.fire({ kind: 'removeVariable' });
        }
    }
}

export class MutableChatRequestModel implements ChatRequestModel, EditableChatRequestModel, Disposable {
    protected readonly _onDidChangeEmitter = new Emitter<ChatChangeEvent>();
    onDidChange: Event<ChatChangeEvent> = this._onDidChangeEmitter.event;
    protected _id: string;
    protected _session: MutableChatModel;
    protected _request: ChatRequest;
    protected _response: MutableChatResponseModel;
    protected _changeSet?: ChangeSetImpl;
    protected _context: ChatContext;
    protected _agentId?: string;
    protected _data: { [key: string]: unknown };
    protected _isEditing = false;
    protected _message: ParsedChatRequest;

    protected readonly toDispose = new DisposableCollection();
    readonly editContextManager: ChatContextManagerImpl;

    constructor(
        session: MutableChatModel,
        messageOrData: ParsedChatRequest | SerializableChatRequestData,
        agentIdOrResponseData?: string | SerializableChatResponseData,
        context: ChatContext = { variables: [] },
        data: { [key: string]: unknown } = {}
    ) {
        this._session = session;

        // Check if we're restoring from serialized data
        if (this.isSerializedRequestData(messageOrData)) {
            this.restoreFromSerializedData(messageOrData, agentIdOrResponseData as SerializableChatResponseData | undefined);
        } else {
            // Normal creation path
            this._request = messageOrData.request;
            this._id = generateUuid();
            this._response = new MutableChatResponseModel(this._id, agentIdOrResponseData as string | undefined);
            this._context = context;
            this._agentId = agentIdOrResponseData as string | undefined;
            this._data = data;
            // Store the parsed message
            this._message = messageOrData;
        }

        this.editContextManager = new ChatContextManagerImpl(this._context);
        this.editContextManager.onDidChange(this._onDidChangeEmitter.fire, this._onDidChangeEmitter, this.toDispose);

        // Wire response changes to propagate through request to session
        this._response.onDidChange(() => {
            this._onDidChangeEmitter.fire({ kind: 'responseChanged' });
        }, this, this.toDispose);

        // Wire interaction needed events to propagate through request to session
        this._response.onInteractionNeeded(contentPart => {
            this._onDidChangeEmitter.fire({ kind: 'interactionNeeded', contentPart });
        }, this, this.toDispose);

        this.toDispose.push(this._onDidChangeEmitter);
    }

    /**
     * Type guard to determine if we're receiving serialized data
     */
    protected isSerializedRequestData(data: ParsedChatRequest | SerializableChatRequestData): data is SerializableChatRequestData {
        return 'id' in data && 'text' in data && !('request' in data);
    }

    /**
     * Restore this request model from serialized data
     */
    protected restoreFromSerializedData(
        reqData: SerializableChatRequestData,
        respData?: SerializableChatResponseData
    ): void {
        this._id = reqData.id;
        this._request = {
            text: reqData.text,
            capabilityOverrides: reqData.capabilityOverrides,
            genericCapabilitySelections: reqData.genericCapabilitySelections,
            serverToolSelections: reqData.serverToolSelections
        };
        this._agentId = reqData.agentId;
        this._data = {};
        this._context = { variables: [] };

        if (reqData.parsedRequest) {
            this._message = this.deserializeParsedRequest(
                reqData.parsedRequest,
                this._request
            );
        } else {
            this._message = {
                request: this._request,
                parts: [new ParsedChatRequestTextPart(
                    { start: 0, endExclusive: reqData.text.length },
                    reqData.text
                )],
                toolRequests: new Map(),
                variables: []
            };
        }

        // Restore response if present
        if (respData) {
            this._response = new MutableChatResponseModel(this._id, this._agentId, respData);
        } else {
            this._response = new MutableChatResponseModel(this._id, this._agentId);
        }
    }

    /**
     * Deserialize ParsedChatRequest from serialized data.
     * Creates placeholder tool requests - actual tools will be restored by ChatService.
     */
    protected deserializeParsedRequest(
        data: SerializableParsedRequest,
        request: ChatRequest
    ): ParsedChatRequest {
        const parts: ParsedChatRequestPart[] = data.parts.map(partData => {
            switch (partData.kind) {
                case 'text':
                    return new ParsedChatRequestTextPart(
                        partData.range,
                        partData.text
                    );
                case 'var': {
                    const varPart = new ParsedChatRequestVariablePart(
                        partData.range,
                        partData.variableName,
                        partData.variableArg
                    );
                    if (partData.variableValue !== undefined) {
                        varPart.resolution = {
                            variable: {
                                id: partData.variableId,
                                name: partData.variableName,
                                description: partData.variableDescription
                            },
                            arg: partData.variableArg,
                            value: partData.variableValue
                        };
                    }
                    return varPart;
                }
                case 'function':
                    // Create placeholder - will be restored by ChatService
                    return new ParsedChatRequestFunctionPart(
                        partData.range,
                        this.createPlaceholderToolRequest(partData.toolRequestId),
                        partData.deferred === true
                    );
                case 'agent':
                    return new ParsedChatRequestAgentPart(
                        partData.range,
                        partData.agentId,
                        partData.agentName
                    );
                default:
                    throw new Error(`Unknown part kind: ${(partData as SerializableParsedRequestPart).kind}`);
            }
        });

        // Create placeholder tool requests map - will be via restoreToolRequests later
        const toolRequests = new Map<string, ToolRequest>();
        const deferredToolIds = new Set<string>();
        for (const toolData of data.toolRequests) {
            toolRequests.set(toolData.id, this.createPlaceholderToolRequest(toolData.id));
            if (toolData.deferred) {
                deferredToolIds.add(toolData.id);
            }
        }

        const variables: ResolvedAIVariable[] = data.variables.map(varData => ({
            variable: {
                id: varData.variableId,
                name: varData.variableName,
                description: varData.variableDescription
            },
            arg: varData.arg,
            value: varData.value
        }));

        return {
            request,
            parts,
            toolRequests,
            deferredToolIds: deferredToolIds.size > 0 ? deferredToolIds : undefined,
            variables
        };
    }

    /**
     * Creates a placeholder tool request that will be replaced during restoration.
     */
    protected createPlaceholderToolRequest(toolId: string): ToolRequest {
        return {
            id: toolId,
            name: toolId,
            parameters: { type: 'object' as const, properties: {} },
            handler: async () => {
                throw new Error(`Tool request '${toolId}' not yet restored. This is a placeholder.`);
            }
        };
    }

    /**
     * Restores the tool requests in the parsed request by replacing placeholders with actual tools.
     * Called after deserialization to upgrade placeholder tools to real tools from the registry.
     */
    restoreToolRequests(toolRequests: Map<string, ToolRequest>): void {
        this._message.toolRequests.clear();
        for (const [id, tool] of toolRequests) {
            this._message.toolRequests.set(id, tool);
        }

        for (const part of this._message.parts) {
            if (part instanceof ParsedChatRequestFunctionPart) {
                const actualTool = toolRequests.get(part.toolRequest.id);
                if (actualTool) {
                    Object.assign(part, { toolRequest: actualTool });
                }
            }
        }
    }

    get message(): ParsedChatRequest {
        return this._message;
    }

    get changeSet(): ChangeSetImpl | undefined {
        return this._changeSet;
    }

    set changeSet(changeSet: ChangeSetImpl) {
        this._changeSet?.dispose();
        this._changeSet = changeSet;
        this.toDispose.push(changeSet);
        changeSet.onDidChange(() => this._onDidChangeEmitter.fire({ kind: 'updateChangeSet', elements: changeSet.getElements(), title: changeSet.title }), this, this.toDispose);
        this._onDidChangeEmitter.fire({ kind: 'updateChangeSet', elements: changeSet.getElements(), title: changeSet.title });
    }

    get isEditing(): boolean {
        return this._isEditing;
    }

    enableEdit(): void {
        this._isEditing = true;
        this.emitEditRequest(this);
    }

    get data(): { [key: string]: unknown } | undefined {
        return this._data;
    }

    addData(key: string, value: unknown): void {
        this._data[key] = value;
    }

    getDataByKey<T = unknown>(key: string): T {
        return this._data[key] as T;
    }

    removeData(key: string): void {
        delete this._data[key];
    }

    get id(): string {
        return this._id;
    }

    get session(): MutableChatModel {
        return this._session;
    }

    get request(): ChatRequest {
        return this._request;
    }

    get response(): MutableChatResponseModel {
        return this._response;
    }

    get context(): ChatContext {
        return this._context;
    }

    get agentId(): string | undefined {
        return this._agentId;
    }

    cancelEdit(): void {
        if (this.isEditing) {
            this._isEditing = false;
            this.emitCancelEdit(this);

            this.clearEditContext();
        }
    }

    submitEdit(newRequest: ChatRequest): void {
        if (this.isEditing) {
            this._isEditing = false;
            const variables = this.editContextManager.getVariables() ?? [];

            this.emitSubmitEdit(this, {
                ...newRequest,
                referencedRequestId: this.id,
                variables
            });

            this.clearEditContext();
        }
    }

    cancel(): void {
        this.response.cancel();
    }

    toSerializable(): SerializableChatRequestData {
        return {
            id: this.id,
            text: this.request.text,
            agentId: this.agentId,
            changeSet: this._changeSet ? {
                title: this._changeSet.title,
                elements: this._changeSet.getElements().map(elem => elem.toSerializable?.()).filter((elem): elem is SerializableChangeSetElement => elem !== undefined)
            } : undefined,
            parsedRequest: this.message ? ParsedChatRequest.toSerializable(this.message) : undefined,
            capabilityOverrides: this.request.capabilityOverrides,
            genericCapabilitySelections: this.request.genericCapabilitySelections,
            serverToolSelections: this.request.serverToolSelections
        };
    }

    dispose(): void {
        this.toDispose.dispose();
    }

    protected clearEditContext(): void {
        this.editContextManager.setVariables(this.context.variables.map(AIVariableResolutionRequest.fromResolved));
    }

    protected emitEditRequest(request: MutableChatRequestModel): void {
        const branch = this.session.getBranch(request.id);
        if (!branch) {
            throw new Error(`Cannot find hierarchy for requestId: ${request.id}`);
        }
        this._onDidChangeEmitter.fire({
            kind: 'enableEdit',
            request,
            branch,
        });
    }

    protected emitCancelEdit(request: MutableChatRequestModel): void {
        const branch = this.session.getBranch(request.id);
        if (!branch) {
            throw new Error(`Cannot find branch for requestId: ${request.id}`);
        }
        this._onDidChangeEmitter.fire({
            kind: 'cancelEdit',
            request,
            branch,
        });
    }

    protected emitSubmitEdit(request: MutableChatRequestModel, newRequest: ChatRequest): void {
        const branch = this.session.getBranch(request.id);
        if (!branch) {
            throw new Error(`Cannot find branch for requestId: ${request.id}`);
        }
        this._onDidChangeEmitter.fire({
            kind: 'submitEdit',
            request,
            branch,
            newRequest
        });
    }
}

export class ErrorChatResponseContentImpl implements ErrorChatResponseContent {
    readonly kind = 'error';
    protected _error: Error;
    constructor(error: Error) {
        this._error = error;
    }
    get error(): Error {
        return this._error;
    }
    asString(): string | undefined {
        return undefined;
    }
    asDisplayString(): string | undefined {
        return this._error.message;
    }
    toSerializable(): SerializableChatResponseContentData<ErrorContentData> {
        return {
            kind: 'error',
            data: {
                message: this._error.message,
                stack: this._error.stack
            }
        };
    }
}

export class TextChatResponseContentImpl implements TextChatResponseContent {
    readonly kind = 'text';
    protected _content: string;

    constructor(content: string) {
        this._content = content;
    }

    get content(): string {
        return this._content;
    }

    asString(): string {
        return this._content;
    }

    asDisplayString(): string | undefined {
        return this.asString();
    }

    merge(nextChatResponseContent: TextChatResponseContent): boolean {
        this._content += nextChatResponseContent.content;
        return true;
    }

    toLanguageModelMessage(): TextMessage {
        return {
            actor: 'ai',
            type: 'text',
            text: this.content
        };
    }

    toSerializable(): SerializableChatResponseContentData<TextContentData> {
        return {
            kind: 'text',
            data: { content: this._content }
        };
    }
}

export class ThinkingChatResponseContentImpl implements ThinkingChatResponseContent {
    readonly kind = 'thinking';
    protected _content: string;
    protected _signature: string;

    constructor(content: string, signature: string) {
        this._content = content;
        this._signature = signature;
    }

    get content(): string {
        return this._content;
    }
    get signature(): string {
        return this._signature;
    }

    asString(): string {
        return JSON.stringify({
            type: 'thinking',
            thinking: this.content,
            signature: this.signature
        });
    }

    asDisplayString(): string | undefined {
        return `<Thinking>${this.content}</Thinking>`;
    }

    merge(nextChatResponseContent: ThinkingChatResponseContent): boolean {
        this._content += nextChatResponseContent.content;
        this._signature += nextChatResponseContent.signature;
        return true;
    }

    toLanguageModelMessage(): ThinkingMessage {
        return {
            actor: 'ai',
            type: 'thinking',
            thinking: this.content,
            signature: this.signature
        };
    }

    toSerializable(): SerializableChatResponseContentData<ThinkingContentData> {
        return {
            kind: 'thinking',
            data: {
                content: this._content,
                signature: this._signature
            }
        };
    }
}

export class CompactionChatResponseContentImpl implements CompactionChatResponseContent {
    readonly kind = 'compaction';
    protected _provider: string;
    protected _data: unknown;
    protected _summary?: string;

    constructor(provider: string, data: unknown, summary?: string) {
        this._provider = provider;
        this._data = data;
        this._summary = summary;
    }

    get provider(): string {
        return this._provider;
    }
    get data(): unknown {
        return this._data;
    }
    get summary(): string | undefined {
        return this._summary;
    }

    asString(): string | undefined {
        return undefined;
    }

    asDisplayString(): string | undefined {
        return this._summary;
    }

    toLanguageModelMessage(): CompactionMessage {
        return {
            actor: 'ai',
            type: 'compaction',
            provider: this._provider,
            data: this._data,
            summary: this._summary
        };
    }

    toSerializable(): SerializableChatResponseContentData<CompactionContentData> {
        return {
            kind: 'compaction',
            data: {
                provider: this._provider,
                data: this._data,
                summary: this._summary
            }
        };
    }
}

export class MarkdownChatResponseContentImpl implements MarkdownChatResponseContent {
    readonly kind = 'markdownContent';
    protected _content: MarkdownStringImpl = new MarkdownStringImpl();

    constructor(content: string) {
        this._content.appendMarkdown(content);
    }

    get content(): MarkdownString {
        return this._content;
    }

    asString(): string {
        return this._content.value;
    }

    asDisplayString(): string | undefined {
        return this.asString();
    }

    merge(nextChatResponseContent: MarkdownChatResponseContent): boolean {
        this._content.appendMarkdown(nextChatResponseContent.content.value);
        return true;
    }

    toLanguageModelMessage(): TextMessage {
        return {
            actor: 'ai',
            type: 'text',
            text: this.content.value
        };
    }

    toSerializable(): SerializableChatResponseContentData<MarkdownContentData> {
        return {
            kind: 'markdownContent',
            data: { content: this._content.value }
        };
    }
}

export class InformationalChatResponseContentImpl implements InformationalChatResponseContent {
    readonly kind = 'informational';
    protected _content: MarkdownStringImpl;

    constructor(content: string) {
        this._content = new MarkdownStringImpl(content);
    }

    get content(): MarkdownString {
        return this._content;
    }

    asString(): string | undefined {
        return undefined;
    }

    merge(nextChatResponseContent: InformationalChatResponseContent): boolean {
        this._content.appendMarkdown(nextChatResponseContent.content.value);
        return true;
    }

    toSerializable(): SerializableChatResponseContentData<InformationalContentData> {
        return {
            kind: 'informational',
            data: { content: this._content.value }
        };
    }
}

export class CodeChatResponseContentImpl implements CodeChatResponseContent {
    readonly kind = 'code';
    protected _code: string;
    protected _language?: string;
    protected _location?: Location;

    constructor(code: string, language?: string, location?: Location) {
        this._code = code;
        this._language = language;
        this._location = location;
    }

    get code(): string {
        return this._code;
    }

    get language(): string | undefined {
        return this._language;
    }

    get location(): Location | undefined {
        return this._location;
    }

    asString(): string {
        return `\`\`\`${this._language ?? ''}\n${this._code}\n\`\`\``;
    }

    merge(nextChatResponseContent: CodeChatResponseContent): boolean {
        this._code += `${nextChatResponseContent.code}`;
        return true;
    }

    toSerializable(): SerializableChatResponseContentData<CodeContentData> {
        return {
            kind: 'code',
            data: {
                code: this._code,
                language: this._language,
                location: this._location
            }
        };
    }
}

export class ToolCallChatResponseContentImpl implements ToolCallChatResponseContent, InteractiveContent {
    readonly kind = 'toolCall';
    protected _id?: string;
    protected _name?: string;
    protected _arguments?: string;
    protected _finished?: boolean;
    protected _result?: ToolCallResult;
    protected _data?: Record<string, string>;
    protected _confirmationTimeout?: number;
    protected _needsUserConfirmation: Promise<void>;
    protected _needsUserConfirmationResolver?: () => void;
    protected _isAwaitingUserConfirmation = false;
    protected _isAwaitingUserInput = false;
    protected _confirmed: Promise<boolean>;
    protected _confirmationResolver?: (value: boolean) => void;
    protected _confirmationRejecter?: (reason?: unknown) => void;
    protected _whenFinished: Promise<void>;
    protected _finishedResolver?: () => void;

    protected readonly _onDidChangeEmitter = new Emitter<void>();
    readonly onDidChange: Event<void> = this._onDidChangeEmitter.event;

    constructor(
        id?: string,
        name?: string,
        arg_string?: string,
        finished?: boolean,
        result?: ToolCallResult,
        data?: Record<string, string>
    ) {
        this._id = id;
        this._name = name;
        this._arguments = arg_string;
        this._finished = finished;
        this._result = result;
        this._data = data;
        this._confirmed = this.createConfirmationPromise();
        this._whenFinished = this.createFinishedPromise();
        this._needsUserConfirmation = new Promise<void>(resolve => {
            this._needsUserConfirmationResolver = resolve;
        });
    }

    get id(): string | undefined {
        return this._id;
    }

    get name(): string | undefined {
        return this._name;
    }

    get arguments(): string | undefined {
        return this._arguments;
    }

    get finished(): boolean {
        return this._finished === undefined ? false : this._finished;
    }
    get result(): ToolCallResult | undefined {
        return this._result;
    }

    get data(): Record<string, string> | undefined {
        return this._data;
    }

    get confirmationTimeout(): number | undefined {
        return this._confirmationTimeout;
    }

    set confirmationTimeout(value: number | undefined) {
        this._confirmationTimeout = value;
    }

    get interactionId(): string | undefined {
        return this._id;
    }

    get isResolved(): boolean {
        return this.finished;
    }

    get confirmed(): Promise<boolean> {
        return this._confirmed;
    }

    get needsUserConfirmation(): Promise<void> {
        return this._needsUserConfirmation;
    }

    get isAwaitingUserConfirmation(): boolean {
        return this._isAwaitingUserConfirmation && !this.finished;
    }

    get isAwaitingUserInput(): boolean {
        return this._isAwaitingUserInput && !this.finished;
    }

    get isAwaitingInteraction(): boolean {
        return this.isAwaitingUserConfirmation || this.isAwaitingUserInput;
    }

    get whenFinished(): Promise<void> {
        return this._whenFinished;
    }

    get whenResolved(): Promise<void> {
        return this._whenFinished;
    }

    createConfirmationPromise(): Promise<boolean> {
        if (!this._confirmationResolver) {
            this._confirmed = new Promise<boolean>((resolve, reject) => {
                this._confirmationResolver = resolve;
                this._confirmationRejecter = reject;
            });
        }
        return this._confirmed;
    }

    createFinishedPromise(): Promise<void> {
        if (this._finished) {
            return Promise.resolve();
        }
        if (!this._finishedResolver) {
            this._whenFinished = new Promise<void>(resolve => {
                this._finishedResolver = resolve;
            });
        }
        return this._whenFinished;
    }

    /**
     * Confirm the tool execution
     */
    confirm(): void {
        this._isAwaitingUserConfirmation = false;
        if (this._confirmationResolver) {
            this._confirmationResolver(true);
        }
    }

    deny(reason?: string): void {
        this._isAwaitingUserConfirmation = false;
        if (this._confirmationResolver) {
            this._finished = true;
            this._result = { denied: true, reason };
            this._confirmationResolver(false);
            this.resolveFinished();
        }
    }

    requestUserConfirmation(): void {
        this._isAwaitingUserConfirmation = true;
        if (this._needsUserConfirmationResolver) {
            this._needsUserConfirmationResolver();
            this._needsUserConfirmationResolver = undefined;
        }
    }

    requestUserInput(): void {
        this._isAwaitingUserInput = true;
    }

    userInputHandled(): void {
        this._isAwaitingUserInput = false;
    }

    updateResult(result: ToolCallResult): void {
        this._result = result;
        this._onDidChangeEmitter.fire();
    }

    cancelConfirmation(reason?: unknown): void {
        this._isAwaitingUserConfirmation = false;
        if (this._confirmationRejecter) {
            this._confirmationRejecter(reason);
        }
    }

    complete(result: ToolCallResult): void {
        this._finished = true;
        this._result = result;
        this.resolveFinished();
        this._onDidChangeEmitter.fire();
    }

    protected resolveFinished(): void {
        if (this._finishedResolver) {
            this._finishedResolver();
            this._finishedResolver = undefined;
        }
    }

    asString(): string {
        return '';
    }

    asDisplayString(): string {
        return `Tool call: ${this._name}(${this._arguments ?? ''})`;
    }

    merge(nextChatResponseContent: ToolCallChatResponseContent | ToolCallArgumentsDeltaContent): boolean {
        // Handle argument delta updates
        if (ToolCallArgumentsDeltaContent.is(nextChatResponseContent)) {
            if (nextChatResponseContent.id === this.id) {
                this._arguments = (this._arguments ?? '') + nextChatResponseContent.delta;
                return true;
            }
            return false;
        }

        // Handle full tool call updates
        if (nextChatResponseContent.id === this.id) {
            const wasFinished = this._finished;
            this._finished = nextChatResponseContent.finished;
            this._result = nextChatResponseContent.result;
            const args = nextChatResponseContent.arguments;
            if (args && args.length > 0) {
                this._arguments = args;
            }
            this._data = { ...nextChatResponseContent.data, ...this._data };
            if (!wasFinished && this._finished) {
                this.resolveFinished();
            }
            return true;
        }
        if (nextChatResponseContent.name !== undefined) {
            return false;
        }
        if (nextChatResponseContent.arguments === undefined) {
            return false;
        }
        this._arguments += `${nextChatResponseContent.arguments}`;
        return true;
    }

    protected parseArgumentsSafe(): object {
        try {
            return JSON.parse(this._arguments!);
        } catch (error) {
            console.warn(`Failed to parse tool call arguments for tool '${this._name}': ${error instanceof Error ? error.message : String(error)}`);
            return {};
        }
    }

    toLanguageModelMessage(): [ToolUseMessage, ToolResultMessage] {
        // Format denial results as a human-readable message for the LLM
        let content = this.result;
        if (ToolCallChatResponseContent.isDenialResult(this.result)) {
            content = this.result.reason
                ? `The user denied the tool with reason: ${this.result.reason}.`
                : 'The user denied the tool.';
        }

        return [{
            actor: 'ai',
            type: 'tool_use',
            id: this.id ?? '',
            input: this.arguments && this.arguments.length !== 0 ? this.parseArgumentsSafe() : {},
            name: this.name ?? '',
            data: this.data
        }, {
            actor: 'user',
            type: 'tool_result',
            tool_use_id: this.id ?? '',
            content,
            name: this.name ?? ''
        }];
    }

    toSerializable(): SerializableChatResponseContentData<ToolCallContentData> {
        // `finished` is intentionally not persisted: on restore there is no live handler,
        // so every restored tool call is necessarily finished. Whether the result is
        // partial or final must be encoded in the result itself.
        const data: ToolCallContentData = {
            id: this._id,
            name: this._name,
            arguments: this._arguments,
            result: this._result
        };
        if (this._data && Object.keys(this._data).length > 0) {
            data.data = this._data;
        }
        return { kind: 'toolCall', data };
    }
}

export class ServerToolCallChatResponseContentImpl implements ServerToolCallChatResponseContent {
    readonly kind = 'serverToolCall';
    protected _id?: string;
    protected _name?: string;
    protected _arguments?: string;
    protected _finished: boolean;
    protected _result?: ToolCallResult;
    protected _data?: Record<string, string>;

    constructor(
        id?: string,
        name?: string,
        arg_string?: string,
        finished?: boolean,
        result?: ToolCallResult,
        data?: Record<string, string>
    ) {
        this._id = id;
        this._name = name;
        this._arguments = arg_string;
        this._finished = finished ?? false;
        this._result = result;
        this._data = data;
    }

    get id(): string | undefined {
        return this._id;
    }

    get name(): string | undefined {
        return this._name;
    }

    get arguments(): string | undefined {
        return this._arguments;
    }

    get finished(): boolean {
        return this._finished;
    }

    get result(): ToolCallResult | undefined {
        return this._result;
    }

    get data(): Record<string, string> | undefined {
        return this._data;
    }

    asString(): string {
        return '';
    }

    asDisplayString(): string {
        return `Server tool call: ${this._name}(${this._arguments ?? ''})`;
    }

    merge(nextChatResponseContent: ChatResponseContent): boolean {
        if (!ServerToolCallChatResponseContent.is(nextChatResponseContent) || nextChatResponseContent.id !== this._id) {
            return false;
        }
        this._finished = nextChatResponseContent.finished;
        if (nextChatResponseContent.result !== undefined) {
            this._result = nextChatResponseContent.result;
        }
        const args = nextChatResponseContent.arguments;
        if (args && args.length > 0) {
            this._arguments = args;
        }
        if (nextChatResponseContent.name) {
            this._name = nextChatResponseContent.name;
        }
        this._data = { ...nextChatResponseContent.data, ...this._data };
        return true;
    }

    protected parseArgumentsSafe(): object {
        try {
            return JSON.parse(this._arguments!);
        } catch (error) {
            console.warn(`Failed to parse server tool call arguments for tool '${this._name}': ${error instanceof Error ? error.message : String(error)}`);
            return {};
        }
    }

    toLanguageModelMessage(): ServerToolUseMessage {
        return {
            actor: 'ai',
            type: 'server_tool_use',
            id: this._id ?? '',
            name: this._name ?? '',
            input: this._arguments && this._arguments.length !== 0 ? this.parseArgumentsSafe() : {},
            result: this._result,
            data: this._data
        };
    }

    toSerializable(): SerializableChatResponseContentData<ServerToolCallContentData> {
        // Like tool calls, `finished` is not persisted: a restored server tool call is always finished.
        const data: ServerToolCallContentData = {
            id: this._id,
            name: this._name,
            arguments: this._arguments,
            result: this._result
        };
        if (this._data && Object.keys(this._data).length > 0) {
            data.data = this._data;
        }
        return { kind: 'serverToolCall', data };
    }
}

export const COMMAND_CHAT_RESPONSE_COMMAND: Command = {
    id: 'ai-chat.command-chat-response.generic'
};
export class CommandChatResponseContentImpl implements CommandChatResponseContent {
    readonly kind = 'command';

    constructor(public command?: Command, public customCallback?: CustomCallback, protected args?: unknown[]) { }

    get arguments(): unknown[] {
        return this.args ?? [];
    }

    asString(): string {
        return this.command?.id || this.customCallback?.label || 'command';
    }

    toSerializable(): SerializableChatResponseContentData<CommandContentData> {
        return {
            kind: 'command',
            data: {
                commandId: this.command?.id,
                commandLabel: this.customCallback?.label,
                arguments: this.args
            }
        };
    }
}

export class HorizontalLayoutChatResponseContentImpl implements HorizontalLayoutChatResponseContent {
    readonly kind = 'horizontal';
    protected _content: ChatResponseContent[];

    constructor(content: ChatResponseContent[] = []) {
        this._content = content;
    }

    get content(): ChatResponseContent[] {
        return this._content;
    }

    asString(): string {
        return this._content.map(child => child.asString && child.asString()).join(' ');
    }

    asDisplayString(): string | undefined {
        return this.asString();
    }

    merge(nextChatResponseContent: ChatResponseContent): boolean {
        if (HorizontalLayoutChatResponseContent.is(nextChatResponseContent)) {
            this._content.push(...nextChatResponseContent.content);
        } else {
            this._content.push(nextChatResponseContent);
        }
        return true;
    }

    toSerializable(): SerializableChatResponseContentData<HorizontalLayoutContentData> {
        return {
            kind: 'horizontal',
            data: {
                content: this._content.map(child => {
                    const serialized = child.toSerializable?.();
                    if (!serialized) {
                        return {
                            kind: child.kind,
                            fallbackMessage: child.asString?.(),
                            data: undefined
                        };
                    }
                    return {
                        ...serialized,
                        fallbackMessage: child.asString?.()
                    };
                })
            }
        };
    }
}

/**
 * Options bag for constructing a {@link QuestionResponseContentImpl}.
 */
export interface QuestionResponseContentOptions {
    selectedOption?: { text: string; value?: string };
    selectedOptions?: { text: string; value?: string }[];
    multiSelect?: boolean;
    header?: string;
    /** Called when the user dismisses a single-select question without choosing an option. */
    onSkip?: () => void;
}

/**
 * Default implementation for the QuestionResponseContent.
 * Can be created with or without handler/request for read-only (restored) mode.
 */
export class QuestionResponseContentImpl implements QuestionResponseContent, InteractiveContent {
    readonly kind = 'question';
    public multiSelect?: boolean;
    public header?: string;
    public onSkip?: () => void;
    protected _selectedOptions: { text: string; value?: string }[] | undefined;
    protected _resolvedResolver?: () => void;
    readonly whenResolved: Promise<void>;

    constructor(
        question: string,
        options: { text: string, value?: string, description?: string }[],
        request: MutableChatRequestModel | undefined,
        handler: QuestionResponseHandler | undefined,
        questionOptions?: Omit<QuestionResponseContentOptions, 'multiSelect'> & { multiSelect?: false }
    );
    constructor(
        question: string,
        options: { text: string, value?: string, description?: string }[],
        request: MutableChatRequestModel | undefined,
        handler: MultiSelectQuestionResponseHandler | undefined,
        questionOptions: QuestionResponseContentOptions & { multiSelect: true }
    );
    constructor(
        public question: string,
        public options: { text: string, value?: string, description?: string }[],
        public request: MutableChatRequestModel | undefined,
        public handler: QuestionResponseHandler | MultiSelectQuestionResponseHandler | undefined,
        questionOptions?: QuestionResponseContentOptions
    ) {
        this.multiSelect = questionOptions?.multiSelect;
        this.header = questionOptions?.header;
        this.onSkip = questionOptions?.onSkip;
        this._selectedOptions = questionOptions?.selectedOptions ??
            (questionOptions?.selectedOption ? [questionOptions.selectedOption] : undefined);
        if (this._selectedOptions) {
            this.whenResolved = Promise.resolve();
        } else {
            this.whenResolved = new Promise<void>(resolve => {
                this._resolvedResolver = resolve;
            });
        }
        if (!this.isReadOnly && this.request) {
            this.request.response.fireInteractionNeeded(this);
        }
    }

    get isReadOnly(): boolean {
        return !this.handler || !this.request;
    }

    get interactionId(): string | undefined {
        return `question-${this.question}`;
    }

    get isResolved(): boolean {
        return this.selectedOption !== undefined;
    }

    get isAwaitingInteraction(): boolean {
        // A skipped question resolves with an empty selection: isResolved stays false,
        // but the question no longer awaits interaction.
        return !this.isReadOnly && this._selectedOptions === undefined;
    }

    set selectedOption(option: { text: string; value?: string } | undefined) {
        this._selectedOptions = option ? [option] : undefined;
        this._resolvedResolver?.();
        if (this.request) {
            this.request.response.response.responseContentChanged();
        }
    }
    get selectedOption(): { text: string; value?: string } | undefined {
        return this._selectedOptions?.[0];
    }

    set selectedOptions(options: { text: string; value?: string }[] | undefined) {
        this._selectedOptions = options;
        this._resolvedResolver?.();
        if (this.request) {
            this.request.response.response.responseContentChanged();
        }
    }
    get selectedOptions(): { text: string; value?: string }[] | undefined {
        return this._selectedOptions;
    }

    asString?(): string | undefined {
        const answer = this._selectedOptions && this._selectedOptions.length > 0
            ? `Answer: ${this._selectedOptions.map(o => o.text).join(', ')}`
            : 'No answer';
        return `Question: ${this.question}\n${answer}`;
    }
    merge?(): boolean {
        return false;
    }
    toSerializable(): SerializableChatResponseContentData<QuestionContentData> {
        const data: QuestionContentData = {
            question: this.question,
            options: this.options,
        };
        if (this.multiSelect) {
            data.selectedOptions = this._selectedOptions;
        } else if (this._selectedOptions?.[0]) {
            data.selectedOption = this._selectedOptions[0];
        }
        if (this.header !== undefined) {
            data.header = this.header;
        }
        if (this.multiSelect !== undefined) {
            data.multiSelect = this.multiSelect;
        }
        return {
            kind: 'question',
            data
        };
    }
}

class ChatResponseImpl implements ChatResponse {
    protected readonly _onDidChangeEmitter = new Emitter<void>();
    onDidChange: Event<void> = this._onDidChangeEmitter.event;
    protected _content: ChatResponseContent[];
    protected _responseRepresentation: string;
    protected _responseRepresentationForDisplay: string;
    protected readonly contentChangeListeners = new Map<ChatResponseContent, Disposable>();

    constructor() {
        this._content = [];
    }

    get content(): ChatResponseContent[] {
        return this._content;
    }

    clearContent(): void {
        this.contentChangeListeners.forEach(listener => listener.dispose());
        this.contentChangeListeners.clear();
        this._content = [];
        this._updateResponseRepresentation();
        this._onDidChangeEmitter.fire();
    }

    addContents(contents: ChatResponseContent[]): void {
        contents.forEach(c => this.doAddContent(c));
        this._onDidChangeEmitter.fire();
    }

    addContent(nextContent: ChatResponseContent): void {
        // TODO: Support more complex merges affecting different content than the last, e.g. via some kind of ProcessorRegistry
        // TODO: Support more of the built-in VS Code behavior, see
        //   https://github.com/microsoft/vscode/blob/a2cab7255c0df424027be05d58e1b7b941f4ea60/src/vs/workbench/contrib/chat/common/chatModel.ts#L188-L244
        this.doAddContent(nextContent);
        this._onDidChangeEmitter.fire();
    }

    protected doAddContent(nextContent: ChatResponseContent): void {
        if (ToolCallArgumentsDeltaContent.is(nextContent)) {
            // Delta content targets an existing tool call by ID
            const targetTool = this._content.find(c => ToolCallChatResponseContent.is(c) && c.id === nextContent.id);
            if (targetTool !== undefined && ChatResponseContent.hasMerge(targetTool)) {
                targetTool.merge(nextContent);
            }
            // If no matching tool call found, silently drop the delta (the tool call might not exist yet)
        } else if (ToolCallChatResponseContent.is(nextContent) && nextContent.id !== undefined) {
            const fittingTool = this._content.find(c => ToolCallChatResponseContent.is(c) && c.id === nextContent.id);
            if (fittingTool !== undefined) {
                fittingTool.merge?.(nextContent);
            } else {
                this._content.push(nextContent);
                // Forward content-level change events (e.g. partial-result updates from a
                // renderer) so auto-save can persist them. Without this, mutations that
                // don't go through addContent/merge are invisible to listeners.
                // The subscription is tracked so that clearContent() can dispose it: the stream
                // parser clears and re-adds the content per token, which would otherwise stack
                // up one listener per token on the same content object (#17858).
                this.contentChangeListeners.get(nextContent)?.dispose();
                this.contentChangeListeners.set(nextContent, nextContent.onDidChange(() => this._onDidChangeEmitter.fire()));
            }
        } else if (ServerToolCallChatResponseContent.is(nextContent) && nextContent.id !== undefined) {
            // Server tool calls are matched by id (the start and result blocks arrive as separate stream parts).
            const fittingTool = this._content.find(c => ServerToolCallChatResponseContent.is(c) && c.id === nextContent.id);
            if (fittingTool !== undefined && ChatResponseContent.hasMerge(fittingTool)) {
                fittingTool.merge(nextContent);
            } else {
                this._content.push(nextContent);
            }
        } else {
            const lastElement = this._content.length > 0
                ? this._content[this._content.length - 1]
                : undefined;
            if (lastElement?.kind === nextContent.kind && ChatResponseContent.hasMerge(lastElement)) {
                const mergeSuccess = lastElement.merge(nextContent);
                if (!mergeSuccess) {
                    this._content.push(nextContent);
                }
            } else {
                this._content.push(nextContent);
            }
        }
        this._updateResponseRepresentation();
    }

    responseContentChanged(): void {
        this._updateResponseRepresentation();
        this._onDidChangeEmitter.fire();
    }

    protected _updateResponseRepresentation(): void {
        this._responseRepresentation = this.responseRepresentationsToString(this._content, 'asString');
        this._responseRepresentationForDisplay = this.responseRepresentationsToString(this.content, 'asDisplayString');
    }

    protected responseRepresentationsToString(content: ChatResponseContent[], collect: 'asString' | 'asDisplayString'): string {
        return content
            .map(responseContent => {
                if (collect === 'asDisplayString') {
                    if (ChatResponseContent.hasDisplayString(responseContent)) {
                        return responseContent.asDisplayString();
                    }
                }
                if (ChatResponseContent.hasAsString(responseContent)) {
                    return responseContent.asString();
                }
                if (TextChatResponseContent.is(responseContent)) {
                    return responseContent.content;
                }
                console.warn(
                    'Was not able to map responseContent to a string',
                    responseContent
                );
                return undefined;
            })
            .filter(text => (text !== undefined && text !== ''))
            .join('\n\n');
    }

    asString(): string {
        return this._responseRepresentation;
    }

    asDisplayString(): string {
        return this._responseRepresentationForDisplay;
    }
}

export class MutableChatResponseModel implements ChatResponseModel {
    protected readonly _onDidChangeEmitter = new Emitter<void>();
    onDidChange: Event<void> = this._onDidChangeEmitter.event;

    protected readonly _onInteractionNeededEmitter = new Emitter<InteractiveContent & ChatResponseContent>();
    readonly onInteractionNeeded: Event<InteractiveContent & ChatResponseContent> = this._onInteractionNeededEmitter.event;

    protected _pendingInteractions: (InteractiveContent & ChatResponseContent)[] = [];

    data = {};

    protected _id: string;
    protected _requestId: string;
    protected _progressMessages: ChatProgressMessage[];
    protected _response: ChatResponseImpl;
    protected _isComplete: boolean;
    // Reference count of in-flight requests for user input (agent questions and pending tool-call
    // confirmations). The response is "waiting for input" while this is greater than zero, so that
    // multiple parallel confirmations all have to resolve before the waiting state clears.
    protected _waitingForInputCount: number = 0;
    protected _agentId?: string;
    protected _isError: boolean;
    protected _errorObject: Error | undefined;
    protected _cancellationToken: CancellationTokenSource;
    protected _promptVariantId?: string;
    protected _isPromptVariantEdited?: boolean;
    protected _languageModel?: string;
    protected _tokenUsage?: ResponseTokenUsage;
    protected _tokenUsageEntries: ResponseTokenUsage[] = [];

    constructor(
        requestId: string,
        agentId?: string,
        serializedData?: SerializableChatResponseData
    ) {
        this._requestId = requestId;
        this._agentId = agentId;
        this._cancellationToken = new CancellationTokenSource();

        // Check if we're restoring from serialized data
        if (serializedData) {
            this.restoreFromSerializedData(serializedData);
        } else {
            // Normal creation path
            this._id = generateUuid();
            this._progressMessages = [];
            this._isComplete = false;
            this._waitingForInputCount = 0;
            this._isError = false;
        }

        const response = new ChatResponseImpl();
        response.onDidChange(() => this._onDidChangeEmitter.fire());
        this._response = response;
    }

    /**
     * Restore this response model from serialized data
     */
    protected restoreFromSerializedData(data: SerializableChatResponseData): void {
        this._id = data.id;
        // Always mark restored responses as complete since there's no active agent
        this._isComplete = true;
        this._isError = data.isError;

        // Do not restore waitingForInput state - when a session is restored,
        // the agent that was waiting for input is no longer running
        this._waitingForInputCount = 0;
        // TODO: Restore progressMessages?
        this._progressMessages = [];
        this._promptVariantId = data.promptVariantId;
        this._isPromptVariantEdited = data.isPromptVariantEdited ?? false;
        this._languageModel = data.languageModel;
        this._tokenUsage = data.tokenUsage;

        if (data.errorMessage) {
            this._errorObject = new Error(data.errorMessage);
        }

        // Note: Content restoration will be handled by ChatService using deserializer registry
    }

    get id(): string {
        return this._id;
    }

    get requestId(): string {
        return this._requestId;
    }

    get progressMessages(): ChatProgressMessage[] {
        return this._progressMessages;
    }

    addProgressMessage(message: { content: string } & Partial<Omit<ChatProgressMessage, 'kind'>>): ChatProgressMessage {
        const id = message.id ?? generateUuid();
        const existingMessage = this.getProgressMessage(id);
        if (existingMessage) {
            this.updateProgressMessage({ id, ...message });
            return existingMessage;
        }
        const newMessage: ChatProgressMessage = {
            kind: 'progressMessage',
            id,
            status: message.status ?? 'inProgress',
            show: message.show ?? 'untilFirstContent',
            ...message,
        };
        this._progressMessages.push(newMessage);
        this._onDidChangeEmitter.fire();
        return newMessage;
    }

    getProgressMessage(id: string): ChatProgressMessage | undefined {
        return this._progressMessages.find(message => message.id === id);
    }

    updateProgressMessage(message: { id: string } & Partial<Omit<ChatProgressMessage, 'kind'>>): void {
        const progressMessage = this.getProgressMessage(message.id);
        if (progressMessage) {
            Object.assign(progressMessage, message);
            this._onDidChangeEmitter.fire();
        }
    }

    get response(): ChatResponseImpl {
        return this._response;
    }

    get isComplete(): boolean {
        return this._isComplete;
    }

    get isCanceled(): boolean {
        return this._cancellationToken.token.isCancellationRequested;
    }

    get isWaitingForInput(): boolean {
        return this._waitingForInputCount > 0;
    }

    get agentId(): string | undefined {
        return this._agentId;
    }

    get promptVariantId(): string | undefined {
        return this._promptVariantId;
    }

    get isPromptVariantEdited(): boolean {
        return this._isPromptVariantEdited ?? false;
    }

    get languageModel(): string | undefined {
        return this._languageModel;
    }

    get tokenUsage(): ResponseTokenUsage | undefined {
        return this._tokenUsage;
    }

    setTokenUsage(usage: ResponseTokenUsage): void {
        this._tokenUsage = usage;
        this.addTokenUsageEntry(usage);
        this._onDidChangeEmitter.fire();
    }

    addTokenUsageEntry(usage: ResponseTokenUsage): void {
        this._tokenUsageEntries.push(usage);
    }

    get tokenUsageEntries(): readonly ResponseTokenUsage[] {
        return this._tokenUsageEntries;
    }

    setPromptVariantInfo(variantId: string | undefined, isEdited: boolean): void {
        this._promptVariantId = variantId;
        this._isPromptVariantEdited = isEdited;
        this._onDidChangeEmitter.fire();
    }

    /** Records the identifier of the language model that produced this response. */
    setLanguageModel(languageModel: string | undefined): void {
        this._languageModel = languageModel;
        this._onDidChangeEmitter.fire();
    }

    overrideAgentId(agentId: string): void {
        this._agentId = agentId;
    }

    /** Reset all pending-input state; the response no longer accepts user interaction. */
    protected resetPendingInput(): void {
        this._waitingForInputCount = 0;
        this._pendingInteractions = [];
    }

    complete(): void {
        this._isComplete = true;
        this.resetPendingInput();
        this._onDidChangeEmitter.fire();
    }

    cancel(): void {
        this._cancellationToken.cancel();
        this._isComplete = true;
        this.resetPendingInput();

        // Ensure any pending tool confirmations are canceled when the chat is canceled
        try {
            const content = this._response.content;
            for (const item of content) {
                if (ToolCallChatResponseContent.is(item)) {
                    item.cancelConfirmation(new Error('Chat request canceled'));
                }
            }
        } catch (e) {
            // best-effort: ignore errors while canceling confirmations
        }

        this._onDidChangeEmitter.fire();
    }

    get cancellationToken(): CancellationToken {
        return this._cancellationToken.token;
    }

    waitForInput(): void {
        this._waitingForInputCount++;
        this._onDidChangeEmitter.fire();
    }

    stopWaitingForInput(): void {
        if (this._waitingForInputCount > 0) {
            this._waitingForInputCount--;
        }
        this._onDidChangeEmitter.fire();
    }

    get pendingInteractions(): ReadonlyArray<InteractiveContent & ChatResponseContent> {
        // Filter by the momentary awaiting state so consumers never see interactions
        // whose actionable phase has passed (e.g. a confirmed tool call that is still
        // executing and therefore not yet resolved).
        return this._pendingInteractions.filter(part => part.isAwaitingInteraction ?? !part.isResolved);
    }

    fireInteractionNeeded(contentPart: InteractiveContent & ChatResponseContent): void {
        if (!this._isComplete && !contentPart.isResolved && !this._pendingInteractions.includes(contentPart)) {
            this._pendingInteractions = [...this._pendingInteractions, contentPart];
            // Stop tracking on settlement either way: custom InteractiveContent implementations
            // may reject whenResolved, and the rejection must not escape the model unhandled.
            const stopTracking = (): void => {
                this._pendingInteractions = this._pendingInteractions.filter(part => part !== contentPart);
            };
            contentPart.whenResolved.then(stopTracking, stopTracking);
        }
        this._onInteractionNeededEmitter.fire(contentPart);
    }

    notifyChanged(): void {
        this._onDidChangeEmitter.fire();
    }

    error(error: Error): void {
        this._isComplete = true;
        this.resetPendingInput();
        this._isError = true;
        this._errorObject = error;
        this._onDidChangeEmitter.fire();
    }
    get errorObject(): Error | undefined {
        return this._errorObject;
    }
    get isError(): boolean {
        return this._isError;
    }

    toSerializable(): SerializableChatResponseData {
        return {
            id: this.id,
            requestId: this.requestId,
            isComplete: this.isComplete,
            isError: this.isError,
            errorMessage: this.errorObject?.message,
            promptVariantId: this._promptVariantId,
            isPromptVariantEdited: this._isPromptVariantEdited,
            languageModel: this._languageModel,
            tokenUsage: this._tokenUsage,
            content: this.response.content.map(c => {
                const serialized = c.toSerializable?.();
                if (!serialized) {
                    // Fallback if toSerializable not implemented
                    return {
                        kind: c.kind,
                        fallbackMessage: c.asString?.(),
                        data: undefined
                    };
                }
                return {
                    ...serialized,
                    fallbackMessage: c.asString?.()
                };
            })
        };
    }
}

export class ErrorChatResponseModel extends MutableChatResponseModel {
    constructor(requestId: string, error: Error, agentId?: string) {
        super(requestId, agentId);
        this.error(error);
    }
}

export class ProgressChatResponseContentImpl implements ProgressChatResponseContent {
    readonly kind = 'progress';
    protected _message: string;

    constructor(message: string) {
        this._message = message;
    }

    get message(): string {
        return this._message;
    }

    asString(): string {
        return JSON.stringify({
            type: 'progress',
            message: this.message
        });
    }

    asDisplayString(): string | undefined {
        return `<Progress>${this.message}</Progress>`;
    }

    merge(nextChatResponseContent: ProgressChatResponseContent): boolean {
        this._message = nextChatResponseContent.message;
        return true;
    }

    toLanguageModelMessage(): TextMessage {
        return {
            actor: 'ai',
            type: 'text',
            text: this.message
        };
    }
    toSerializable(): SerializableChatResponseContentData<ProgressContentData> {
        return {
            kind: 'progress',
            data: { message: this._message }
        };
    }
}

/**
 * Fallback content for unknown content types.
 * Used when a deserializer is not available (e.g., content from removed extension).
 */
export interface UnknownChatResponseContent extends ChatResponseContent {
    kind: 'unknown';
    originalKind: string;
    fallbackMessage?: string;
    data: unknown;
}

export class UnknownChatResponseContentImpl implements UnknownChatResponseContent {
    readonly kind = 'unknown';

    constructor(
        public readonly originalKind: string,
        public readonly fallbackMessage: string | undefined,
        public readonly data: unknown
    ) { }

    asString(): string | undefined {
        return this.fallbackMessage;
    }

    toSerializable(): SerializableChatResponseContentData {
        return {
            kind: 'unknown',
            data: this.data
        };
    }
}