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

import { CompletionDirection, SeparatorMode } from "@typeagent/agent-sdk";
import { mergeSeparatorMode } from "@typeagent/agent-sdk/helpers/command";
import {
    ExecutableAction,
    HistoryContext,
} from "../explanation/requestAction.js";
import {
    Construction,
    ConstructionMatchResult,
    ConstructionPart,
    WildcardMode,
} from "./constructions.js";
import { MatchPart, MatchSet, isMatchPart } from "./matchPart.js";
import { Transforms } from "./transforms.js";

import registerDebug from "debug";
import {
    MatchConfig,
    createMatchPartsCache,
    getMatchPartsCacheStats,
} from "./constructionMatch.js";
import {
    ConstructionCacheJSON,
    constructionCacheJSONVersion,
} from "./constructionJSONTypes.js";
import {
    getLanguageTools,
    needsSeparatorInAutoMode,
} from "../utils/language.js";
const debugConst = registerDebug("typeagent:const");
const debugConstMatchStat = registerDebug("typeagent:const:match:stat");
const debugCompletion = registerDebug("typeagent:const:completion");

// Agent Cache define the namespace policy.  At the cache, it just combine the keys into a string for lookup.
function getConstructionNamespace(namespaceKeys: string[]) {
    // Combine the namespace keys into a string using | as the separator.  Use to filter easily when
    // during match for schemas are disabled or not or hash mismatches.
    return namespaceKeys.join("|");
}

function getNamespaceKeys(constructionNamespace: string) {
    // Convert the namespace into an array of translator names for filtering.
    return constructionNamespace.split("|");
}

type AddConstructionResult =
    | {
          added: true;
          existing: Construction[];
          construction: Construction;
      }
    | {
          added: false;
          existing: Construction[];
      };

type Constructions = {
    constructions: Construction[];
    // For assigning runtime id
    maxId: number;
};
export type NamespaceKeyFilter = (namespaceKey: string) => boolean;
export type MatchOptions = {
    // namespace keys to filter.  If undefined, all constructions are used.
    namespaceKeys?: string[] | undefined;
    wildcard?: boolean; // default is true
    entityWildcard?: boolean; // default is true
    rejectReferences?: boolean; // default is true
    conflicts?: boolean; // default is false
    history?: HistoryContext | undefined;
};

export type CompletionProperty = {
    actions: ExecutableAction[];
    names: string[];
};

export type CompletionResult = {
    completions: string[];
    properties?: CompletionProperty[] | undefined;
    // Characters consumed by the grammar before the completion point.
    matchedPrefixLength?: number | undefined;
    // What kind of separator is required between the already-typed prefix
    // and the completion text.  See SeparatorMode in @typeagent/agent-sdk.
    separatorMode?: SeparatorMode | undefined;
    // True when the completions form a closed set — if the user types
    // something not in the list, no further completions can exist
    // beyond it.  False or undefined means the parser can continue
    // past unrecognized input and find more completions.
    closedSet?: boolean | undefined;
    // True when the result would differ if queried with the opposite
    // direction.  When false, the caller can skip re-fetching on
    // direction change.
    directionSensitive?: boolean | undefined;
    // True when the completions are offered at a position where a
    // wildcard was finalized at end-of-input.  The wildcard's extent
    // is ambiguous — the user may still be typing within it — so the
    // caller should allow the anchor to slide forward on further input
    // rather than re-fetching or giving up.
    openWildcard?: boolean | undefined;
};

// Architecture: docs/architecture/completion.md — §2 Cache Layer
export function mergeCompletionResults(
    first: CompletionResult | undefined,
    second: CompletionResult | undefined,
): CompletionResult | undefined {
    if (first === undefined) {
        return second;
    }
    if (second === undefined) {
        return first;
    }
    // Eagerly discard shorter-prefix completions — consistent with the
    // grammar matcher's approach.  Only the source(s) with the longest
    // matchedPrefixLength contribute completions.
    const firstLen = first.matchedPrefixLength ?? 0;
    const secondLen = second.matchedPrefixLength ?? 0;
    if (firstLen > secondLen) {
        return first;
    }
    if (secondLen > firstLen) {
        return second;
    }
    // Same prefix length — merge completions from both sources.
    const matchedPrefixLength =
        first.matchedPrefixLength !== undefined ||
        second.matchedPrefixLength !== undefined
            ? firstLen
            : undefined;
    return {
        completions: [...first.completions, ...second.completions],
        properties: first.properties
            ? second.properties
                ? [...first.properties, ...second.properties]
                : first.properties
            : second.properties,
        matchedPrefixLength,
        separatorMode: mergeSeparatorMode(
            first.separatorMode,
            second.separatorMode,
        ),
        // Closed set only when both sources are closed sets.
        closedSet:
            first.closedSet !== undefined || second.closedSet !== undefined
                ? (first.closedSet ?? false) && (second.closedSet ?? false)
                : undefined,
        // Direction-sensitive if either source is.
        directionSensitive:
            first.directionSensitive !== undefined ||
            second.directionSensitive !== undefined
                ? (first.directionSensitive ?? false) ||
                  (second.directionSensitive ?? false)
                : undefined,
        // Open wildcard if either source has one.
        openWildcard:
            first.openWildcard !== undefined ||
            second.openWildcard !== undefined
                ? (first.openWildcard ?? false) ||
                  (second.openWildcard ?? false)
                : undefined,
    };
}

export class ConstructionCache {
    private readonly matchSetsByUid = new Map<string, MatchSet>();

    // Construction and transforms use different namespaces.
    private readonly constructionNamespaces = new Map<string, Constructions>();
    private readonly transformNamespaces = new Map<string, Transforms>();
    public constructor(public readonly explainerName: string) {}

    public get count() {
        let count = 0;
        for (const constructionNamespace of this.constructionNamespaces.values()) {
            count += constructionNamespace.constructions.length;
        }
        return count;
    }

    public getFilteredCount(filter: NamespaceKeyFilter) {
        let count = 0;
        for (const [
            namespace,
            constructionNamespace,
        ] of this.constructionNamespaces.entries()) {
            const keys = getNamespaceKeys(namespace);
            if (keys.every((key) => filter(key))) {
                count += constructionNamespace.constructions.length;
            }
        }
        return count;
    }

    private addMatchSet(matchSet: MatchSet, mergeMatchSet: boolean) {
        const merge = mergeMatchSet && matchSet.canBeMerged;
        const uid = merge ? matchSet.mergedUid : matchSet.unmergedUid;
        let newMatchSet: MatchSet | undefined = this.matchSetsByUid.get(uid);
        if (newMatchSet !== undefined) {
            // If merge, then add to the existing match set.
            // If non-merge, then the uid will have determine the equivalent match set to reuse
            if (newMatchSet !== matchSet && merge) {
                for (const match of matchSet.matches) {
                    // Merge matches
                    newMatchSet.matches.add(match);
                }
                // match set is modified, clear the regexp
                newMatchSet.clearRegexp();
            }
        } else {
            newMatchSet = matchSet.clone(merge, this.matchSetsByUid.size);
            this.matchSetsByUid.set(uid, newMatchSet);
        }
        return newMatchSet;
    }

    private ensureConstructionNamespace(namespace: string) {
        const constructionNamespace =
            this.constructionNamespaces.get(namespace);
        if (constructionNamespace !== undefined) {
            return constructionNamespace;
        }
        const newCacheNamespace = {
            constructions: [],
            transforms: new Transforms(),
            maxId: 0,
        };
        this.constructionNamespaces.set(namespace, newCacheNamespace);
        return newCacheNamespace;
    }

    private mergeTransformNamespaces(
        transformNamespaces: Map<string, Transforms>,
        cacheConflicts?: boolean,
    ) {
        for (const [namespace, transforms] of transformNamespaces) {
            const transformNamespace = this.transformNamespaces.get(namespace);
            if (transformNamespace === undefined) {
                this.transformNamespaces.set(namespace, transforms);
            } else {
                transformNamespace.merge(transforms, cacheConflicts);
            }
        }
    }

    public addConstruction(
        namespaceKeys: string[],
        construction: Construction,
        mergeMatchSets: boolean,
        cacheConflicts?: boolean,
    ): AddConstructionResult {
        const mergedParts = construction.parts.map((p) =>
            isMatchPart(p)
                ? new MatchPart(
                      p.matchSet
                          ? this.addMatchSet(p.matchSet, mergeMatchSets)
                          : undefined,
                      p.optional,
                      p.wildcardMode,
                      p.transformInfos,
                  )
                : p,
        );

        const namespace = getConstructionNamespace(namespaceKeys);
        const constructionNamespace =
            this.ensureConstructionNamespace(namespace);
        this.mergeTransformNamespaces(
            construction.transformNamespaces,
            cacheConflicts,
        );

        // Detect if there are existing rules
        const existingRules = constructionNamespace.constructions.filter((c) =>
            c.isSupersetOf(mergedParts, construction.implicitParameters),
        );
        if (existingRules.length) {
            return { added: false, existing: existingRules };
        }

        // Create a new rule and remove all the existing rule that the new rule is a superset of
        // REVIEW: do we want to share transforms globally?
        const newConstruction = new Construction(
            mergedParts,
            this.transformNamespaces,
            construction.emptyArrayParameters,
            construction.implicitParameters,
            construction.implicitActionName,
            constructionNamespace.maxId++,
        );

        const removedRules: Construction[] = [];
        constructionNamespace.constructions =
            constructionNamespace.constructions.filter((c) => {
                const isSupersetOf = newConstruction.isSupersetOf(
                    c.parts,
                    c.implicitParameters,
                );
                if (isSupersetOf) {
                    removedRules.push(c);
                    return false;
                }
                return true;
            });

        constructionNamespace.constructions.push(newConstruction);
        return {
            added: true,
            existing: removedRules,
            construction: newConstruction,
        };
    }

    public forceRegexp() {
        this.matchSetsByUid.forEach((matchSet) => matchSet.forceRegexp());
    }

    public delete(namespace: string, id: number) {
        const constructionNamespace =
            this.constructionNamespaces.get(namespace);
        if (constructionNamespace === undefined) {
            return -1;
        }

        const count = constructionNamespace.constructions.length;
        constructionNamespace.constructions =
            constructionNamespace.constructions.filter((c) => c.id !== id);

        // TODO: GC match sets
        return count - constructionNamespace.constructions.length;
    }

    private getMatches(
        request: string,
        matchConfig: MatchConfig,
        constructionNamespace: Constructions,
    ): ConstructionMatchResult[] {
        return constructionNamespace.constructions.flatMap((construction) => {
            return construction.match(request, matchConfig);
        });
    }

    public prune(filter: NamespaceKeyFilter) {
        let count = 0;
        for (const namespace of this.constructionNamespaces.keys()) {
            const keys = getNamespaceKeys(namespace);
            if (!keys.every((key) => filter(key))) {
                this.constructionNamespaces.delete(namespace);
                debugConst(`Prune: ${namespace} deleted`);
                count++;
            }
        }
        return count;
    }

    // For matching
    public match(
        request: string,
        options?: MatchOptions,
        partial?: boolean,
        needMatchedStarts?: boolean,
    ): ConstructionMatchResult[] {
        const namespaceKeys = options?.namespaceKeys;
        if (namespaceKeys?.length === 0) {
            return [];
        }
        const config = {
            enableWildcard: options?.wildcard ?? true, // default to true.
            enableEntityWildcard: options?.entityWildcard ?? true, // default to true.
            rejectReferences: options?.rejectReferences ?? true, // default to true.
            history: options?.history,
            conflicts: options?.conflicts,
            matchPartsCache: createMatchPartsCache(request),
            partial: partial ?? false, // default to false.
            needMatchedStarts: needMatchedStarts ?? false,
        };

        // If the useTranslators is undefined use all the translators
        // otherwise filter the translators based on the useTranslators
        const matches: ConstructionMatchResult[] = [];
        const filter = namespaceKeys ? new Set(namespaceKeys) : undefined;
        for (const [
            name,
            constructionNamespace,
        ] of this.constructionNamespaces.entries()) {
            const keys = getNamespaceKeys(name);
            if (filter && keys.some((key) => !filter.has(key))) {
                continue;
            }

            matches.push(
                ...this.getMatches(request, config, constructionNamespace),
            );
        }
        debugConstMatchStat(getMatchPartsCacheStats(config.matchPartsCache));
        return matches;
    }

    public completion(
        requestPrefix: string,
        options?: MatchOptions,
        direction?: CompletionDirection, // defaults to forward-like behavior when omitted
    ): CompletionResult | undefined {
        debugCompletion(`Request completion for prefix: '${requestPrefix}'`);
        const namespaceKeys = options?.namespaceKeys;
        debugCompletion(`Request completion namespace keys`, namespaceKeys);

        // Resolve direction to a boolean: true when the user is actively
        // backing up and no trailing separator has committed the last token.
        const backward =
            direction === "backward" && !/[\s\p{P}]$/u.test(requestPrefix);
        const results = this.match(requestPrefix, options, true, backward);

        debugCompletion(
            `Request completion construction match: ${results.length}`,
        );

        if (results.length === 0) {
            return undefined;
        }

        // Track the furthest character position consumed across all
        // matching constructions.  When a longer match is found, all
        // previously accumulated completions from shorter matches are
        // discarded — mirroring the grammar matcher's approach.
        let maxPrefixLength = 0;
        const completionProperty: CompletionProperty[] = [];
        const requestText: string[] = [];
        let separatorMode: SeparatorMode | undefined;
        // Whether the accumulated completions form a closed set.
        // Starts true; set to false when property/wildcard completions
        // are added (entity values are external).  Reset to true when
        // maxPrefixLength advances (old candidates discarded).
        let closedSet: boolean = true;
        // Direction-sensitive when the opposite direction would produce
        // different completions.  True when at least one construction
        // at maxPrefixLength has matched parts to back up to and the
        // prefix doesn't end with a commit signal (separator).
        const noTrailingSeparator = !/[\s\p{P}]$/u.test(requestPrefix);
        let directionSensitive = false;
        const rejectReferences = options?.rejectReferences ?? true;
        const langTools = getLanguageTools("en");

        function updateMaxPrefixLength(prefixLength: number): void {
            if (prefixLength > maxPrefixLength) {
                maxPrefixLength = prefixLength;
                requestText.length = 0;
                completionProperty.length = 0;
                separatorMode = undefined;
                closedSet = true;
                directionSensitive = false;
            }
        }

        for (const result of results) {
            const { construction, partialPartCount, partialMatchedCurrent } =
                result;
            if (partialPartCount === undefined) {
                throw new Error(
                    "Internal Error: Partial part count is undefined",
                );
            }

            // --- Step 1: Determine which part to complete and the
            //     prefix length up to that point. ---
            let completionPart: ConstructionPart | undefined;
            let candidatePrefixLength: number;

            if (backward) {
                // Walk matchedStarts backwards to find the last part
                // that actually matched (skip optional parts = -1).
                const matchedStarts = result.matchedStarts;
                candidatePrefixLength = -1;
                if (matchedStarts !== undefined && partialPartCount > 0) {
                    for (let i = partialPartCount - 1; i >= 0; i--) {
                        if (matchedStarts[i] >= 0) {
                            completionPart = construction.parts[i];
                            candidatePrefixLength = matchedStarts[i];
                            break;
                        }
                    }
                }
                if (candidatePrefixLength < 0) {
                    continue; // Nothing matched to back up to
                }
            } else {
                // Forward: exact match means nothing to complete.
                if (partialPartCount === construction.parts.length) {
                    updateMaxPrefixLength(requestPrefix.length);
                    if (noTrailingSeparator && partialPartCount >= 1) {
                        directionSensitive = true;
                    }
                    continue;
                }
                completionPart = construction.parts[partialPartCount];
                candidatePrefixLength = partialMatchedCurrent ?? 0;
            }

            // --- Step 2: Check against maxPrefixLength ---
            updateMaxPrefixLength(candidatePrefixLength);
            if (candidatePrefixLength !== maxPrefixLength) {
                continue; // Shorter than the best match — skip
            }

            // Direction-sensitive when a matched part exists to back
            // up to and the prefix has no trailing commit signal.
            if (noTrailingSeparator && partialPartCount >= 1) {
                directionSensitive = true;
            }

            // --- Step 3: Offer literal completions from the part ---
            if (
                completionPart !== undefined &&
                completionPart.wildcardMode <= WildcardMode.Enabled
            ) {
                const partCompletions = completionPart.getCompletion();
                if (partCompletions) {
                    for (const completionText of partCompletions) {
                        if (
                            completionPart.capture &&
                            rejectReferences &&
                            langTools?.possibleReferentialPhrase(completionText)
                        ) {
                            continue;
                        }
                        requestText.push(completionText);
                        if (
                            candidatePrefixLength > 0 &&
                            completionText.length > 0
                        ) {
                            const needsSep = needsSeparatorInAutoMode(
                                requestPrefix[candidatePrefixLength - 1],
                                completionText[0],
                            );
                            separatorMode = mergeSeparatorMode(
                                separatorMode,
                                needsSep ? "spacePunctuation" : "optional",
                            );
                        }
                    }
                }
            }

            // --- Step 4: Offer property completions for entity parts ---
            if (completionPart !== undefined) {
                const partPropertyNames = completionPart.getPropertyNames();
                if (
                    partPropertyNames !== undefined &&
                    partPropertyNames.length > 0
                ) {
                    // Filter out properties that appear in multiple parts
                    // so we only offer single-part properties.
                    const allPropertyNames = new Map<string, number>();
                    for (const part of construction.parts) {
                        const names = part.getPropertyNames();
                        if (names === undefined) {
                            continue;
                        }
                        for (const name of names) {
                            const count = allPropertyNames.get(name) ?? 0;
                            allPropertyNames.set(name, count + 1);
                        }
                    }

                    const queryPropertyNames = partPropertyNames.filter(
                        (name: string) => allPropertyNames.get(name) === 1,
                    );
                    if (queryPropertyNames.length > 0) {
                        completionProperty.push({
                            actions: result.match.actions,
                            names: queryPropertyNames,
                        });
                        if (candidatePrefixLength > 0) {
                            const needsSep = needsSeparatorInAutoMode(
                                requestPrefix[candidatePrefixLength - 1],
                                "a",
                            );
                            separatorMode = mergeSeparatorMode(
                                separatorMode,
                                needsSep ? "spacePunctuation" : "optional",
                            );
                        }
                        closedSet = false;
                    }
                }
            }
        }

        // Advance past trailing separators so that the reported prefix
        // length includes any trailing whitespace the user typed.
        // When advancing, demote separatorMode to "optional" — the
        // trailing space is already consumed.
        if (maxPrefixLength < requestPrefix.length) {
            const trailing = requestPrefix.substring(maxPrefixLength);
            if (/^[\s\p{P}]+$/u.test(trailing)) {
                maxPrefixLength = requestPrefix.length;
                separatorMode = "optional";
            }
        }

        return {
            completions: requestText,
            properties: completionProperty,
            matchedPrefixLength: maxPrefixLength,
            separatorMode,
            closedSet,
            directionSensitive,
        };
    }

    public get matchSets(): IterableIterator<MatchSet> {
        return this.matchSetsByUid.values();
    }

    public toJSON() {
        return {
            version: constructionCacheJSONVersion,
            explainerName: this.explainerName,
            matchSets: Array.from(this.matchSets),
            constructionNamespaces: Array.from(
                this.constructionNamespaces.entries(),
            ).map(([name, constructionNamespace]) => ({
                name,
                constructions: constructionNamespace.constructions,
            })),
            transformNamespaces: Array.from(
                this.transformNamespaces.entries(),
            ).map(([name, transforms]) => ({
                name,
                transforms,
            })),
        };
    }

    public static fromJSON(originalJSON: ConstructionCacheJSON) {
        const json = ensureVersion(originalJSON);
        const store = new ConstructionCache(json.explainerName);

        // Load the match sets
        const allMatchSets = new Map<string, MatchSet>();
        for (const matchSet of json.matchSets) {
            const newMatchSet = new MatchSet(
                matchSet.matches,
                matchSet.basename,
                matchSet.canBeMerged,
                matchSet.namespace,
                matchSet.index,
            );
            const uid = matchSet.canBeMerged
                ? newMatchSet.mergedUid
                : newMatchSet.unmergedUid;
            store.matchSetsByUid.set(uid, newMatchSet);
            allMatchSets.set(newMatchSet.fullName, newMatchSet);
        }

        // load the constructions and transforms for each translator
        json.constructionNamespaces.forEach(({ name, constructions }) => {
            const newConstructions = constructions.map((construction, index) =>
                Construction.fromJSON(
                    construction,
                    allMatchSets,
                    store.transformNamespaces,
                    index,
                ),
            );
            store.constructionNamespaces.set(name, {
                constructions: newConstructions,
                maxId: newConstructions.length,
            });
            debugConst(newConstructions.join("\n  "));
        });

        json.transformNamespaces.forEach(({ name, transforms }) => {
            store.transformNamespaces.set(
                name,
                Transforms.fromJSON(transforms),
            );
        });
        return store;
    }

    // for viewers
    public getConstructionNamespace(namespace: string) {
        return this.constructionNamespaces.get(namespace);
    }

    public getConstructionNamespaces() {
        return Array.from(this.constructionNamespaces.keys());
    }

    public getTransformNamespaces() {
        return this.transformNamespaces;
    }
}

function ensureVersion(json: any): ConstructionCacheJSON {
    if (json.version === constructionCacheJSONVersion) {
        return json as ConstructionCacheJSON;
    }

    throw new Error(
        `Unsupported version of ConstructionCache: ${json.version}`,
    );
}
