// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

import { ChatModel, openai } from "aiclient";
import * as kp from "knowpro";
import {
    conversation as kpLib,
    TextEmbeddingModelWithCache,
    TextEmbeddingCache,
} from "knowledge-processor";
import * as ms from "memory-storage";
import { error, PromptSection, Result, success } from "typechat";
import { createEmbeddingModelWithCache } from "./common.js";

export interface MemorySettings {
    languageModel: ChatModel;
    embeddingModel: TextEmbeddingModelWithCache;
    embeddingSize: number;
    conversationSettings: kp.ConversationSettings;
    queryTranslator?: kp.SearchQueryTranslator | undefined;
    answerGenerator?: kp.IAnswerGenerator | undefined;
    fileSaveSettings?: IndexFileSettings | undefined;
    /**
     * Experimental: use improved usage of 'scopes' in queries
     * Improved scoping leverages Structured Tags
     * @see {kp.searchConversationWithLanguage2}
     */
    useScopedSearch?: boolean | undefined;
}

export function createMemorySettings(
    embeddingCacheSize = 64,
    getPersistentCache?: () => TextEmbeddingCache | undefined,
    languageModel?: ChatModel,
): MemorySettings {
    languageModel ??= openai.createChatModelDefault("conversation-memory");
    /**
     * Our index already has embeddings for every term in the podcast
     * Create a caching embedding model that can just leverage those embeddings
     * @returns embedding model, size of embedding
     */
    const [embeddingModel, embeddingSize] = createEmbeddingModelWithCache(
        embeddingCacheSize,
        getPersistentCache,
    );

    const conversationSettings = kp.createConversationSettings(
        embeddingModel,
        embeddingSize,
    );
    conversationSettings.semanticRefIndexSettings.knowledgeExtractor =
        kp.createKnowledgeExtractor(languageModel);

    const queryTranslator = kp.createSearchQueryTranslator(languageModel);
    const memorySettings: MemorySettings = {
        languageModel,
        embeddingModel,
        embeddingSize,
        conversationSettings,
        queryTranslator,
    };
    return memorySettings;
}

/**
 * Settings for how memory index file are stored
 */
export type IndexFileSettings = {
    /**
     * Directory that contains memory files
     */
    dirPath: string;
    /**
     * Base filename for memory files
     * Specific filenames used to persist memory use this as a prefix
     */
    baseFileName: string;
};

export type IndexingState = {
    lastMessageOrdinal: kp.MessageOrdinal;
    lastSemanticRefOrdinal: kp.SemanticRefOrdinal;
};

export type TermSynonyms = {
    term: string;
    relatedTerms: string[];
};

export function addSynonymsAsAliases(
    aliases: kp.TermToRelatedTermsMap,
    synonyms: TermSynonyms[],
): void {
    for (const ts of synonyms) {
        let relatedTerm: kp.Term = { text: ts.term.toLowerCase() };
        for (const synonym of ts.relatedTerms) {
            aliases.addRelatedTerm(synonym.toLowerCase(), relatedTerm);
        }
    }
}

export function addSynonymsFileAsAliases(
    aliases: kp.TermToRelatedTermsMap,
    filePath: string,
) {
    const synonyms = ms.readJsonFile<TermSynonyms[]>(filePath);
    if (synonyms && synonyms.length > 0) {
        addSynonymsAsAliases(aliases, synonyms);
    }
}

export function addNoiseWordsFromFile(
    noise: Set<string>,
    filePath: string,
): void {
    const words = ms.readAllLines(filePath);
    if (words) {
        words.forEach((word) => noise.add(word));
    }
}

export class MessageMetadata
    implements kp.IMessageMetadata, kp.IKnowledgeSource
{
    public get source(): string | string[] | undefined {
        return undefined;
    }
    public get dest(): string | string[] | undefined {
        return undefined;
    }

    public getKnowledge(): kpLib.KnowledgeResponse | undefined {
        return undefined;
    }
}

/**
 * A Message in a Memory {@link Memory}
 * Message implements {@link kp.IMessage}
 */
export class Message<TMeta extends MessageMetadata = MessageMetadata>
    implements kp.IMessage
{
    /**
     * Text chunk associated with this messages
     */
    public textChunks: string[];

    constructor(
        /**
         * {@link MessageMetadata}
         */
        public metadata: TMeta,
        textChunks: string | string[],
        public tags: kp.MessageTag[] = [],
        public timestamp: string | undefined = undefined,
        public knowledge: kpLib.KnowledgeResponse | undefined = undefined,
        public deletionInfo: kp.DeletionInfo | undefined = undefined,
    ) {
        if (Array.isArray(textChunks)) {
            this.textChunks = textChunks;
        } else {
            this.textChunks = [textChunks];
        }
    }

    public addContent(content: string, chunkOrdinal = 0) {
        if (
            chunkOrdinal > this.textChunks.length - 1 ||
            this.textChunks.length == 0
        ) {
            this.textChunks.push(content);
        } else {
            this.textChunks[chunkOrdinal] += content;
        }
    }

    public addKnowledge(
        newKnowledge: kpLib.KnowledgeResponse,
    ): kpLib.KnowledgeResponse {
        if (this.knowledge !== undefined) {
            this.knowledge.entities = this.mergeEntities([
                ...this.knowledge.entities,
                ...newKnowledge.entities,
            ]);
            this.knowledge.topics = kp.mergeTopics([
                ...this.knowledge.topics,
                ...newKnowledge.topics,
            ]);
            this.knowledge.actions.push(...newKnowledge.actions);
            this.knowledge.inverseActions.push(...newKnowledge.inverseActions);
        } else {
            this.knowledge = newKnowledge;
        }
        return this.knowledge;
    }

    getKnowledge(): kpLib.KnowledgeResponse | undefined {
        let metaKnowledge = this.metadata.getKnowledge();
        if (!metaKnowledge) {
            return this.knowledge;
        }
        if (!this.knowledge) {
            return metaKnowledge;
        }
        const combinedKnowledge: kpLib.KnowledgeResponse = {
            ...this.knowledge,
        };
        combinedKnowledge.entities.push(...metaKnowledge.entities);
        combinedKnowledge.entities = this.mergeEntities(
            combinedKnowledge.entities,
        );
        combinedKnowledge.topics.push(...metaKnowledge.topics);
        combinedKnowledge.topics = kp.mergeTopics(combinedKnowledge.topics);

        combinedKnowledge.actions.push(...metaKnowledge.actions);
        combinedKnowledge.inverseActions.push(...metaKnowledge.inverseActions);
        return combinedKnowledge;
    }

    private mergeEntities(
        entities: kpLib.ConcreteEntity[],
    ): kpLib.ConcreteEntity[] {
        // TODO: using mergeConcreteEntitiesEx to avoid forcing the data to be lower case.
        // Replace with mergeConcreteEntities once it has switch over to the Ex version.
        return kp.mergeConcreteEntitiesEx(entities);
    }
}

/**
 * A memory containing a sequence of messages {@link Message}
 * Memory is modeled as a conversation {@link kp.IConversation}
 */
export abstract class Memory<
    TSettings extends MemorySettings = MemorySettings,
    TMessage extends Message = Message,
> {
    public nameTag: string;
    public tags: string[];
    public settings: TSettings;
    public noiseTerms: Set<string>;

    constructor(settings: TSettings, nameTag?: string, tags?: string[]) {
        this.settings = settings;
        this.nameTag = nameTag ?? "";
        this.tags = tags ?? [];
        this.noiseTerms = new Set<string>();
        this.settings.queryTranslator ??= kp.createSearchQueryTranslator(
            this.settings.languageModel,
        );
        this.settings.embeddingModel.getPersistentCache = () =>
            this.getPersistentEmbeddingCache();

        this.addStandardNoiseTerms();
    }

    /**
     * The conversation representing this memory
     */
    public abstract get conversation(): kp.IConversation<TMessage>;

    /**
     * Search memory using a select expression. Searches for and returns both conversation knowledge and messages
     * @param {kp.SearchSelectExpr} selectExpr - The select expression to use for searching.
     * @returns {Promise<kp.ConversationSearchResult | undefined>} - The search result or undefined if not found.
     */
    public async search(
        selectExpr: kp.SearchSelectExpr,
        options?: kp.SearchOptions,
    ): Promise<kp.ConversationSearchResult | undefined> {
        return kp.searchConversation(
            this.conversation,
            selectExpr.searchTermGroup,
            selectExpr.when,
            options,
        );
    }

    /**
     * Search knowledge extracted from this conversation
     * @param selectExpr
     * @param options
     * @returns {Promise<Map<kp.KnowledgeType, kp.SemanticRefSearchResult> | undefined>} - knowledge matches or undefined if none
     */
    public async searchKnowledge(
        selectExpr: kp.SearchSelectExpr,
        options?: kp.SearchOptions,
    ): Promise<Map<kp.KnowledgeType, kp.SemanticRefSearchResult> | undefined> {
        return kp.searchConversationKnowledge(
            this.conversation,
            selectExpr.searchTermGroup,
            selectExpr.when,
            options,
        );
    }

    /**
     * Find matching topics in this conversation
     * Wrapper around the more powerful searchKnowledge
     * @param searchTerms
     * @param topK
     * @param exactMatch
     * @returns
     */
    public async searchTopics(
        searchTerms: string | string[],
        topK?: number,
        exactMatch: boolean = false,
    ): Promise<kp.Topic[] | undefined> {
        const searchTermGroup = kp.createTopicSearchTermGroup(
            searchTerms,
            exactMatch,
        );
        const matches = await this.searchKnowledge({
            searchTermGroup,
            when: { knowledgeType: "topic" },
        });
        const topicMatches = matches?.get("topic");
        if (!topicMatches || topicMatches.semanticRefMatches.length === 0) {
            return undefined;
        }
        /**
         * Search will find every matching instance of the occurrence of a topic during the conversation
         * Since we are only interested in the topics themselves,return a distinct list
         */
        const distinct = kp.getDistinctTopicMatches(
            this.conversation.semanticRefs!,
            topicMatches.semanticRefMatches,
            topK,
        );
        return distinct.map((t) => t.knowledge as kp.Topic);
    }

    /**
     * Find matching entities in this conversation
     * Wrapper around the more powerful searchKnowledge
     * @param name (Optional) entity name to match
     * @param type (Optional) entity type to match
     * @param facetName (Optional) facet name to match
     * @param facetValue (Optional) facet value  to match
     * @param topK
     * @param exactMatch
     * @returns
     */
    public async searchEntities(
        name: string | undefined,
        type?: string | undefined,
        facetName?: string | undefined,
        facetValue?: string | undefined,
        topK?: number,
        exactMatch: boolean = false,
    ): Promise<kpLib.ConcreteEntity[] | undefined> {
        const searchTermGroup = kp.createEntitySearchTermGroup(
            name,
            type,
            facetName,
            facetValue,
            exactMatch,
        );
        const matches = await this.searchKnowledge({
            searchTermGroup,
            when: { knowledgeType: "entity" },
        });
        const entityMatches = matches?.get("entity");
        if (!entityMatches || entityMatches.semanticRefMatches.length === 0) {
            return undefined;
        }

        const distinct = kp.getDistinctEntityMatches(
            this.conversation.semanticRefs!,
            entityMatches.semanticRefMatches,
            topK,
        );
        return distinct.map((t) => t.knowledge as kpLib.ConcreteEntity);
    }

    /***
     * Run a natural language query against this memory.
     * Natural language is translated to knowpro {@link kp.SearchSelectExpr | search expressions}
     * These expressions are used to search memory, which returns knowledge and messages relevant to the query.
     * Knowledge and messages are scored for relevance.
     * @param {string} searchText - The natural language query text.
     * @param {kp.LanguageSearchOptions} [options] - Optional search options.
     * @param {kp.LanguageSearchDebugContext} [debugContext] - Optional debug context.
     * @returns {Promise<Result<kp.ConversationSearchResult[]>>} - The search results.
     */
    public async searchWithLanguage(
        searchText: string,
        options?: kp.LanguageSearchOptions,
        langSearchFilter?: kp.LanguageSearchFilter,
        debugContext?: kp.LanguageSearchDebugContext,
    ): Promise<Result<kp.ConversationSearchResult[]>> {
        options = this.adjustLanguageSearchOptions(options);
        if (!this.useScoped) {
            return kp.searchConversationWithLanguage(
                this.conversation,
                searchText,
                this.getQueryTranslator(),
                options,
                langSearchFilter,
                debugContext,
            );
        } else {
            return kp.searchConversationWithLanguage2(
                this.conversation,
                searchText,
                this.getQueryTranslator(),
                options,
                langSearchFilter,
                debugContext,
            );
        }
    }

    /**
     * Run a query on the memory using topics
     * Simple wrapper around {@link search}
     * @param topicTerms
     * @param {kp.WhenFilter} when
     * @param {kp.SearchOptions} options
     * @returns {kp.ConversationSearchResult}
     */
    public async searchWithTopics(
        topicTerms: string | string[],
        when?: kp.WhenFilter,
        options?: kp.SearchOptions,
    ): Promise<kp.ConversationSearchResult | undefined> {
        const selectExpr: kp.SearchSelectExpr = {
            searchTermGroup: kp.createTopicSearchTermGroup(topicTerms),
            when,
        };
        return this.search(selectExpr, options);
    }

    /**
     * Run a query on the memory using entity matching
     * Simple wrapper around {@link search}
     * @param name (Optional) entity name to match
     * @param type (Optional) entity type to match
     * @param facetName (Optional) facet name to match
     * @param facetValue (Optional) facet value  to match
     * @param {kp.WhenFilter} when (Optional) Scoping filter
     * @param {kp.SearchOptions} options
     * @returns
     */
    public async searchWithEntities(
        name: string | undefined,
        type?: string | undefined,
        facetName?: string | undefined,
        facetValue?: string | undefined,
        when?: kp.WhenFilter,
        options?: kp.SearchOptions,
    ): Promise<kp.ConversationSearchResult | undefined> {
        const selectExpr: kp.SearchSelectExpr = {
            searchTermGroup: kp.createEntitySearchTermGroup(
                name,
                type,
                facetName,
                facetValue,
            ),
            when,
        };
        return this.search(selectExpr, options);
    }

    /**
     * Translate a natural language query into a query expression.
     * @param {string} searchText - The natural language query text.
     * @param {kp.LanguageSearchOptions} [options] - Optional search options.
     * @returns {Promise<Result<kp.querySchema.SearchQuery>>} - The translated query expression.
     */
    public async searchQueryFromLanguage(
        searchText: string,
        options?: kp.LanguageSearchOptions,
    ): Promise<Result<kp.querySchema.SearchQuery>> {
        options = this.adjustLanguageSearchOptions(options);
        if (!this.useScoped) {
            return kp.searchQueryFromLanguage(
                this.conversation,
                this.getQueryTranslator(),
                searchText,
                this.getModelInstructions(),
            );
        } else {
            return kp.searchQueryFromLanguage2(
                this.conversation,
                this.getQueryTranslator(),
                searchText,
                this.getModelInstructions(),
            );
        }
    }

    /**
     * Get an answer from a natural language question. The returned result object will indicate success or failure.  
     * If the returned result is a success:
     * - A single question can be turned into multiple search queries, although most single phrase questions are a single query.
     * - The result is a tuple of (a) Raw search results (b) The natural language answer generated from the search results.
     * @see kp.AnswerResponse. 
     * - If the question was answered, the response {@link kp.AnswerType} will be "Answered". Else "NoAnswer" and a reason is provided in {@link kp.AnswerResponse.whyNoAnswer}
     
     * @param {string} question - The natural language question.
     * @param {kp.LanguageSearchOptions} [searchOptions] - Optional search options.
     * @param progress - Optional progress callback.
     * @returns {Promise<Result<[kp.ConversationSearchResult, kp.AnswerResponse][]>>} Tuple: [Raw Search Results, Answers generated using Search Results].
     */
    public async getAnswerFromLanguage(
        question: string,
        searchOptions?: kp.LanguageSearchOptions,
        langSearchFilter?: kp.LanguageSearchFilter,
        progress?: (
            searchResult: kp.ConversationSearchResult,
            chunk: kp.AnswerContext,
            index: number,
            result: Result<kp.AnswerResponse>,
        ) => void,
        answerContextOptions?: kp.AnswerContextOptions,
    ): Promise<Result<[kp.ConversationSearchResult, kp.AnswerResponse][]>> {
        const searchResults = await this.searchWithLanguage(
            question,
            searchOptions,
            langSearchFilter,
        );
        if (!searchResults.success) {
            return searchResults;
        }

        const answers: [kp.ConversationSearchResult, kp.AnswerResponse][] = [];
        for (let i = 0; i < searchResults.data.length; ++i) {
            const searchResult = searchResults.data[i];
            const answerResult = await this.getAnswerFromSearchResults(
                searchResult,
                searchResult.rawSearchQuery,
                progress !== undefined
                    ? (chunk, index, result) => {
                          progress(searchResult, chunk, index, result);
                      }
                    : undefined,
                answerContextOptions,
            );
            if (!answerResult.success) {
                return answerResult;
            }
            answers.push([searchResult, answerResult.data]);
        }

        return success(answers);
    }

    /**
     * Get an answer from search results.
     * @param {kp.ConversationSearchResult} searchResult - The search result.
     * @param {string} [question] - Optional question text.
     * @param progress - Optional progress callback.
     * @returns {Promise<Result<kp.AnswerResponse>>} - The answer response.
     */
    public async getAnswerFromSearchResults(
        searchResult: kp.ConversationSearchResult,
        question?: string,
        progress?: (
            chunk: kp.AnswerContext,
            index: number,
            result: Result<kp.AnswerResponse>,
        ) => void,
        answerContextOptions?: kp.AnswerContextOptions,
    ): Promise<Result<kp.AnswerResponse>> {
        question ??= searchResult.rawSearchQuery;
        if (!question) {
            return error("No searchResult.rawSearchQuery or question provided");
        }
        const answerGenerator = this.ensureAnswerGenerator();
        return kp.generateAnswer(
            this.conversation,
            answerGenerator,
            question,
            searchResult,
            progress,
            answerContextOptions,
        );
    }

    protected beginIndexing(): void {
        // Turn off caching during indexing because:
        // - LRU caches will not be useful
        // - Any persistent caches will be rebuilt anyway
        this.settings.embeddingModel.cacheEnabled = false;
    }

    protected endIndexing(): void {
        // See note in beginIndexing for why this was turned off during indexing
        this.settings.embeddingModel.cacheEnabled = true;
    }

    public getModelInstructions(): PromptSection[] | undefined {
        return undefined;
    }

    protected getPersistentEmbeddingCache(): TextEmbeddingCache | undefined {
        return undefined;
    }

    protected addStandardNoiseTerms() {
        addNoiseWordsFromFile(
            this.noiseTerms,
            ms.getAbsolutePathFromUrl(import.meta.url, "noiseTerms.txt"),
        );
    }

    protected ensureAnswerGenerator(): kp.IAnswerGenerator {
        if (this.settings.answerGenerator === undefined) {
            this.settings.answerGenerator = new kp.AnswerGenerator(
                kp.createAnswerGeneratorSettings(this.settings.languageModel),
            );
        }
        return this.settings.answerGenerator;
    }

    private getQueryTranslator(): kp.SearchQueryTranslator {
        const queryTranslator = this.settings.queryTranslator;
        if (!queryTranslator) {
            throw new Error("No query translator provided");
        }
        return queryTranslator;
    }

    private adjustLanguageSearchOptions(options?: kp.LanguageSearchOptions) {
        options ??= kp.createLanguageSearchOptionsTypical();
        const instructions = this.getModelInstructions();
        if (instructions) {
            if (options.modelInstructions) {
                options.modelInstructions.push(...instructions);
            } else {
                options.modelInstructions = instructions;
            }
        }
        options.compileOptions.termFilter = (t) =>
            !this.noiseTerms.has(t.toLowerCase());
        return options;
    }

    private get useScoped(): boolean {
        return this.settings.useScopedSearch ?? false;
    }
}

/**
 * Create settings for text memory.
 * @param embeddingCacheSize Default size of the embedding cache.
 * @param getPersistentCache Function to retrieve the persistent cache.
 * @returns Memory settings object.
 */

export function createTextMemorySettings(
    embeddingCacheSize = 64,
    getPersistentCache?: () => TextEmbeddingCache | undefined,
) {
    return {
        ...createMemorySettings(embeddingCacheSize, getPersistentCache),
    };
}
