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

import {
    InlineTranslatorSchemaDef,
    createJsonTranslatorFromSchemaDef,
} from "typechat-utils";
import { Result, success } from "typechat";
import registerDebug from "debug";
import { ActionConfigProvider } from "./actionConfigProvider.js";
import {
    generateSchemaTypeDefinition,
    SchemaCreator as sc,
} from "@typeagent/action-schema";
import { getCombinedActionSchemaTypeName } from "./agentTranslators.js";
import { PromptLogger } from "telemetry";
const debugSwitchSearch = registerDebug("typeagent:switch:search");

function createSelectionActionTypeDefinition(
    schemaName: string,
    provider: ActionConfigProvider,
) {
    const actionConfig = provider.getActionConfig(schemaName);
    // Skip injected schemas except for chat; investigate whether we can get chat always on first pass
    if (actionConfig.injected && schemaName !== "chat") {
        // No need to select for injected schemas
        return undefined;
    }
    const actionSchemaFile =
        provider.getActionSchemaFileForConfig(actionConfig);

    const actionNames: string[] = [];
    const actionComments: string[] = [];
    for (const [
        name,
        info,
    ] of actionSchemaFile.parsedActionSchema.actionSchemas.entries()) {
        actionNames.push(name);
        actionComments.push(
            ` "${name}"${info.comments ? ` - ${info.comments[0].trim()}` : ""}`,
        );
    }

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

    const typeName = `${getCombinedActionSchemaTypeName(actionConfig)}Assistant`;

    const schema = sc.type(
        typeName,
        sc.obj({
            assistant: sc.field(
                sc.string(schemaName),
                ` ${actionConfig.description}`,
            ),
            action: sc.field(sc.string(actionNames), actionComments),
        }),
    );
    return schema;
}

function createSelectionSchema(
    schemaName: string,
    provider: ActionConfigProvider,
): InlineTranslatorSchemaDef | undefined {
    const schema = createSelectionActionTypeDefinition(schemaName, provider);
    if (schema === undefined) {
        return undefined;
    }
    const typeName = schema.name;
    return {
        kind: "inline",
        typeName,
        schema: generateSchemaTypeDefinition(schema),
    };
}

const selectSchemaCache = new Map<
    string,
    InlineTranslatorSchemaDef | undefined
>();
function getSelectionSchema(
    schemaName: string,
    provider: ActionConfigProvider,
): InlineTranslatorSchemaDef | undefined {
    if (selectSchemaCache.has(schemaName)) {
        return selectSchemaCache.get(schemaName);
    }

    const result = createSelectionSchema(schemaName, provider);
    selectSchemaCache.set(schemaName, result);
    return result;
}

const unknownAssistantSelectionSchemaDef: InlineTranslatorSchemaDef = {
    kind: "inline",
    typeName: "UnknownAssistantSelection",
    schema: `
export type UnknownAssistantSelection = {
    assistant: "unknown";
    action: "unknown";
};`,
};

type AssistantSelectionSchemaEntry = {
    name: string;
    schema: InlineTranslatorSchemaDef;
};
export function getAssistantSelectionSchemas(
    schemaNames: Iterable<string>,
    provider: ActionConfigProvider,
) {
    const schemaEntries: AssistantSelectionSchemaEntry[] = [];
    for (const schemaName of schemaNames) {
        const schema = getSelectionSchema(schemaName, provider);
        if (schema !== undefined) {
            schemaEntries.push({ name: schemaName, schema });
        }
    }

    return schemaEntries;
}

export type AssistantSelection = {
    assistant: string;
    action: string;
};

type TranslatorPartition = {
    names: string[];
    translator: {
        translate: (
            request: string,
            signal?: AbortSignal,
        ) => Promise<Result<AssistantSelection>>;
    };
};

/**
 * Run all translator partitions in parallel and return the first non-"unknown"
 * result in partition order, or "unknown" if all partitions return "unknown".
 *
 * The provided signal is forwarded to each partition's translator so that an
 * abort tears down the in-flight LLM HTTP requests, rather than just stopping
 * the loop from awaiting them.
 */
export async function selectFromPartitions(
    partitions: TranslatorPartition[],
    request: string,
    signal?: AbortSignal,
): Promise<Result<AssistantSelection>> {
    partitions.forEach(({ names }) =>
        debugSwitchSearch(`Switch: searching ${names.join(", ")}`),
    );
    // Start all translations in parallel.
    const promises = partitions.map(({ translator }) =>
        translator
            .translate(request, signal)
            .catch((err): Result<AssistantSelection> => {
                // Let AbortError propagate so callers see cancellation rather
                // than a generic failure result. Non-abort errors are converted
                // to a failure result so a single bad partition doesn't block
                // the others from being checked.
                if (
                    err instanceof Error &&
                    (err.name === "AbortError" ||
                        (err as DOMException).name === "AbortError")
                ) {
                    throw err;
                }
                return {
                    success: false,
                    message: err instanceof Error ? err.message : String(err),
                };
            }),
    );
    // Await results in partition order; return as soon as outcome is decided.
    for (const promise of promises) {
        // Check for cancellation before and after awaiting each partition result
        // so an abort that fires mid-await is caught promptly.
        signal?.throwIfAborted();
        const result = await promise;
        signal?.throwIfAborted();
        if (!result.success) {
            return result;
        }
        if (result.data.assistant !== "unknown") {
            return result;
        }
    }
    return success({
        assistant: "unknown",
        action: "unknown",
    });
}

// GPT-4 has 8192 token window, with an estimated 4 chars per token, so use only 3 times to leave room for output.
const assistantSelectionLimit = 8192 * 3;

export function loadAssistantSelectionJsonTranslator(
    schemaNames: Iterable<string>,
    provider: ActionConfigProvider,
    promptLogger?: PromptLogger,
) {
    const schemas = getAssistantSelectionSchemas(schemaNames, provider);

    let currentLength = 0;
    let current: AssistantSelectionSchemaEntry[] = [];
    const limit = assistantSelectionLimit; // TODO: this should be adjusted based on model used.
    const partitions = [current];
    for (const entry of schemas) {
        const schema = entry.schema.schema;
        if (currentLength + schema.length > limit) {
            if (current.length === 0) {
                throw new Error(
                    `The assistant section schema for '${entry.name}' is too large to fit in the limit ${limit}`,
                );
            }
            current = [];
            currentLength = 0;
            partitions.push(current);
        }
        current.push(entry);
        currentLength += schema.length;
    }
    const translators = partitions.map((entries): TranslatorPartition => {
        const innerTranslator =
            createJsonTranslatorFromSchemaDef<AssistantSelection>(
                "AllAssistantSelection",
                entries
                    .map((entry) => entry.schema)
                    .concat(unknownAssistantSelectionSchemaDef),
                {
                    instructions: [
                        {
                            role: "system",
                            content:
                                "Select the assistant to handle the request",
                        },
                    ],
                    promptLogger,
                },
            );
        return {
            names: entries.map((entry) => entry.name),
            // Adapt typechat's translate(request, promptPreamble?, signal?) to
            // the (request, signal?) shape selectFromPartitions expects. The
            // signal is smuggled through to model.complete via the prompt's
            // ModelParamSection (see jsonTranslator.ts).
            translator: {
                translate: (request: string, signal?: AbortSignal) =>
                    innerTranslator.translate(request, undefined, signal),
            },
        };
    });

    return {
        translate: (request: string, signal?: AbortSignal) =>
            selectFromPartitions(translators, request, signal),
    };
}
