import { DisposeAggregator, DisposedError, EventRelay, withLock } from "lifecycle-utils";
import { defaultChatSystemPrompt } from "../../config.js";
import { appendUserMessageToChatHistory } from "../../utils/appendUserMessageToChatHistory.js";
import { LlamaChat } from "../LlamaChat/LlamaChat.js";
export class LlamaChatSession {
    /** @internal */ _disposeAggregator = new DisposeAggregator();
    /** @internal */ _autoDisposeSequence;
    /** @internal */ _contextShift;
    /** @internal */ _chatHistory;
    /** @internal */ _lastEvaluation;
    /** @internal */ _chat;
    onDispose = new EventRelay();
    /**
     * @param options
     */
    constructor({ contextSequence, chatWrapper = "auto", systemPrompt = defaultChatSystemPrompt, autoDisposeSequence = true, contextShift }) {
        if (contextSequence == null)
            throw new Error("contextSequence cannot be null");
        if (contextSequence.disposed)
            throw new DisposedError();
        this._contextShift = contextShift;
        this._chatHistory = [{
                type: "system",
                text: systemPrompt
            }];
        this._chat = new LlamaChat({
            autoDisposeSequence,
            chatWrapper,
            contextSequence
        });
        this._autoDisposeSequence = autoDisposeSequence;
        this._disposeAggregator.add(this._chat.onDispose.createListener(() => {
            this.dispose();
        }));
        this._disposeAggregator.add(this.onDispose.dispatchEvent);
    }
    dispose({ disposeSequence = this._autoDisposeSequence } = {}) {
        if (this._chat == null)
            return;
        this._chat.dispose({ disposeSequence });
        this._chat = null;
        this._disposeAggregator.dispose();
    }
    /** @hidden */
    [Symbol.dispose]() {
        return this.dispose();
    }
    get disposed() {
        return this._chat == null || this._chat.disposed;
    }
    get sequence() {
        if (this._chat == null)
            throw new DisposedError();
        return this._chat.sequence;
    }
    get context() {
        return this.sequence.context;
    }
    get model() {
        return this.sequence.model;
    }
    /**
     * @param prompt
     * @param [options]
     */
    async prompt(prompt, { functions, documentFunctionParams, onToken, signal, maxTokens, temperature, topK, topP, grammar, trimWhitespaceSuffix = false, repeatPenalty } = {}) {
        const { responseText } = await this.promptWithMeta(prompt, {
            // this is a workaround to allow passing both `functions` and `grammar`
            functions: functions,
            documentFunctionParams: documentFunctionParams,
            onToken, signal, maxTokens, temperature, topK, topP, grammar, trimWhitespaceSuffix, repeatPenalty
        });
        return responseText;
    }
    /**
     * @param prompt
     * @param [options]
     */
    async promptWithMeta(prompt, { functions, documentFunctionParams, onToken, signal, maxTokens, temperature, topK, topP, grammar, trimWhitespaceSuffix = false, repeatPenalty, evaluationPriority } = {}) {
        this._ensureNotDisposed();
        if (grammar != null && grammar._llama !== this.model._llama)
            throw new Error("The LlamaGrammar used by passed to this function was created with a different Llama instance than the one used by this sequence's model. Make sure you use the same Llama instance for both the model and the grammar.");
        return await withLock(this, "prompt", async () => {
            this._ensureNotDisposed();
            if (this._chat == null)
                throw new DisposedError();
            let lastEvaluation = this._lastEvaluation;
            let newChatHistory = appendUserMessageToChatHistory(this._chatHistory, prompt);
            let newContextWindowChatHistory = lastEvaluation?.contextWindow == null
                ? undefined
                : appendUserMessageToChatHistory(lastEvaluation?.contextWindow, prompt);
            newChatHistory.push({
                type: "model",
                response: []
            });
            if (newContextWindowChatHistory != null)
                newContextWindowChatHistory.push({
                    type: "model",
                    response: []
                });
            // eslint-disable-next-line no-constant-condition
            while (true) {
                const { functionCall, lastEvaluation: currentLastEvaluation, metadata } = await this._chat.generateResponse(newChatHistory, {
                    functions,
                    documentFunctionParams,
                    grammar: grammar,
                    onToken,
                    signal,
                    repeatPenalty,
                    topK,
                    topP,
                    maxTokens,
                    temperature,
                    trimWhitespaceSuffix,
                    contextShift: {
                        ...this._contextShift,
                        lastEvaluationMetadata: lastEvaluation?.contextShiftMetadata
                    },
                    evaluationPriority,
                    lastEvaluationContextWindow: {
                        history: newContextWindowChatHistory,
                        minimumOverlapPercentageToPreventContextShift: 0.01
                    }
                });
                this._ensureNotDisposed();
                lastEvaluation = currentLastEvaluation;
                newChatHistory = lastEvaluation.cleanHistory;
                if (functionCall != null) {
                    const functionDefinition = functions?.[functionCall.functionName];
                    if (functionDefinition == null)
                        throw new Error(`The model tried to call function "${functionCall.functionName}" which is not defined`);
                    const functionCallResult = await functionDefinition.handler(functionCall.params);
                    this._ensureNotDisposed();
                    newChatHistory = addFunctionCallToChatHistory({
                        chatHistory: newChatHistory,
                        functionName: functionCall.functionName,
                        functionDescription: functionDefinition.description,
                        callParams: functionCall.params,
                        callResult: functionCallResult,
                        raw: functionCall.raw + this._chat.chatWrapper.generateFunctionCallResult(functionCall.functionName, functionCall.params, functionCallResult)
                    });
                    newContextWindowChatHistory = addFunctionCallToChatHistory({
                        chatHistory: lastEvaluation.contextWindow,
                        functionName: functionCall.functionName,
                        functionDescription: functionDefinition.description,
                        callParams: functionCall.params,
                        callResult: functionCallResult,
                        raw: functionCall.raw + this._chat.chatWrapper.generateFunctionCallResult(functionCall.functionName, functionCall.params, functionCallResult)
                    });
                    lastEvaluation.cleanHistory = newChatHistory;
                    lastEvaluation.contextWindow = newContextWindowChatHistory;
                    continue;
                }
                this._lastEvaluation = lastEvaluation;
                this._chatHistory = newChatHistory;
                const lastModelResponseItem = getLastModelResponseItem(newChatHistory);
                return {
                    response: lastModelResponseItem.response,
                    responseText: lastModelResponseItem.response
                        .filter((item) => typeof item === "string")
                        .join(""),
                    stopReason: metadata.stopReason,
                    remainingGenerationAfterStop: metadata.remainingGenerationAfterStop
                };
            }
        });
    }
    getChatHistory() {
        return structuredClone(this._chatHistory);
    }
    getLastEvaluationContextWindow() {
        if (this._lastEvaluation == null)
            return null;
        return structuredClone(this._lastEvaluation?.contextWindow);
    }
    setChatHistory(chatHistory) {
        this._chatHistory = structuredClone(chatHistory);
        this._lastEvaluation = undefined;
    }
    /** @internal */
    _ensureNotDisposed() {
        if (this.disposed)
            throw new DisposedError();
    }
}
function addFunctionCallToChatHistory({ chatHistory, functionName, functionDescription, callParams, callResult, raw }) {
    const newChatHistory = chatHistory.slice();
    if (newChatHistory.length === 0 || newChatHistory[newChatHistory.length - 1].type !== "model")
        newChatHistory.push({
            type: "model",
            response: []
        });
    const lastModelResponseItem = newChatHistory[newChatHistory.length - 1];
    const newLastModelResponseItem = { ...lastModelResponseItem };
    newChatHistory[newChatHistory.length - 1] = newLastModelResponseItem;
    const modelResponse = newLastModelResponseItem.response.slice();
    newLastModelResponseItem.response = modelResponse;
    modelResponse.push({
        type: "functionCall",
        name: functionName,
        description: functionDescription,
        params: callParams,
        result: callResult,
        raw
    });
    return newChatHistory;
}
function getLastModelResponseItem(chatHistory) {
    if (chatHistory.length === 0 || chatHistory[chatHistory.length - 1].type !== "model")
        throw new Error("Expected chat history to end with a model response");
    return chatHistory[chatHistory.length - 1];
}
//# sourceMappingURL=LlamaChatSession.js.map