import { AIMessage, AIMessageChunk, ChatMessage, ChatMessageChunk, FunctionMessageChunk, HumanMessageChunk, isAIMessage, SystemMessageChunk, } from "@langchain/core/messages";
import { BaseChatModel, } from "@langchain/core/language_models/chat_models";
import { GigaChat as GigaChatClient } from "gigachat";
import { zodToJsonSchema } from "zod-to-json-schema";
import { isLangChainTool } from "@langchain/core/utils/function_calling";
import { RunnablePassthrough, RunnableSequence, } from "@langchain/core/runnables";
import { ChatGenerationChunk } from "@langchain/core/outputs";
import { isZodSchema } from "@langchain/core/utils/types";
import { JsonOutputKeyToolsParser } from "@langchain/core/output_parsers/openai_tools";
import { v4 as uuidv4 } from "uuid";
function removeEmpty(obj) {
    const newObj = {};
    for (const key in obj) {
        if (obj[key] !== undefined)
            newObj[key] = obj[key];
    }
    return newObj;
}
function extractGenericMessageCustomRole(message) {
    if (message.role !== "system" &&
        message.role !== "assistant" &&
        message.role !== "user" &&
        message.role !== "function" &&
        message.role !== "function_in_progress" &&
        message.role !== "search_result") {
        console.warn(`Unknown message role: ${message.role}`);
    }
    return message.role;
}
function extractMessageContentString(content) {
    if (content.constructor === String) {
        return content;
    }
    else if (content.constructor === Array) {
        return content
            .filter((part) => part.type === "text")
            .map((part) => ("text" in part ? part.text : ""))
            .join(" ");
    }
    return "";
}
function messageToGigaChatRole(message) {
    const type = message._getType();
    switch (type) {
        case "system":
            return "system";
        case "ai":
            return "assistant";
        case "human":
            return "user";
        case "function":
            return "function";
        case "tool":
            return "function";
        case "generic": {
            if (!ChatMessage.isInstance(message))
                throw new Error("Invalid generic chat message");
            return extractGenericMessageCustomRole(message);
        }
        default:
            throw new Error(`Unknown message type: ${type}`);
    }
}
function gigachatResponseToChatMessage(completion, includeRawResponse) {
    const choice = completion.choices[0];
    const rawToolCalls = choice.message.function_call;
    switch (choice.message.role) {
        case "assistant": {
            const toolCalls = [];
            const additional_kwargs = {};
            if (choice.message.function_call) {
                toolCalls.push({
                    name: choice.message.function_call.name,
                    // eslint-disable-next-line @typescript-eslint/no-explicit-any
                    args: choice.message.function_call.arguments,
                    id: uuidv4(),
                    type: "tool_call",
                });
                additional_kwargs.function_call = {
                    name: choice.message.function_call?.name,
                    arguments: JSON.stringify(choice.message.function_call?.arguments),
                };
                additional_kwargs.tool_calls = rawToolCalls;
                additional_kwargs.function_state_id = choice.message.functions_state_id;
            }
            if (includeRawResponse !== undefined) {
                additional_kwargs.__raw_response = choice;
            }
            return new AIMessage({
                content: choice.message.content || "",
                tool_calls: toolCalls,
                additional_kwargs,
                response_metadata: {
                    xHeaders: completion.xHeaders,
                },
                usage_metadata: {
                    input_tokens: completion.usage.prompt_tokens,
                    output_tokens: completion.usage.completion_tokens,
                    total_tokens: completion.usage.total_tokens,
                },
                id: completion.xHeaders["xRequestID"] ?? uuidv4(),
            });
        }
        default:
            return new ChatMessage(choice.message.content || "", choice.message.role ?? "unknown");
    }
}
function _convertDeltaToMessageChunk(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
chunk, index, defaultRole, includeRawResponse) {
    const { delta } = chunk.choices[0];
    const role = delta.role ?? defaultRole;
    const content = delta.content ?? "";
    let additional_kwargs;
    if (delta.function_call) {
        additional_kwargs = {
            function_call: {
                name: delta.function_call.name,
                arguments: JSON.stringify(delta.function_call.arguments),
            },
        };
    }
    else {
        additional_kwargs = {};
    }
    if (includeRawResponse !== undefined) {
        additional_kwargs.__raw_response = chunk;
    }
    if (role === "user") {
        return new HumanMessageChunk({ content });
    }
    else if (role === "assistant") {
        const toolCallChunks = [];
        if (delta.function_call) {
            toolCallChunks.push({
                name: delta.function_call.name,
                args: JSON.stringify(delta.function_call.arguments),
                type: "tool_call_chunk",
                id: uuidv4(),
                index,
            });
        }
        return new AIMessageChunk({
            content,
            tool_call_chunks: toolCallChunks,
            response_metadata: {
                xHeaders: chunk.xHeaders,
            },
            additional_kwargs,
            id: chunk.xHeaders["xRequestID"] ?? uuidv4(),
        });
    }
    else if (role === "system") {
        return new SystemMessageChunk({ content });
    }
    else if (role === "function") {
        return new FunctionMessageChunk({
            content,
            additional_kwargs,
        });
    }
    else {
        return new ChatMessageChunk({
            content,
            role: role ?? defaultRole ?? "assistant",
        });
    }
}
// eslint-disable-next-line @typescript-eslint/no-explicit-any
function isGigaChatTool(tool) {
    return "name" in tool && "parameters" in tool;
}
/**
 * Integration with a chat model.
 */
export class GigaChat extends BaseChatModel {
    static lc_name() {
        return "GigaChat";
    }
    get lc_secrets() {
        return {
            credentials: "GIGACHAT_CREDENTIALS",
            access_token: "GIGACHAT_ACCESS_TOKEN",
            password: "GIGACHAT_PASSWORD",
            key_file_password: "GIGACHAT_KEY_FILE_PASSWORD",
        };
    }
    get lc_aliases() {
        return {
            credentials: "GIGACHAT_CREDENTIALS",
            access_token: "GIGACHAT_ACCESS_TOKEN",
            user: "GIGACHAT_USER",
            password: "GIGACHAT_PASSWORD",
            scope: "GIGACHAT_SCOPE",
            key_file_password: "GIGACHAT_KEY_FILE_PASSWORD",
        };
    }
    _convertMessageToPayload(_messages) {
        return _messages.map((_message) => {
            const role = messageToGigaChatRole(_message);
            let content = extractMessageContentString(_message.content);
            if (role === "function") {
                content = JSON.stringify(content);
            }
            let function_call;
            if (isAIMessage(_message) && _message.tool_calls?.length) {
                function_call = {
                    name: _message.tool_calls[0].name,
                    arguments: _message.tool_calls[0].args,
                };
            }
            else if (_message.additional_kwargs.function_call) {
                function_call = {
                    name: _message.additional_kwargs.function_call.name,
                    arguments: JSON.parse(_message.additional_kwargs.function_call.arguments),
                };
            }
            const message = {
                role,
                content,
                function_call,
                attachments: _message.additional_kwargs.attachments ?? undefined,
                functions_state_id: _message.additional_kwargs.functions_state_id ??
                    undefined,
            };
            return message;
        });
    }
    getLsParams(options) {
        const params = this.invocationParams(options);
        return {
            ls_provider: "giga-chat-model",
            ls_model_name: this.model,
            ls_model_type: "chat",
            ls_temperature: params.temperature ?? undefined,
            ls_max_tokens: params.max_tokens ?? undefined,
            ls_stop: options.stop,
        };
    }
    /**
     * Get the parameters used to invoke the model
     */
    invocationParams(options) {
        const tool_choice = options?.tool_choice;
        return {
            model: options?.model ?? this.model,
            temperature: options?.temperature ?? this.temperature,
            max_tokens: options?.maxTokens ?? this.maxTokens,
            top_p: options?.topP ?? this.topP,
            repetitionPenalty: options?.repetitionPenalty ?? this.repetitionPenalty,
            update_interval: options?.updateInterval ?? this.updateInterval,
            stop_sequences: options?.stop ?? this.stopSequence,
            stream: this.streaming,
            functions: this.formatStructuredToolToGigaChat(options?.tools),
            function_call: tool_choice,
            ...this.invocationKwargs,
        };
    }
    constructor(fields) {
        super(fields ?? {});
        Object.defineProperty(this, "lc_serializable", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: true
        });
        Object.defineProperty(this, "model", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "GigaChat"
        });
        Object.defineProperty(this, "useApiForTokens", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: false
        });
        Object.defineProperty(this, "streaming", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: false
        });
        Object.defineProperty(this, "verbose", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: false
        });
        Object.defineProperty(this, "temperature", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "maxTokens", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "topP", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "repetitionPenalty", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "updateInterval", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "stopSequence", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "invocationKwargs", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "clientConfig", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "_client", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this.model = fields?.model ?? this.model;
        this.useApiForTokens = fields?.useApiForTokens ?? this.useApiForTokens;
        this.streaming = fields?.streaming ?? this.streaming;
        this.verbose = fields?.verbose ?? this.verbose;
        this.temperature = fields?.temperature ?? this.temperature;
        this.maxTokens = fields?.maxTokens ?? this.maxTokens;
        this.topP = fields?.topP ?? this.topP;
        this.repetitionPenalty =
            fields?.repetitionPenalty ?? this.repetitionPenalty;
        this.updateInterval = fields?.updateInterval ?? this.updateInterval;
        this.stopSequence = fields?.stopSequence ?? this.stopSequence;
        this.invocationKwargs = fields?.invocationKwargs ?? this.invocationKwargs;
        this.clientConfig = {
            baseUrl: fields?.baseUrl,
            authUrl: fields?.authUrl,
            credentials: fields?.credentials,
            scope: fields?.scope,
            accessToken: fields?.accessToken,
            model: fields?.model,
            profanityCheck: fields?.profanityCheck,
            user: fields?.user,
            password: fields?.password,
            timeout: fields?.timeout,
            verbose: fields?.verbose,
            flags: fields?.flags,
            httpsAgent: fields?.httpsAgent,
        };
        this.clientConfig = removeEmpty(this.clientConfig);
        this._client = new GigaChatClient(this.clientConfig);
    }
    _llmType() {
        return "giga-chat-model";
    }
    bindTools(tools, kwargs) {
        return this.bind({
            tools: this.formatStructuredToolToGigaChat(tools),
            ...kwargs,
        });
    }
    /**
     * Formats LangChain StructuredTools to GigaChat Functions.
     *
     * @param {ChatGigaChatToolType[] | undefined} tools The tools to format
     * @returns {_Function[] | undefined} The formatted tools, or undefined if none are passed.
     */
    formatStructuredToolToGigaChat(tools) {
        if (!tools || !tools.length) {
            return undefined;
        }
        return tools.map((tool) => {
            if (isGigaChatTool(tool)) {
                return tool;
            }
            if (isLangChainTool(tool)) {
                return {
                    name: tool.name,
                    description: tool.description,
                    parameters: zodToJsonSchema(tool.schema),
                };
            }
            throw new Error(`Unknown tool type passed to GigaChat: ${JSON.stringify(tool, null, 2)}`);
        });
    }
    _combineLLMOutput(...llmOutputs) {
        return llmOutputs.reduce((acc, llmOutput) => {
            if (llmOutput && llmOutput.usage) {
                acc.usage.completion_tokens += llmOutput.usage.completion_tokens ?? 0;
                acc.usage.prompt_tokens += llmOutput.usage.prompt_tokens ?? 0;
                acc.usage.total_tokens += llmOutput.usage.total_tokens ?? 0;
            }
            return acc;
        }, {
            usage: {
                completion_tokens: 0,
                prompt_tokens: 0,
                total_tokens: 0,
            },
        });
    }
    identifyingParams() {
        return {
            model_name: this.model,
            ...this.invocationParams(),
        };
    }
    async *_streamResponseChunks(messages, options, runManager) {
        const params = this.invocationParams(options);
        const formattedMessages = this._convertMessageToPayload(messages);
        const stream = await this.createStreamWithRetry({
            ...params,
            messages: formattedMessages,
            stream: true,
        }, options.signal);
        if (!stream) {
            return;
        }
        let index = 0;
        for await (const data of stream) {
            if (options.signal?.aborted) {
                throw new Error("AbortError: User aborted the request.");
            }
            const chunk = _convertDeltaToMessageChunk(data, index);
            const generationChunk = new ChatGenerationChunk({
                message: chunk,
                text: data.choices[0].delta.content ?? "",
            });
            yield generationChunk;
            await runManager?.handleLLMNewToken(data.choices[0].delta.content ?? "", undefined, undefined, undefined, undefined, { chunk: generationChunk });
            index += 1;
        }
    }
    /**
     * Creates a streaming request with retry.
     * @param request The parameters for creating a completion.
     * @returns A streaming request.
     */
    async createStreamWithRetry(request, signal) {
        const makeCompletionRequest = async () => {
            try {
                return this._client?.stream(request, signal);
            }
            catch (error) {
                console.error(error);
                throw error;
            }
        };
        return this.caller.call(makeCompletionRequest);
    }
    async completionWithRetry(request, options) {
        const makeCompletionRequest = async () => {
            try {
                return await this._client?.chat(request);
            }
            catch (error) {
                console.error(error);
                throw error;
            }
        };
        return this.caller.callWithOptions({ signal: options.signal ?? undefined }, makeCompletionRequest);
    }
    /** @ignore */
    async _generateNonStreaming(messages, params, requestOptions) {
        const response = await this.completionWithRetry({
            ...params,
            stream: false,
            messages: this._convertMessageToPayload(messages),
        }, requestOptions);
        const generation = gigachatResponseToChatMessage(response);
        return {
            generations: [
                {
                    message: generation,
                    text: extractMessageContentString(generation.content),
                    generationInfo: {
                        finish_reason: response.choices[0].finish_reason,
                    },
                },
            ],
            llmOutput: {
                tokenUsage: {
                    input_tokens: response.usage.prompt_tokens,
                    output_tokens: response.usage.completion_tokens,
                    total_tokens: response.usage.total_tokens,
                },
            },
        };
    }
    /** @ignore */
    async _generate(messages, options, runManager) {
        if (this.stopSequence && options.stop) {
            throw new Error(`"stopSequence" parameter found in input and default params`);
        }
        const params = this.invocationParams(options);
        if (params.stream) {
            let finalChunk;
            const stream = this._streamResponseChunks(messages, options, runManager);
            for await (const chunk of stream) {
                if (finalChunk === undefined) {
                    finalChunk = chunk;
                }
                else {
                    finalChunk = finalChunk.concat(chunk);
                }
            }
            if (finalChunk === undefined) {
                throw new Error("No chunks returned from GigaChat API.");
            }
            return {
                generations: [
                    {
                        text: finalChunk.text,
                        message: finalChunk.message,
                    },
                ],
            };
        }
        else {
            return this._generateNonStreaming(messages, params, options);
        }
    }
    withStructuredOutput(outputSchema, config) {
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        const schema = outputSchema;
        const name = config?.name;
        const method = config?.method;
        const includeRaw = config?.includeRaw;
        if (method === "jsonMode" || method === "jsonSchema") {
            throw new Error(`Anthropic only supports "functionCalling" as a method.`);
        }
        let functionName = name ?? "extract";
        let outputParser;
        let tools;
        if (isZodSchema(schema)) {
            const jsonSchema = zodToJsonSchema(schema);
            tools = [
                {
                    name: functionName,
                    description: jsonSchema.description ?? "A function available to call.",
                    parameters: jsonSchema,
                },
            ];
            outputParser = new JsonOutputKeyToolsParser({
                returnSingle: true,
                keyName: functionName,
                zodSchema: schema,
            });
        }
        else {
            let gigachatTools;
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            const schema_ = schema;
            if (typeof schema_.name === "string" &&
                typeof schema_.description === "string" &&
                typeof schema_.parameters === "object" &&
                schema_.parameters != null) {
                gigachatTools = schema;
                functionName = schema_.name;
            }
            else {
                gigachatTools = {
                    name: functionName,
                    description: schema.description ?? "",
                    parameters: schema,
                };
            }
            tools = [gigachatTools];
            outputParser = new JsonOutputKeyToolsParser({
                returnSingle: true,
                keyName: functionName,
            });
        }
        const llm = this.bindTools(tools, {
            tool_choice: { name: functionName },
            ...config,
        });
        if (!includeRaw) {
            return llm.pipe(outputParser).withConfig({
                runName: "GigaChatStructuredOutput",
            });
        }
        const parserAssign = RunnablePassthrough.assign({
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            parsed: (input, config) => outputParser.invoke(input.raw, config),
        });
        const parserNone = RunnablePassthrough.assign({
            parsed: () => null,
        });
        const parsedWithFallback = parserAssign.withFallbacks({
            fallbacks: [parserNone],
        });
        return RunnableSequence.from([
            {
                raw: llm,
            },
            parsedWithFallback,
        ]).withConfig({
            runName: "StructuredOutputRunnable",
        });
    }
}
