// *****************************************************************************
// 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/chatRequestParser.ts

import { inject, injectable, named } from '@theia/core/shared/inversify';
import { ChatAgentService } from './chat-agent-service';
import { ChatAgentLocation } from './chat-agents';
import { ChatContext, ChatRequest } from './chat-model';
import {
    chatAgentLeader,
    chatFunctionLeader,
    ParsedChatRequestAgentPart,
    ParsedChatRequestFunctionPart,
    ParsedChatRequestTextPart,
    ParsedChatRequestVariablePart,
    chatVariableLeader,
    chatSubcommandLeader,
    OffsetRange,
    ParsedChatRequest,
    ParsedChatRequestPart,
} from './parsed-chat-request';
import {
    AIVariable, AIVariableService, createAIResolveVariableCache, getAllResolvedAIVariables, parseFunctionReference, PromptService, ToolInvocationRegistry, ToolRequest
} from '@theia/ai-core';
import { ILogger } from '@theia/core';

const agentReg = /^@([\w_\-\.]+)(?=(\s|$|\b))/i; // An @-agent
// A ~ tool function. The optional `?` prefix marks the tool as deferred,
// e.g. `~?functionId` (chat format) or `~{?functionId}` (prompt format).
const functionReg = /^~(\??[\w_\-\.]+)(?=(\s|$|\b))/i;
const functionPromptFormatReg = /^\~\{\s*(.*?)\s*\}/i;
const variableReg = /^#([\w_\-]+)(?::([\w_\-_\/\\.:]+))?(?=(\s|$|\b))/i; // A #-variable with an optional : arg (#file:workspace/path/name.ext)
// A /-command (/commandname) with optional arguments parsed separately. The command name must be
// terminated by whitespace or the end of the input, so that path segments such as `/home/user` are
// not mistaken for a command.
// The colon and the period are in the charset so a qualified skill is invocable as
// `/<qualifier>:<skill>`; the qualifier comes from a plugin identifier, so it usually has periods.
const commandReg = /^\/([\w_\-.:]+)(?=\s|$)/;
const nextCommandReg = /\s+\/([\w_\-.:]+)(?=\s|$)/g;

export const ChatRequestParser = Symbol('ChatRequestParser');
export interface ChatRequestParser {
    parseChatRequest(request: ChatRequest, location: ChatAgentLocation, context: ChatContext): Promise<ParsedChatRequest>;
}

function offsetRange(start: number, endExclusive: number): OffsetRange {
    if (start > endExclusive) {
        throw new Error(`Invalid range: start=${start} endExclusive=${endExclusive}`);
    }
    return { start, endExclusive };
}
@injectable()
export class ChatRequestParserImpl implements ChatRequestParser {

    @inject(PromptService)
    protected readonly promptService: PromptService;

    constructor(
        @inject(ChatAgentService) private readonly agentService: ChatAgentService,
        @inject(AIVariableService) private readonly variableService: AIVariableService,
        @inject(ToolInvocationRegistry) private readonly toolInvocationRegistry: ToolInvocationRegistry,
        @inject(ILogger) @named('ai-chat:ChatRequestParserImpl')
        protected readonly logger: ILogger
    ) { }

    async parseChatRequest(request: ChatRequest, location: ChatAgentLocation, context: ChatContext): Promise<ParsedChatRequest> {
        // Parse the request into parts
        const { parts, toolRequests, deferredToolIds } = this.parseParts(request, location);

        // Resolve all variables and add them to the variable parts.
        // Parse resolved variable texts again for tool requests.
        // These are not added to parts as they are not visible in the initial chat message.
        // However, they need to be added to the result to be considered by the executing agent.
        const variableCache = createAIResolveVariableCache();
        for (const part of parts) {
            if (part instanceof ParsedChatRequestVariablePart) {
                const resolvedVariable = await this.variableService.resolveVariable(
                    { variable: part.variableName, arg: part.variableArg },
                    context,
                    variableCache
                );
                if (resolvedVariable) {
                    part.resolution = resolvedVariable;
                    // Resolve tool requests in resolved variables
                    this.parseFunctionsFromVariableText(resolvedVariable.value, toolRequests, deferredToolIds);
                } else {
                    this.logger.warn(`Failed to resolve variable ${part.variableName}${part.variableArg ? ':' + part.variableArg : ''} for ${location}`);
                }
            }
        }

        // Get resolved variables from variable cache after all variables have been resolved.
        // We want to return all recursively resolved variables, thus use the whole cache.
        const resolvedVariables = await getAllResolvedAIVariables(variableCache);

        return { request, parts, toolRequests, deferredToolIds, variables: resolvedVariables };
    }

    protected parseParts(request: ChatRequest, location: ChatAgentLocation): {
        parts: ParsedChatRequestPart[];
        toolRequests: Map<string, ToolRequest>;
        deferredToolIds: Set<string>;
        variables: Map<string, AIVariable>;
    } {
        const parts: ParsedChatRequestPart[] = [];
        const variables = new Map<string, AIVariable>();
        const toolRequests = new Map<string, ToolRequest>();
        const deferredToolIds = new Set<string>();
        if (!request.text) {
            return { parts, toolRequests, deferredToolIds, variables };
        }
        const message = request.text;
        for (let i = 0; i < message.length; i++) {
            const previousChar = message.charAt(i - 1);
            const char = message.charAt(i);
            let newPart: ParsedChatRequestPart | undefined;

            if (previousChar.match(/\s/) || i === 0) {
                if (char === chatSubcommandLeader) {
                    // Try to parse as command - commands are syntactic sugar for #prompt:commandName|args
                    const commandPart = this.tryToParseCommand(
                        message.slice(i),
                        i,
                        parts
                    );
                    if (commandPart) {
                        newPart = commandPart;
                        const variable = this.variableService.getVariable(commandPart.variableName);
                        if (variable) {
                            variables.set(variable.name, variable);
                        }
                    }
                } else if (char === chatFunctionLeader) {
                    const functionPart = this.tryToParseFunction(
                        message.slice(i),
                        i
                    );
                    newPart = functionPart;
                    if (functionPart) {
                        toolRequests.set(functionPart.toolRequest.id, functionPart.toolRequest);
                        if (functionPart.deferred) {
                            deferredToolIds.add(functionPart.toolRequest.id);
                        }
                    }
                } else if (char === chatVariableLeader) {
                    const variablePart = this.tryToParseVariable(
                        message.slice(i),
                        i,
                        parts
                    );
                    newPart = variablePart;
                    if (variablePart) {
                        const variable = this.variableService.getVariable(variablePart.variableName);
                        if (variable) {
                            variables.set(variable.name, variable);
                        }
                    }
                } else if (char === chatAgentLeader) {
                    newPart = this.tryToParseAgent(
                        message.slice(i),
                        i,
                        parts,
                        location
                    );
                }
            }

            if (newPart) {
                if (i !== 0) {
                    // Insert a part for all the text we passed over, then insert the new parsed part
                    const previousPart = parts.at(-1);
                    const previousPartEnd = previousPart?.range.endExclusive ?? 0;
                    parts.push(
                        new ParsedChatRequestTextPart(
                            offsetRange(previousPartEnd, i),
                            message.slice(previousPartEnd, i)
                        )
                    );
                }

                parts.push(newPart);
                i = newPart.range.endExclusive - 1;
            }
        }

        const lastPart = parts.at(-1);
        const lastPartEnd = lastPart?.range.endExclusive ?? 0;
        if (lastPartEnd < message.length) {
            parts.push(
                new ParsedChatRequestTextPart(
                    offsetRange(lastPartEnd, message.length),
                    message.slice(lastPartEnd, message.length)
                )
            );
        }
        return { parts, toolRequests, deferredToolIds, variables };
    }

    /**
     * Parse text for tool requests and add them to the given map
     */
    private parseFunctionsFromVariableText(text: string, toolRequests: Map<string, ToolRequest>, deferredToolIds: Set<string>): void {
        for (let i = 0; i < text.length; i++) {
            const char = text.charAt(i);

            // Check for function markers at start of words
            if (char === chatFunctionLeader) {
                const functionPart = this.tryToParseFunction(text.slice(i), i);
                if (functionPart) {
                    // Add the found tool request to the given map
                    toolRequests.set(functionPart.toolRequest.id, functionPart.toolRequest);
                    if (functionPart.deferred) {
                        deferredToolIds.add(functionPart.toolRequest.id);
                    }
                }
            }
        }
    }

    private tryToParseAgent(
        message: string,
        offset: number,
        parts: ReadonlyArray<ParsedChatRequestPart>,
        location: ChatAgentLocation
    ): ParsedChatRequestAgentPart | ParsedChatRequestVariablePart | undefined {
        const nextAgentMatch = message.match(agentReg);
        if (!nextAgentMatch) {
            return;
        }

        const [full, name] = nextAgentMatch;
        const agentRange = offsetRange(offset, offset + full.length);

        let agents = this.agentService.getAgents().filter(a => a.name === name);
        if (!agents.length) {
            const fqAgent = this.agentService.getAgent(name);
            if (fqAgent) {
                agents = [fqAgent];
            }
        }

        // If there is more than one agent with this name, and the user picked it from the suggest widget, then the selected agent should be in the
        // context and we use that one. Otherwise just pick the first.
        const agent = agents[0];
        if (!agent || !agent.locations.includes(location)) {
            return;
        }

        if (parts.some(p => p instanceof ParsedChatRequestAgentPart)) {
            // Only one agent allowed
            return;
        }

        return new ParsedChatRequestAgentPart(agentRange, agent.id, agent.name);
    }

    private tryToParseVariable(
        message: string,
        offset: number,
        _parts: ReadonlyArray<ParsedChatRequestPart>
    ): ParsedChatRequestVariablePart | undefined {
        const nextVariableMatch = message.match(variableReg);
        if (!nextVariableMatch) {
            return;
        }

        const [full, name] = nextVariableMatch;
        const variableArg = nextVariableMatch[2];
        const varRange = offsetRange(offset, offset + full.length);

        return new ParsedChatRequestVariablePart(varRange, name, variableArg);
    }

    /**
     * Try to parse a command at the start of the given message.
     *
     * Commands are syntactic sugar for `#prompt:commandName|args`.
     * The prompt variable resolver will handle the actual resolution.
     */
    protected tryToParseCommand(
        message: string,
        offset: number,
        _parts: ReadonlyArray<ParsedChatRequestPart>
    ): ParsedChatRequestVariablePart | undefined {
        const nextCommandMatch = message.match(commandReg);
        if (!nextCommandMatch) {
            return;
        }

        const [commandText, commandName] = nextCommandMatch;
        if (!this.isCommandCandidate(commandName)) {
            // Not a command we know about, so leave the text alone. Otherwise anything the user
            // types after a `/word` token, e.g. a Unix path, would be swallowed as command
            // arguments and silently dropped when the non-existing command fails to resolve.
            return;
        }
        let commandEnd = commandText.length;
        let commandArgs: string | undefined;

        // Arguments never span multiple lines: a command only consumes the remainder of its own line.
        const lineBreakOffset = message.indexOf('\n', commandEnd);
        const lineEnd = lineBreakOffset === -1 ? message.length : lineBreakOffset;
        const nextCommandOffset = this.findNextCommandOffset(message, commandEnd, lineEnd);
        const argsEnd = nextCommandOffset ?? lineEnd;
        const rawArgs = message.slice(commandEnd, argsEnd);
        const args = rawArgs.trim();
        if (args) {
            commandArgs = args;
            // Advance the consumed range only up to the end of the trimmed argument, not up to the
            // next command. This keeps the whitespace between the argument and a following command
            // out of this part's range, so parseParts emits it as a separator text part. Otherwise
            // the two resolved prompt fragments would run into each other with no space between them.
            commandEnd = argsEnd - (rawArgs.length - rawArgs.trimEnd().length);
        }

        const commandRange = offsetRange(offset, offset + commandEnd);

        const variableArg = commandArgs ? `${commandName}|${commandArgs}` : commandName;
        return new ParsedChatRequestVariablePart(commandRange, 'prompt', variableArg);
    }

    private findNextCommandOffset(message: string, startOffset: number, endOffset: number): number | undefined {
        nextCommandReg.lastIndex = startOffset;
        let match = nextCommandReg.exec(message);
        while (match && match.index < endOffset) {
            if (this.isCommandCandidate(match[1])) {
                return match.index + match[0].indexOf(chatSubcommandLeader);
            }
            match = nextCommandReg.exec(message);
        }
        return undefined;
    }

    /**
     * Whether a `/name` token should be treated as a command. Only names that actually resolve to a
     * command or prompt fragment are accepted, which keeps unknown `/word` tokens (e.g. `/tmp`,
     * `/path/to/file`) as plain text.
     */
    protected isCommandCandidate(commandName: string): boolean {
        return this.promptService.isKnownCommand(commandName);
    }

    private tryToParseFunction(message: string, offset: number): ParsedChatRequestFunctionPart | undefined {
        // Support both the chat and prompt formats for functions
        const nextFunctionMatch = message.match(functionPromptFormatReg) || message.match(functionReg);
        if (!nextFunctionMatch) {
            return;
        }

        const [full, rawId] = nextFunctionMatch;
        const { id, deferred } = parseFunctionReference(rawId);

        const maybeToolRequest = this.toolInvocationRegistry.getFunction(id);
        if (!maybeToolRequest) {
            return;
        }

        const functionRange = offsetRange(offset, offset + full.length);
        return new ParsedChatRequestFunctionPart(functionRange, maybeToolRequest, deferred);
    }
}