/*
 * Copyright 2025 the original author or authors.
 * <p>
 * Licensed under the Moderne Source Available License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * <p>
 * https://docs.moderne.io/licensing/moderne-source-available-license
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import * as rpc from "vscode-jsonrpc/node";
import {Recipe} from "../../recipe";
import {Cursor, Tree} from "../../tree";
import {TreeVisitor} from "../../visitor";
import {MarkersKind, SearchResult} from "../../markers";
import {Visit} from "./visit";
import {ExecutionContext} from "../../execution";
import {DATA_TABLE_STORE, DataTableStore} from "../../data-table";
import {withMetrics} from "./metrics";

export interface BatchVisitItem {
    visitor: string
    visitorOptions?: Map<string, any>
}

export interface BatchVisitResult {
    modified: boolean
    deleted: boolean
    hasNewMessages: boolean
    searchResultIds: string[]
}

export interface BatchVisitResponse {
    results: BatchVisitResult[]
}

export interface BatchVisitRequest {
    sourceFileType: string
    treeId: string
    p: string
    cursor?: string[]
    visitors: BatchVisitItem[]
}

async function collectSearchResultIds(tree: Tree | null | undefined): Promise<Set<string>> {
    const ids = new Set<string>();
    if (!tree) return ids;

    await new class extends TreeVisitor<any, Set<string>> {
        protected async visitMarker(marker: any, ctx: Set<string>): Promise<any> {
            if (marker && marker.kind === MarkersKind.SearchResult) {
                ctx.add((marker as SearchResult).id.toString());
            }
            return super.visitMarker(marker, ctx);
        }
    }().visit(tree, ids);
    return ids;
}

export class BatchVisit {
    static handle(connection: rpc.MessageConnection,
                  localObjects: Map<string, any>,
                  preparedRecipes: Map<String, Recipe>,
                  recipeCursors: WeakMap<Recipe, Cursor>,
                  getObject: (id: string, sourceFileType?: string) => any,
                  captureRefCheckpoint: (treeId: string) => void,
                  getCursor: (cursorIds: string[] | undefined, sourceFileType?: string) => Promise<Cursor>,
                  dataTableStore: () => DataTableStore | undefined,
                  metricsCsv?: string): void {
        connection.onRequest(
            new rpc.RequestType<BatchVisitRequest, BatchVisitResponse, Error>("BatchVisit"),
            withMetrics<BatchVisitRequest, BatchVisitResponse>(
                "BatchVisit",
                metricsCsv,
                (_context) => async (request) => {
                    const p = await getObject(request.p, undefined);
                    const store = dataTableStore();
                    if (store && p instanceof ExecutionContext) {
                        p.messages[DATA_TABLE_STORE] = store;
                    }
                    captureRefCheckpoint(request.treeId);
                    let tree: Tree = await getObject(request.treeId, request.sourceFileType);
                    const cursor = await getCursor(request.cursor, request.sourceFileType);

                    const results: BatchVisitResult[] = [];
                    const knownIds = await collectSearchResultIds(tree);

                    for (const item of request.visitors) {
                        // Instantiate and run visitor
                        const visitor = await Visit.instantiateVisitor(
                            {visitor: item.visitor, visitorOptions: item.visitorOptions},
                            preparedRecipes, recipeCursors, p);
                        // Snapshot ctx message keys so we can flag whether
                        // the visitor put anything new into the context.
                        const preKeys = new Set<string | symbol>(
                            Reflect.ownKeys(p.messages) as (string | symbol)[]
                        );
                        const after = await visitor.visit(tree, p, cursor);

                        const modified = after !== tree;
                        const deleted = after == null;

                        let hasNewMessages = false;
                        for (const k of Reflect.ownKeys(p.messages)) {
                            if (!preKeys.has(k)) {
                                hasNewMessages = true;
                                break;
                            }
                        }

                        // Diff SearchResult IDs against the running set
                        let searchResultIds: string[];
                        if (deleted) {
                            searchResultIds = [];
                        } else {
                            const afterIds = await collectSearchResultIds(after);
                            searchResultIds = [...afterIds].filter(id => !knownIds.has(id));
                            for (const id of searchResultIds) knownIds.add(id);
                        }

                        results.push({modified, deleted, hasNewMessages, searchResultIds});

                        if (deleted) {
                            localObjects.delete(request.treeId);
                            break;
                        }

                        if (modified) {
                            tree = after;
                        }
                    }

                    // Store final tree in localObjects
                    if (tree != null) {
                        localObjects.set(tree.id.toString(), tree);
                        if (tree.id.toString() !== request.treeId) {
                            localObjects.set(request.treeId, tree);
                        }
                    }

                    return {results};
                }
            )
        );
    }
}