// *****************************************************************************
// Copyright (C) 2025 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
// *****************************************************************************

import { inject, injectable, named } from '@theia/core/shared/inversify';
import { AGENT_DELEGATION_FUNCTION_ID, ToolInvocationContext, ToolProvider, ToolRequest } from '@theia/ai-core';
import { Disposable, ILogger } from '@theia/core';
import {
    assertChatContext,
    ChatAgentService,
    ChatAgentServiceFactory,
    ChatChangeEvent,
    ChatRequest,
    ChatResponseContent,
    ChatService,
    ChatServiceFactory,
    ChatToolContext,
    MutableChatModel,
    MutableChatRequestModel,
    MutableChatResponseModel,
    ChatRequestInvocation,
    ThinkingChatResponseContent,
    ToolCallChatResponseContent,
} from '../common';
import { TASK_CONTEXT_VARIABLE } from './task-context-variable';

@injectable()
export class AgentDelegationTool implements ToolProvider {
    static ID = AGENT_DELEGATION_FUNCTION_ID;

    protected readonly pendingDelegations = new Map<string, { prompt: string; invocation: ChatRequestInvocation }>();

    @inject(ChatAgentServiceFactory)
    protected readonly getChatAgentService: () => ChatAgentService;

    @inject(ChatServiceFactory)
    protected readonly getChatService: () => ChatService;

    @inject(ILogger) @named('ai-chat:AgentDelegationTool')
    protected readonly logger: ILogger;

    getTool(): ToolRequest {
        return {
            id: AgentDelegationTool.ID,
            name: AgentDelegationTool.ID,
            description:
                'Delegate a task or question to a specific AI agent. IMPORTANT: When you delegate a task or question to a specific AI agent using this tool, ' +
                'remember that each sub-agent operates solely within its specialized capabilities and tools and does not have access to previous conversation context ' +
                ' or external systems. Therefore, it is crucial to provide all necessary context and detailed information directly within your request to ensure accurate ' +
                'and effective task completion. ' +
                'You may optionally pass a taskContextId to make a specific task context (e.g. a plan) available to the delegated agent via its system prompt.',
            parameters: {
                type: 'object',
                properties: {
                    agentId: {
                        type: 'string',
                        description:
                            'The ID of the AI agent to delegate the task to.',
                    },
                    prompt: {
                        type: 'string',
                        description:
                            'The task, question, or prompt to pass to the specified agent.',
                    },
                    taskContextId: {
                        type: 'string',
                        description: 'Optional task context ID to make available to the delegated agent. The agent will see the task context in its system prompt.',
                    },
                },
                required: ['agentId', 'prompt'],
            },
            handler: (arg_string: string, ctx?: ToolInvocationContext) => {
                assertChatContext(ctx);
                return this.delegateToAgent(arg_string, ctx);
            },
        };
    }

    getDelegation(toolCallId: string): { prompt: string; invocation: ChatRequestInvocation } | undefined {
        return this.pendingDelegations.get(toolCallId);
    }

    private async delegateToAgent(
        arg_string: string,
        ctx: ChatToolContext
    ): Promise<string> {
        if (ctx.cancellationToken?.isCancellationRequested) {
            return 'Operation cancelled by user';
        }

        try {
            const args = JSON.parse(arg_string);
            const { agentId, prompt, taskContextId } = args;

            if (!agentId || !prompt) {
                const errorMsg = 'Both agentId and prompt parameters are required.';
                this.logger.error(errorMsg, { agentId, prompt });
                return errorMsg;
            }

            // Check if the specified agent exists
            const agent = this.getChatAgentService().getAgent(agentId, true);
            if (!agent) {
                const availableAgents = this.getChatAgentService()
                    .getAgents(true)
                    .map(a => a.id);
                const errorMsg = `Agent '${agentId}' not found or not enabled. Available agents: ${availableAgents.join(', ')}`;
                this.logger.error(errorMsg);
                return errorMsg;
            }

            let newSession;
            let childModelDisposable: Disposable | undefined;
            try {
                const chatService = this.getChatService();

                // Store the current active session to restore it after delegation
                const currentActiveSession = chatService.getActiveSession();

                newSession = chatService.createSession(
                    undefined,
                    { focus: false },
                    agent
                );

                // Set root session ID to enable task context sharing across delegation chains
                // Root is either the current root (for nested delegation) or current session (for first-level delegation)
                const rootId = ctx.rootSessionId || ctx.request.session.id;
                newSession.rootSessionId = rootId;
                newSession.model.rootSessionId = rootId;

                // Track the immediate parent (the delegating session) so the UI can render the full
                // delegation hierarchy, not just a flat list under the root.
                const parentId = ctx.request.session.id;
                newSession.parentSessionId = parentId;
                newSession.model.parentSessionId = parentId;

                if (taskContextId) {
                    newSession.model.context.addVariables({
                        variable: TASK_CONTEXT_VARIABLE,
                        arg: taskContextId
                    });
                }

                // Immediately restore the original active session to avoid confusing the user
                if (currentActiveSession) {
                    chatService.setActiveSession(currentActiveSession.id, { focus: false });
                }

                // Setup bubbling of child session events to parent session
                childModelDisposable = this.setupChildSessionBubbling(newSession.model as MutableChatModel, ctx.request.session, ctx.response);
            } catch (sessionError) {
                const errorMsg = `Failed to create chat session for agent '${agentId}': ${sessionError instanceof Error ? sessionError.message : sessionError}`;
                this.logger.error(errorMsg, sessionError);
                return errorMsg;
            }

            // Send the request
            const chatRequest: ChatRequest = {
                text: `@${agentId} ${prompt}`,
            };

            let response: ChatRequestInvocation | undefined;
            try {
                if (ctx.cancellationToken?.isCancellationRequested) {
                    return 'Operation cancelled by user';
                }

                const chatService = this.getChatService();
                response = await chatService.sendRequest(
                    newSession.id,
                    chatRequest
                );

                if (ctx.cancellationToken) {
                    ctx.cancellationToken.onCancellationRequested(
                        async () => {
                            if (response) {
                                ((await response?.requestCompleted) as MutableChatRequestModel).cancel();
                            }
                        }
                    );
                }
            } catch (sendError) {
                const errorMsg = `Failed to send request to agent '${agentId}': ${sendError instanceof Error ? sendError.message : sendError}`;
                this.logger.error(errorMsg, sendError);
                return errorMsg;
            }

            if (response) {
                // Store the invocation in the registry so the renderer can access it
                if (ctx.toolCallId) {
                    this.pendingDelegations.set(ctx.toolCallId, {
                        prompt,
                        invocation: response
                    });
                    // Clean up when the delegated session is deleted
                    const chatService = this.getChatService();
                    const toolCallId = ctx.toolCallId;
                    const sessionEventListener = chatService.onSessionEvent(event => {
                        if (event.type === 'deleted' && event.sessionId === newSession.id) {
                            this.pendingDelegations.delete(toolCallId);
                            sessionEventListener.dispose();
                        }
                    });
                }

                try {
                    // Wait for completion to return the final result as tool output
                    const result = await response.responseCompleted;
                    const filteredContent = result.response.content.filter(c => !ThinkingChatResponseContent.is(c) && !ToolCallChatResponseContent.is(c));
                    const stringResult = filteredContent
                        .map(c => ChatResponseContent.hasAsString(c) ? c.asString() : undefined)
                        .filter((text): text is string => text !== undefined && text !== '')
                        .join('\n\n');

                    // Clean up the parent-child link after completion (event bubbling is no longer needed)
                    childModelDisposable?.dispose();

                    // Return the raw text to the top-level Agent, as a tool result
                    return stringResult;
                } catch (completionError) {
                    if (
                        completionError instanceof Error &&
                        completionError.message.includes('cancelled')
                    ) {
                        return 'Operation cancelled by user';
                    }
                    const errorMsg = `Failed to complete response from agent '${agentId}': ${completionError instanceof Error ? completionError.message : completionError}`;
                    this.logger.error(errorMsg, completionError);
                    return errorMsg;
                }
            } else {
                const errorMsg = `Delegation to agent '${agentId}' has failed: no response returned.`;
                this.logger.error(errorMsg);
                return errorMsg;
            }
        } catch (error) {
            this.logger.error('Failed to delegate to agent', error);
            return JSON.stringify({
                error: `Failed to parse arguments or delegate to agent: ${error instanceof Error ? error.message : 'Unknown error'}`,
            });
        }
    }

    /**
     * Sets up all event bubbling from a delegated child session to the parent session:
     * - Interaction forwarding: child interactionNeeded events are forwarded to the parent response
     * - ChangeSet bubbling: child changeset changes are forwarded to the parent model
     */
    private setupChildSessionBubbling(
        childModel: MutableChatModel,
        parentModel: MutableChatModel,
        parentResponse: MutableChatResponseModel
    ): Disposable {

        // Forward interactionNeeded events to the parent response model
        // so the UI (which subscribes to response.onInteractionNeeded) can display them.
        // Also watch for each forwarded interaction's resolution to trigger cleanup.
        const eventForwarding = childModel.onDidChange(event => {
            if (ChatChangeEvent.isInteractionNeededEvent(event)) {
                parentResponse.fireInteractionNeeded(event.contentPart);
                event.contentPart.whenResolved.then(() => parentResponse.notifyChanged());
            }
        });

        // Bubble ChangeSet changes to the parent model
        const changeSetForwarding = childModel.changeSet.onDidChange(() => {
            const delegatedElements = childModel.changeSet.getElements();
            if (delegatedElements.length > 0) {
                parentModel.changeSet.setTitle(childModel.changeSet.title);
                parentModel.changeSet.addElements(...delegatedElements);
            }
        });

        return {
            dispose: () => {
                eventForwarding.dispose();
                changeSetForwarding.dispose();
            }
        };
    }
}
