"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
    var ownKeys = function(o) {
        ownKeys = Object.getOwnPropertyNames || function (o) {
            var ar = [];
            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
            return ar;
        };
        return ownKeys(o);
    };
    return function (mod) {
        if (mod && mod.__esModule) return mod;
        var result = {};
        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
        __setModuleDefault(result, mod);
        return result;
    };
})();
Object.defineProperty(exports, "__esModule", { value: true });
exports.GenericGeminiClient = void 0;
const genai_1 = require("@google/genai");
const fsSync = __importStar(require("fs"));
const pathSync = __importStar(require("path"));
const utils_1 = require("../utils");
const types_1 = require("../types");
const pricing_1 = require("./pricing");
const contextLimits_1 = require("./contextLimits");
function getMimeTypeFromUrl(url) {
    if (url.endsWith(".png"))
        return "image/png";
    if (url.endsWith(".gif"))
        return "image/gif";
    if (url.endsWith(".webp"))
        return "image/webp";
    return "image/jpeg";
}
function getVideoMimeTypeFromUrl(url) {
    if (url.endsWith(".mp4"))
        return "video/mp4";
    if (url.endsWith(".webm"))
        return "video/webm";
    if (url.endsWith(".mov"))
        return "video/quicktime";
    if (url.endsWith(".avi"))
        return "video/x-msvideo";
    return "video/mp4";
}
function pcmToWav(pcmData, sampleRate = 24000, numChannels = 1, bitsPerSample = 16) {
    const dataSize = pcmData.length;
    const headerSize = 44;
    const wavBuffer = Buffer.alloc(headerSize + dataSize);
    wavBuffer.write("RIFF", 0);
    wavBuffer.writeUInt32LE(36 + dataSize, 4);
    wavBuffer.write("WAVE", 8);
    wavBuffer.write("fmt ", 12);
    wavBuffer.writeUInt32LE(16, 16);
    wavBuffer.writeUInt16LE(1, 20);
    wavBuffer.writeUInt16LE(numChannels, 22);
    wavBuffer.writeUInt32LE(sampleRate, 24);
    wavBuffer.writeUInt32LE(sampleRate * numChannels * (bitsPerSample / 8), 28);
    wavBuffer.writeUInt16LE(numChannels * (bitsPerSample / 8), 32);
    wavBuffer.writeUInt16LE(bitsPerSample, 34);
    wavBuffer.write("data", 36);
    wavBuffer.writeUInt32LE(dataSize, 40);
    pcmData.copy(wavBuffer, 44);
    return wavBuffer;
}
class GenericGeminiClient {
    client;
    apiKey;
    constructor(apiKey) {
        this.setKey(apiKey || process.env.GEMINI_API_KEY || "");
    }
    setKey(apiKey) {
        this.apiKey = apiKey;
        this.client = new genai_1.GoogleGenAI({
            apiKey: apiKey || process.env.GEMINI_API_KEY,
        });
    }
    transformContentParts(content) {
        if (typeof content === "string") {
            return [{ text: content }];
        }
        return content
            .map((part) => {
            if (part.type === "text") {
                return { text: part.text };
            }
            if (part.type === "image_url") {
                const url = part.image_url.url;
                if (url.startsWith("data:")) {
                    const [header, base64Data] = url.split(",");
                    const mimeType = header.split(":")[1].split(";")[0];
                    return {
                        inlineData: {
                            data: base64Data,
                            mimeType,
                        },
                    };
                }
                if (url.startsWith("https://generativelanguage.googleapis.com")) {
                    return {
                        fileData: {
                            fileUri: url,
                            mimeType: getMimeTypeFromUrl(url),
                        },
                    };
                }
            }
            if (part.type === "video_url") {
                const mimeType = getVideoMimeTypeFromUrl(part.video_url.url);
                return {
                    fileData: {
                        fileUri: part.video_url.url,
                        mimeType,
                    },
                };
            }
            console.warn(`Unsupported generic message content part type: ${part.type}`);
            return { text: `[Unsupported content type: ${part.type}]` };
        })
            .filter((part) => !!part);
    }
    transformMessages(messages) {
        const googleContents = [];
        let systemInstruction;
        const assistantToolCalls = {};
        for (const msg of messages) {
            if (msg.role === "system") {
                if (typeof msg.content === "string") {
                    systemInstruction =
                        (systemInstruction ? systemInstruction + "\n" : "") + msg.content;
                }
                else {
                    systemInstruction =
                        (systemInstruction ? systemInstruction + "\n" : "") +
                            this.transformContentParts(msg.content)
                                .filter((p) => "text" in p && typeof p.text === "string")
                                .map((p) => p.text)
                                .join("\n");
                }
            }
            else if (msg.role === "user" || msg.role === "assistant") {
                const parts = msg.content
                    ? this.transformContentParts(msg.content)
                    : [];
                if (msg.role === "assistant" &&
                    msg.tool_calls &&
                    msg.tool_calls.length > 0) {
                    for (const toolCall of msg.tool_calls) {
                        parts.push({
                            functionCall: {
                                name: toolCall.function.name,
                                args: JSON.parse(toolCall.function.arguments || "{}"),
                            },
                        });
                        assistantToolCalls[toolCall.id] = toolCall;
                    }
                }
                if (parts.length > 0) {
                    googleContents.push({
                        role: msg.role === "user" ? "user" : "model",
                        parts,
                    });
                }
            }
            else if (msg.role === "tool") {
                if (!msg.tool_call_id) {
                    throw new Error("Tool message must have a tool_call_id.");
                }
                const toolOutputContent = typeof msg.content === "string"
                    ? msg.content
                    : JSON.stringify(msg.content);
                const matchingToolCall = assistantToolCalls[msg.tool_call_id];
                const functionName = matchingToolCall
                    ? matchingToolCall.function.name
                    : "unknown_function";
                if (!matchingToolCall) {
                    console.warn(`Matching assistant tool call not found for tool_call_id: ${msg.tool_call_id}. Using name '${functionName}'.`, msg);
                }
                googleContents.push({
                    role: "user",
                    parts: [
                        {
                            functionResponse: {
                                name: functionName,
                                response: {
                                    result: toolOutputContent,
                                },
                            },
                        },
                    ],
                });
            }
            else {
                console.warn(`Unsupported generic message role: ${msg.role}, skipping.`);
            }
        }
        return { systemInstruction, contents: googleContents };
    }
    cleanSchemaForGemini(schema) {
        if (!schema || typeof schema !== "object") {
            return schema;
        }
        if (Array.isArray(schema)) {
            return schema.map((item) => this.cleanSchemaForGemini(item));
        }
        const cleaned = {};
        for (const key in schema) {
            if (!Object.prototype.hasOwnProperty.call(schema, key)) {
                continue;
            }
            if (key === "additionalProperties" ||
                key === "$ref" ||
                key === "$defs" ||
                key === "positional") {
                continue;
            }
            const value = schema[key];
            if (key === "type" && typeof value === "string") {
                cleaned[key] = value.toUpperCase();
            }
            else if (key === "type" && Array.isArray(value)) {
                cleaned[key] = value.map((t) => typeof t === "string" ? t.toUpperCase() : t);
            }
            else if (typeof value === "object" && value !== null) {
                cleaned[key] = this.cleanSchemaForGemini(value);
            }
            else {
                cleaned[key] = value;
            }
        }
        return cleaned;
    }
    transformTools(tools) {
        if (!tools || tools.length === 0) {
            return undefined;
        }
        const functionDeclarations = tools.map((tool) => {
            const cleanedParameters = this.cleanSchemaForGemini(tool.function.parameters);
            return {
                name: tool.function.name,
                description: tool.function.description || "",
                parameters: cleanedParameters,
            };
        });
        return [{ functionDeclarations }];
    }
    buildThinkingConfig(options) {
        const model = options.model;
        const effort = options.reasoning_effort ?? "low";
        if (types_1.GoogleThinkingLevelModels.includes(model)) {
            const levelMap = {
                low: "low",
                medium: "medium",
                high: "high",
            };
            return {
                thinkingLevel: levelMap[effort] ?? "low",
            };
        }
        if (types_1.GoogleThinkingBudgetModels.includes(model)) {
            const budgetMap = {
                low: 1024,
                medium: 8192,
                high: -1,
            };
            return {
                thinkingBudget: budgetMap[effort] ?? 1024,
            };
        }
        return undefined;
    }
    async createChatCompletion(options) {
        const thinkingConfig = this.buildThinkingConfig(options);
        const { systemInstruction, contents } = this.transformMessages(options.messages);
        try {
            await (0, utils_1.wait)(2000);
            const response = await this.client.models.generateContent({
                model: options.model,
                contents,
                config: {
                    systemInstruction,
                    thinkingConfig,
                    tools: this.transformTools(options.tools),
                    maxOutputTokens: options.max_tokens,
                },
            });
            let toolCalls = [];
            if (response.functionCalls) {
                for (const call of response.functionCalls) {
                    toolCalls.push({
                        id: call.id ||
                            `fc_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`,
                        type: "function",
                        function: {
                            name: call.name,
                            arguments: JSON.stringify(call.args || {}),
                        },
                    });
                }
            }
            if (response?.promptFeedback?.blockReason) {
                throw new Error(`Google GenAI blocked the response due to: ${response.promptFeedback.blockReason}`);
            }
            const choices = response.candidates?.map((candidate) => {
                const message = {
                    role: candidate.content.role === "model" ? "assistant" : "user",
                    content: "",
                    tool_calls: [...toolCalls],
                };
                let textContent = "";
                if (toolCalls.length) {
                    toolCalls = [];
                }
                if (!candidate?.content?.parts) {
                    console.warn("No content parts in candidate:", candidate);
                    return { message };
                }
                candidate?.content?.parts?.forEach((part) => {
                    if ("text" in part && typeof part.text === "string") {
                        textContent += part.text;
                    }
                    else if ("functionCall" in part && part.functionCall) {
                        message.tool_calls.push({
                            id: part.functionCall.id ||
                                `fc_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`,
                            type: "function",
                            function: {
                                name: part.functionCall.name,
                                arguments: JSON.stringify(part.functionCall.args || {}),
                            },
                        });
                    }
                });
                message.content = textContent || null;
                return { message };
            }) || [];
            const usage = response.usageMetadata;
            const usdCost = usage
                ? this.calculateCost(options.model, usage)
                : undefined;
            const cachedTokens = usage?.cachedContentTokenCount ?? 0;
            const usageWithCache = usage
                ? {
                    prompt_tokens: usage.promptTokenCount ?? 0,
                    completion_tokens: usage.candidatesTokenCount ?? 0,
                    total_tokens: usage.totalTokenCount,
                    prompt_tokens_details: { cached_tokens: cachedTokens },
                }
                : undefined;
            return {
                choices,
                model: options.model,
                usage: usageWithCache,
                usd_cost: usdCost,
            };
        }
        catch (error) {
            console.error("Error calling Google GenAI generateContent:", error.message);
            throw error;
        }
    }
    pricesPerMillion() {
        return pricing_1.GeminiTextPricing;
    }
    calculateCost(model, usage) {
        const pricing = this.pricesPerMillion()[model];
        if (!pricing || !usage) {
            return 0;
        }
        let cost = 0;
        const inputDetails = usage.inputTokensDetails;
        if (inputDetails &&
            inputDetails.length > 0 &&
            pricing.input_audio !== undefined) {
            for (const detail of inputDetails) {
                const tokens = detail.tokenCount || 0;
                if (!tokens)
                    continue;
                const modality = (detail.modality || "").toUpperCase();
                if (modality === "AUDIO") {
                    const rate = pricing.input_audio ?? pricing.input ?? 0;
                    cost += (tokens * rate) / 1e6;
                }
                else {
                    const isGt200k = tokens > 200_000 && pricing.input_gt_200k;
                    const rate = isGt200k ? pricing.input_gt_200k : pricing.input ?? 0;
                    cost += (tokens * rate) / 1e6;
                }
            }
        }
        else if ("promptTokenCount" in usage && usage.promptTokenCount) {
            const isGt200k = usage.promptTokenCount > 200_000 && pricing.input_gt_200k;
            const rate = isGt200k ? pricing.input_gt_200k : pricing.input ?? 0;
            cost += (usage.promptTokenCount * rate) / 1e6;
        }
        const outputDetails = usage.outputTokensDetails;
        if (outputDetails &&
            outputDetails.length > 0 &&
            pricing.output_audio !== undefined) {
            for (const detail of outputDetails) {
                const tokens = detail.tokenCount || 0;
                if (!tokens)
                    continue;
                const modality = (detail.modality || "").toUpperCase();
                if (modality === "AUDIO") {
                    const rate = pricing.output_audio ?? pricing.output ?? 0;
                    cost += (tokens * rate) / 1e6;
                }
                else {
                    const isGt200k = tokens > 200_000 && pricing.output_gt_200k;
                    const rate = isGt200k ? pricing.output_gt_200k : pricing.output ?? 0;
                    cost += (tokens * rate) / 1e6;
                }
            }
        }
        else if ("responseTokenCount" in usage && usage.responseTokenCount) {
            const isGt200k = usage.responseTokenCount > 200_000 && pricing.output_gt_200k;
            const rate = isGt200k ? pricing.output_gt_200k : pricing.output ?? 0;
            cost += (usage.responseTokenCount * rate) / 1e6;
        }
        if ("cachedContentTokenCount" in usage &&
            usage.cachedContentTokenCount &&
            pricing.context_caching) {
            if (usage.cachedContentTokenCount > 200_000 &&
                pricing.context_caching_gt_200k) {
                cost +=
                    (usage.cachedContentTokenCount * pricing.context_caching_gt_200k) /
                        1e6;
            }
            else {
                cost += (usage.cachedContentTokenCount * pricing.context_caching) / 1e6;
            }
        }
        return cost;
    }
    async getModels(modality) {
        if (modality) {
            const map = {
                completion: types_1.GoogleReasoningModels,
                embedding: types_1.GoogleEmbeddingModelsList,
                image: types_1.GoogleImageModels,
                audio: types_1.GoogleTTSModels,
                video: types_1.GoogleVideoModels,
            };
            return (map[modality] ?? []).map((id) => ({ id }));
        }
        try {
            const models = await this.client.models.list();
            return models.page.map((m) => ({ id: m.name }));
        }
        catch (error) {
            console.error("Error fetching Google GenAI models:", error);
            throw error;
        }
    }
    async createEmbedding(options) {
        if (!options.model) {
            console.warn("Embedding model not specified, using default 'text-embedding-004'.");
        }
        try {
            const googleEmbedding = await this.client.models.embedContent({
                model: options.model,
                contents: options.input,
            });
            const data = googleEmbedding.embeddings.map((e, index) => ({
                object: "embedding",
                embedding: e.values,
                index,
            }));
            const usage = {
                promptTokenCount: googleEmbedding.metadata.billableCharacterCount || 0,
                totalTokenCount: googleEmbedding.metadata.billableCharacterCount || 0,
            };
            const usdCost = this.calculateCost(options.model, usage);
            return {
                data,
                model: options.model,
                usage: {
                    prompt_tokens: usage.promptTokenCount,
                    total_tokens: usage.totalTokenCount,
                },
                usd_cost: usdCost,
            };
        }
        catch (error) {
            console.error("Error calling Google GenAI embedContent:", error);
            throw error;
        }
    }
    async createAudioTranscription(options) {
        throw new Error("Audio transcription is not yet supported by the Gemini client. Use OpenAI client with Whisper model instead.");
    }
    async createAudioGeneration(options) {
        try {
            const response = await this.client.models.generateContent({
                model: options.model,
                contents: [
                    {
                        role: "user",
                        parts: [{ text: options.input }],
                    },
                ],
                config: {
                    responseModalities: ["AUDIO"],
                    speechConfig: {
                        voiceConfig: {
                            prebuiltVoiceConfig: {
                                voiceName: options.voice || "Puck",
                            },
                        },
                    },
                },
            });
            const audioPart = response.candidates?.[0]?.content?.parts?.find((part) => part.inlineData?.mimeType?.startsWith("audio/"));
            if (!audioPart || !audioPart.inlineData) {
                throw new Error("No audio data returned from Gemini TTS");
            }
            const rawBuffer = Buffer.from(audioPart.inlineData.data, "base64");
            const mimeType = audioPart.inlineData.mimeType || "audio/wav";
            let audioBuffer = rawBuffer;
            if (mimeType.includes("L16") || mimeType.includes("pcm")) {
                const rateMatch = mimeType.match(/rate=(\d+)/);
                const sampleRate = rateMatch ? parseInt(rateMatch[1], 10) : 24000;
                audioBuffer = pcmToWav(rawBuffer, sampleRate);
            }
            return {
                audio: audioBuffer,
                format: "audio/wav",
            };
        }
        catch (error) {
            console.error("Error calling Gemini TTS:", error);
            throw error;
        }
    }
    async createImageGeneration(options) {
        try {
            const isImagen3 = options.model?.includes("imagen");
            if (isImagen3) {
                const response = await this.client.models.generateImages({
                    model: options.model,
                    prompt: options.prompt,
                    config: {
                        numberOfImages: options.n || 1,
                    },
                });
                const generatedImages = response.generatedImages || [];
                const images = generatedImages.map((img) => ({
                    b64_json: img.image?.imageBytes ? img.image.imageBytes : "",
                    revised_prompt: options.prompt,
                }));
                return {
                    created: Math.floor(Date.now() / 1000),
                    data: images,
                    usd_cost: 0.03 * images.length,
                };
            }
            else {
                const response = await this.client.models.generateContent({
                    model: options.model,
                    contents: [
                        {
                            role: "user",
                            parts: [{ text: options.prompt }],
                        },
                    ],
                    config: {
                        responseModalities: ["IMAGE", "TEXT"],
                    },
                });
                const imageParts = response.candidates?.[0]?.content?.parts?.filter((part) => part.inlineData?.mimeType?.startsWith("image/")) || [];
                if (imageParts.length === 0) {
                    throw new Error("No image data returned from Gemini");
                }
                const images = imageParts.map((part) => ({
                    b64_json: part.inlineData.data,
                    revised_prompt: options.prompt,
                }));
                const usageMetadata = response.usageMetadata;
                const usdCost = usageMetadata
                    ? this.calculateCost(options.model, usageMetadata)
                    : undefined;
                return {
                    created: Math.floor(Date.now() / 1000),
                    data: images,
                    usd_cost: usdCost,
                };
            }
        }
        catch (error) {
            console.error("Error calling Gemini image generation:", error);
            throw error;
        }
    }
    async createVideoGeneration(options) {
        try {
            const operation = await this.client.models.generateVideos({
                model: options.model,
                prompt: options.prompt,
                config: {
                    numberOfVideos: options.n || 1,
                    ...(options.duration && {
                        durationSeconds: Math.max(6, options.duration),
                    }),
                    ...(options.resolution && { resolution: options.resolution }),
                    ...(options.aspect_ratio && { aspectRatio: options.aspect_ratio }),
                },
            });
            const duration = options.duration || 5;
            const pricingEntry = pricing_1.GeminiTextPricing[options.model];
            const ratePerSec = pricingEntry?.video_generation ?? 0.4;
            const usdCost = (options.n || 1) * duration * ratePerSec;
            return {
                created: Math.floor(Date.now() / 1000),
                data: [],
                jobId: operation.name,
                usd_cost: usdCost,
            };
        }
        catch (error) {
            console.error("Error calling Gemini video generation:", error);
            throw error;
        }
    }
    async getVideoStatus(options) {
        try {
            const operation = await this.client.operations.getVideosOperation({
                operation: { name: options.jobId },
            });
            if (operation.error) {
                return {
                    jobId: options.jobId,
                    status: "failed",
                    error: JSON.stringify(operation.error),
                };
            }
            if (!operation.done) {
                return {
                    jobId: options.jobId,
                    status: "in_progress",
                };
            }
            const generatedVideos = operation.response?.generatedVideos || [];
            const data = generatedVideos.map((vid) => {
                const videoBytes = vid.video?.videoBytes;
                const uri = vid.video?.uri;
                return {
                    b64_json: videoBytes || undefined,
                    url: uri || undefined,
                    fileUri: uri || undefined,
                };
            });
            return {
                jobId: options.jobId,
                status: "completed",
                data,
            };
        }
        catch (error) {
            console.error("Error checking Gemini video status:", error);
            throw error;
        }
    }
    async downloadVideo(options) {
        return this.downloadFile(options);
    }
    async uploadFile(options) {
        try {
            const blob = new Blob([new Uint8Array(options.data)], { type: options.mimeType });
            const uploadedFile = await this.client.files.upload({
                file: blob,
                config: {
                    mimeType: options.mimeType,
                    displayName: options.displayName,
                    name: options.fileName,
                },
            });
            return {
                fileId: uploadedFile.name,
                uri: uploadedFile.uri,
                url: uploadedFile.downloadUri || uploadedFile.uri,
                mimeType: uploadedFile.mimeType,
                sizeBytes: uploadedFile.sizeBytes
                    ? Number(uploadedFile.sizeBytes)
                    : undefined,
            };
        }
        catch (error) {
            console.error("Error uploading file to Google GenAI Files API:", error);
            throw error;
        }
    }
    async downloadFile(options) {
        const mimeMap = {
            ".mp4": "video/mp4",
            ".webm": "video/webm",
            ".mov": "video/quicktime",
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".gif": "image/gif",
            ".wav": "audio/wav",
            ".mp3": "audio/mpeg",
        };
        try {
            const rawUri = options.uri || options.fileId || "";
            let downloadUrl;
            if (rawUri.startsWith("https://")) {
                const sep = rawUri.includes("?") ? "&" : "?";
                downloadUrl = `${rawUri}${sep}key=${this.apiKey}`;
            }
            else {
                const fileId = rawUri.replace(/^files\//, "");
                downloadUrl = `https://generativelanguage.googleapis.com/v1beta/files/${fileId}:download?alt=media&key=${this.apiKey}`;
            }
            const response = await fetch(downloadUrl);
            if (!response.ok) {
                throw new Error(`HTTP ${response.status} ${response.statusText} downloading ${downloadUrl}`);
            }
            const arrayBuffer = await response.arrayBuffer();
            const data = Buffer.from(arrayBuffer);
            if (options.filePath) {
                fsSync.mkdirSync(pathSync.dirname(options.filePath), {
                    recursive: true,
                });
                fsSync.writeFileSync(options.filePath, data);
            }
            const sourceForExt = options.uri || options.fileId || options.filePath || "";
            const ext = pathSync.extname(sourceForExt.split("?")[0]).toLowerCase();
            const mimeType = mimeMap[ext] || "video/mp4";
            return { data, mimeType };
        }
        catch (error) {
            console.error("Error downloading file from Google GenAI Files API:", error);
            throw error;
        }
    }
    getContextLimit(model) {
        const contextLimit = contextLimits_1.ContextLimits[model];
        if (contextLimit === undefined)
            return undefined;
        const pricing = pricing_1.GeminiTextPricing[model];
        const threshold = pricing && "input_gt_200k" in pricing ? 200_000 : contextLimit;
        return { contextLimit, threshold };
    }
}
exports.GenericGeminiClient = GenericGeminiClient;
//# sourceMappingURL=gemini.js.map