"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.GenericXAIClient = void 0;
const openai_1 = __importDefault(require("openai"));
const pricing_1 = require("./pricing");
const contextLimits_1 = require("./contextLimits");
const types_1 = require("../types");
const xai_1 = require("./pricing/xai");
class GenericXAIClient {
    client;
    apiKey;
    constructor(apiKey = process.env.XAI_API_KEY) {
        this.apiKey = apiKey || "";
        this.client = new openai_1.default({
            apiKey: apiKey || process.env.XAI_API_KEY,
            baseURL: "https://api.x.ai/v1",
        });
    }
    setKey(apiKey) {
        this.apiKey = apiKey;
        this.client = new openai_1.default({
            apiKey,
            baseURL: "https://api.x.ai/v1",
            timeout: 60,
        });
    }
    async createChatCompletion(options) {
        if (xai_1.XaiResponsesOnlyModels.includes(options.model)) {
            return this.createChatResponse(options);
        }
        const xaiMessages = options.messages.map((msg) => {
            if (msg.role === "tool") {
                return {
                    ...msg,
                    content: msg.content || "",
                    role: "tool",
                    tool_call_id: msg.tool_call_id,
                };
            }
            return msg;
        });
        const response = await this.client.chat.completions.create({
            model: options.model,
            messages: xaiMessages,
            max_tokens: options.max_tokens,
            ...(xai_1.XaiReasoningModels.includes(options.model) && options.reasoning_effort && {
                reasoning_effort: options.reasoning_effort,
            }),
            ...(options.tools && {
                tools: options.tools,
                tool_choice: "auto",
            }),
        });
        const usdCost = this.calculateCost(options.model, response.usage);
        return {
            choices: response.choices.map((choice) => ({
                message: {
                    role: choice.message?.role || "assistant",
                    content: choice.message?.content || null,
                    tool_calls: choice.message?.tool_calls
                        ? choice.message.tool_calls
                        : undefined,
                },
            })),
            model: options.model,
            usage: response.usage ? {
                prompt_tokens: response.usage.prompt_tokens ?? 0,
                completion_tokens: response.usage.completion_tokens ?? 0,
                total_tokens: response.usage.total_tokens,
                prompt_tokens_details: {
                    cached_tokens: response.usage.prompt_tokens_details?.cached_tokens ?? 0,
                },
            } : undefined,
            usd_cost: usdCost,
        };
    }
    async createChatResponse(options) {
        const apiKey = this.apiKey || process.env.XAI_API_KEY;
        if (!apiKey) {
            throw new Error("XAI API key not set");
        }
        const systemMessages = options.messages.filter((m) => m.role === "system");
        const nonSystemMessages = options.messages.filter((m) => m.role !== "system");
        const instructions = systemMessages
            .map((m) => (typeof m.content === "string" ? m.content : ""))
            .join("\n")
            .trim() || undefined;
        const input = nonSystemMessages.map((msg) => {
            if (msg.role === "tool") {
                return {
                    type: "function_call_output",
                    call_id: msg.tool_call_id,
                    output: typeof msg.content === "string" ? msg.content : JSON.stringify(msg.content),
                };
            }
            if (msg.role === "assistant" && msg.tool_calls?.length) {
                return msg.tool_calls.map((tc) => ({
                    type: "function_call",
                    id: tc.id.startsWith("fc") ? tc.id : `fc_${tc.id}`,
                    call_id: tc.id,
                    name: tc.function.name,
                    arguments: tc.function.arguments,
                }));
            }
            return {
                role: msg.role,
                content: typeof msg.content === "string" ? msg.content : JSON.stringify(msg.content),
            };
        }).flat();
        const tools = options.tools?.map((tool) => ({
            type: "function",
            name: tool.function.name,
            description: tool.function.description,
            parameters: tool.function.parameters,
            strict: false,
        }));
        const pricing = pricing_1.XaiTextPricing[options.model];
        const supportedLevels = pricing?.reasoningLevels;
        let reasoningEffort = options.reasoning_effort;
        if (supportedLevels?.length) {
            if (!reasoningEffort || !supportedLevels.includes(reasoningEffort)) {
                reasoningEffort = supportedLevels[0];
            }
        }
        const body = {
            model: options.model,
            input,
            ...(instructions && { instructions }),
            ...(options.max_tokens && { max_output_tokens: Math.max(options.max_tokens, 16_000) }),
            ...(reasoningEffort && { reasoning: { effort: reasoningEffort } }),
            ...(tools?.length && { tools, tool_choice: "auto" }),
            store: false,
        };
        const response = await fetch("https://api.x.ai/v1/responses", {
            method: "POST",
            headers: {
                "Content-Type": "application/json",
                Authorization: `Bearer ${apiKey}`,
            },
            body: JSON.stringify(body),
        });
        if (!response.ok) {
            const errorText = await response.text();
            throw new Error(`XAI Responses API error: ${response.status} ${errorText}`);
        }
        const data = await response.json();
        const usage = data.usage
            ? {
                prompt_tokens: data.usage.input_tokens,
                completion_tokens: data.usage.output_tokens,
                total_tokens: data.usage.input_tokens + data.usage.output_tokens,
                prompt_tokens_details: {
                    cached_tokens: data.usage.input_tokens_details?.cached_tokens ?? 0,
                },
            }
            : undefined;
        const usdCost = usage ? this.calculateCost(options.model, usage) : undefined;
        let textContent = null;
        const toolCalls = [];
        for (const item of data.output ?? []) {
            if (item.type === "message") {
                for (const part of item.content ?? []) {
                    if (part.type === "output_text") {
                        textContent = (textContent ?? "") + part.text;
                    }
                }
            }
            else if (item.type === "function_call") {
                toolCalls.push({
                    id: item.call_id,
                    type: "function",
                    function: {
                        name: item.name,
                        arguments: item.arguments,
                    },
                });
            }
        }
        return {
            choices: [
                {
                    message: {
                        role: "assistant",
                        content: textContent,
                        ...(toolCalls.length > 0 && { tool_calls: toolCalls }),
                    },
                },
            ],
            model: options.model,
            usage,
            usd_cost: usdCost,
        };
    }
    pricesPerMillion() {
        return pricing_1.XaiTextPricing;
    }
    calculateCost(model, usage) {
        if (!usage) {
            return undefined;
        }
        const pricing = this.pricesPerMillion()[model];
        if (!pricing) {
            return undefined;
        }
        const inputTokens = usage.prompt_tokens || 0;
        const inputCost = (inputTokens * pricing.input) / 1e6;
        const outputTokens = usage.completion_tokens || 0;
        const outputCost = (outputTokens * pricing.output) / 1e6;
        const cacheToken = usage.prompt_tokens_details?.cached_tokens || 0;
        const cacheCost = (cacheToken * (pricing.cache_hit || 0)) / 1e6;
        const total = inputCost + outputCost + cacheCost;
        return total;
    }
    async getModels(modality) {
        if (modality) {
            const map = {
                completion: Object.values(types_1.Models.xai),
                image: types_1.XaiImageModels,
                video: types_1.XaiVideoModels,
            };
            return (map[modality] ?? []).map((id) => ({ id }));
        }
        return Object.values(types_1.Models.xai).map((id) => ({ id }));
    }
    async createEmbedding(options) {
        throw new Error("XAI provider does not support embeddings");
    }
    async createAudioTranscription(options) {
        throw new Error("Audio transcription is not supported by the XAI provider. Use OpenAI client with Whisper model instead.");
    }
    async createAudioGeneration(options) {
        throw new Error("Audio generation is not supported by the XAI provider. Use OpenAI client with TTS model instead.");
    }
    async createImageGeneration(options) {
        const response = await this.client.images.generate({
            model: options.model || "grok-imagine-image",
            prompt: options.prompt,
            n: options.n,
            size: options.size,
            quality: options.quality,
            style: options.style,
            response_format: options.response_format,
            user: options.user,
        });
        const imageModel = options.model || "grok-imagine-image";
        const costPerImage = pricing_1.XaiImagePricing[imageModel]?.image_generation || 0.02;
        const usdCost = (options.n || 1) * costPerImage;
        return {
            ...response,
            created: response.created || Math.floor(Date.now() / 1000),
            usd_cost: usdCost,
        };
    }
    async createVideoGeneration(options) {
        const model = options.model || "grok-imagine-video";
        const startPayload = {
            model,
            prompt: options.prompt,
        };
        if (options.duration !== undefined) {
            startPayload.duration = options.duration;
        }
        if (options.aspect_ratio) {
            startPayload.aspect_ratio = options.aspect_ratio;
        }
        if (options.resolution) {
            startPayload.resolution = options.resolution;
        }
        if (options.image_url) {
            startPayload.image_url = options.image_url;
        }
        if (options.video_url) {
            startPayload.video_url = options.video_url;
        }
        const startResponse = await fetch("https://api.x.ai/v1/videos/generations", {
            method: "POST",
            headers: {
                "Content-Type": "application/json",
                Authorization: `Bearer ${this.apiKey}`,
            },
            body: JSON.stringify(startPayload),
        });
        if (!startResponse.ok) {
            const errorText = await startResponse.text();
            throw new Error(`XAI video generation start failed: ${startResponse.status} ${errorText}`);
        }
        const startData = await startResponse.json();
        const requestId = startData.request_id;
        if (!requestId) {
            throw new Error("No request_id returned from XAI video generation start");
        }
        const duration = options.duration || 5;
        const pricePerSecond = pricing_1.XaiVideoPricing[model]?.video_generation || 0.07;
        const usdCost = duration * pricePerSecond;
        return {
            created: Math.floor(Date.now() / 1000),
            data: [],
            jobId: requestId,
            usd_cost: usdCost,
        };
    }
    async getVideoStatus(options) {
        const statusResponse = await fetch(`https://api.x.ai/v1/videos/${options.jobId}`, {
            method: "GET",
            headers: {
                Authorization: `Bearer ${this.apiKey}`,
            },
        });
        if (!statusResponse.ok) {
            const errorText = await statusResponse.text();
            throw new Error(`XAI video status check failed: ${statusResponse.status} ${errorText}`);
        }
        const statusData = await statusResponse.json();
        let mappedStatus;
        switch (statusData.status) {
            case "pending":
                mappedStatus = "queued";
                break;
            case "processing":
                mappedStatus = "in_progress";
                break;
            case "succeeded":
                mappedStatus = "completed";
                break;
            case "failed":
                mappedStatus = "failed";
                break;
            case "expired":
                mappedStatus = "expired";
                break;
            default:
                mappedStatus = "queued";
        }
        const response = {
            jobId: options.jobId,
            status: mappedStatus,
        };
        if (mappedStatus === "completed" && statusData.video?.url) {
            response.data = [
                {
                    url: statusData.video.url,
                },
            ];
        }
        return response;
    }
    async downloadVideo(options) {
        const url = options.uri || options.fileId;
        const response = await fetch(url);
        if (!response.ok) {
            throw new Error(`Failed to download video from XAI URL: ${response.status} ${response.statusText}`);
        }
        const arrayBuffer = await response.arrayBuffer();
        return {
            data: Buffer.from(arrayBuffer),
            mimeType: "video/mp4",
        };
    }
    async uploadFile(options) {
        const apiKey = this.apiKey || process.env.XAI_API_KEY;
        if (!apiKey) {
            throw new Error("XAI API key not set");
        }
        const formData = new FormData();
        formData.append("purpose", "assistants");
        const blob = new Blob([new Uint8Array(options.data)], { type: options.mimeType || "application/octet-stream" });
        formData.append("file", blob, options.fileName || "upload");
        const response = await fetch("https://api.x.ai/v1/files", {
            method: "POST",
            headers: { Authorization: `Bearer ${apiKey}` },
            body: formData,
        });
        if (!response.ok) {
            const errorText = await response.text();
            throw new Error(`XAI uploadFile failed: ${response.status} ${errorText}`);
        }
        const data = await response.json();
        return { fileId: data.id, uri: data.uri };
    }
    async downloadFile(options) {
        const apiKey = this.apiKey || process.env.XAI_API_KEY;
        if (!apiKey) {
            throw new Error("XAI API key not set");
        }
        const fileId = options.fileId;
        if (!fileId) {
            throw new Error("downloadFile requires fileId");
        }
        const response = await fetch(`https://api.x.ai/v1/files/${fileId}/content`, {
            method: "GET",
            headers: { Authorization: `Bearer ${apiKey}` },
        });
        if (!response.ok) {
            const errorText = await response.text();
            throw new Error(`XAI downloadFile failed: ${response.status} ${errorText}`);
        }
        const mimeType = response.headers.get("content-type") || undefined;
        const data = Buffer.from(await response.arrayBuffer());
        return {
            data,
            mimeType,
        };
    }
    getContextLimit(model) {
        const contextLimit = contextLimits_1.ContextLimits[model];
        if (contextLimit === undefined)
            return undefined;
        const pricing = pricing_1.XaiTextPricing[model];
        const threshold = pricing && "input_gt_200k" in pricing ? 200_000 : contextLimit;
        return { contextLimit, threshold };
    }
}
exports.GenericXAIClient = GenericXAIClient;
//# sourceMappingURL=xai.js.map