apix/lib/providers/grok-client.ts

/**
 * Grok/xAI Client for Image Generation
 * 
 * Supports two authentication methods:
 * 1. Official API Key from console.x.ai (recommended)
 * 2. Cookie-based auth from logged-in grok.com session
 * 
 * Image Model: FLUX.1 by Black Forest Labs
 */

// Official xAI API endpoint
const XAI_API_BASE = "https://api.x.ai/v1";

// Grok web interface endpoint (for cookie-based auth)
const GROK_WEB_BASE = "https://grok.com";

interface GrokGenerateOptions {
    prompt: string;
    apiKey?: string;
    cookies?: string;
    numImages?: number;
}

interface GrokImageResult {
    url: string;
    data?: string; // base64
    prompt: string;
    model: string;
}

export class GrokClient {
    private apiKey?: string;
    private cookies?: string;

    constructor(options: { apiKey?: string; cookies?: string }) {
        this.apiKey = options.apiKey;
        this.cookies = this.normalizeCookies(options.cookies);
    }

    /**
     * Normalize cookies from string or JSON format
     * Handles cases where user pastes JSON array from extension/devtools
     */
    private normalizeCookies(cookies?: string): string | undefined {
        if (!cookies) return undefined;

        try {
            // Check if it looks like JSON
            if (cookies.trim().startsWith('[')) {
                const parsed = JSON.parse(cookies);
                if (Array.isArray(parsed)) {
                    return parsed
                        .map((c: any) => `${c.name}=${c.value}`)
                        .join('; ');
                }
            }
        } catch (e) {
            // Not JSON, assume string
        }

        return cookies;
    }

    /**
     * Generate images using Grok/xAI
     * Prefers official API if apiKey is provided, falls back to cookie-based
     */
    async generate(prompt: string, numImages: number = 1): Promise<GrokImageResult[]> {
        if (this.apiKey) {
            return this.generateWithAPI(prompt, numImages);
        } else if (this.cookies) {
            return this.generateWithCookies(prompt, numImages);
        } else {
            throw new Error("Grok: No API key or cookies provided. Configure in Settings.");
        }
    }

    /**
     * Generate using official xAI API (recommended)
     * Requires API key from console.x.ai
     */
    private async generateWithAPI(prompt: string, numImages: number): Promise<GrokImageResult[]> {
        console.log(`[Grok API] Generating ${numImages} image(s) for: "${prompt.substring(0, 50)}..."`);

        const response = await fetch(`${XAI_API_BASE}/images/generations`, {
            method: "POST",
            headers: {
                "Content-Type": "application/json",
                "Authorization": `Bearer ${this.apiKey}`
            },
            body: JSON.stringify({
                model: "grok-2-image",
                prompt: prompt,
                n: numImages,
                response_format: "url" // or "b64_json"
            })
        });

        if (!response.ok) {
            const errorText = await response.text();
            console.error("[Grok API] Error:", response.status, errorText);
            throw new Error(`Grok API Error: ${response.status} - ${errorText.substring(0, 200)}`);
        }

        const data = await response.json();
        console.log("[Grok API] Response:", JSON.stringify(data, null, 2));

        // Parse response - xAI uses OpenAI-compatible format
        const images: GrokImageResult[] = (data.data || []).map((img: any) => ({
            url: img.url || (img.b64_json ? `data:image/png;base64,${img.b64_json}` : ''),
            data: img.b64_json,
            prompt: prompt,
            model: "grok-2-image"
        }));

        if (images.length === 0) {
            throw new Error("Grok API returned no images");
        }

        return images;
    }

    /**
     * Generate using Grok web interface (cookie-based)
     * Requires cookies from logged-in grok.com session
     */
    private async generateWithCookies(prompt: string, numImages: number): Promise<GrokImageResult[]> {
        console.log(`[Grok Web] Generating image for: "${prompt.substring(0, 50)}..."`);

        // The Grok web interface uses a chat-based API
        // We need to send a message asking for image generation
        const imagePrompt = `Generate an image: ${prompt}`;

        const response = await fetch(`${GROK_WEB_BASE}/rest/app-chat/conversations/new`, {
            method: "POST",
            headers: {
                "Content-Type": "application/json",
                "Cookie": this.cookies!,
                "Origin": GROK_WEB_BASE,
                "Referer": `${GROK_WEB_BASE}/`,
                "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36",
                "Sec-Fetch-Site": "same-origin",
                "Sec-Fetch-Mode": "cors",
                "Sec-Fetch-Dest": "empty"
            },
            body: JSON.stringify({
                temporary: false,
                modelName: "grok-3",
                message: imagePrompt,
                fileAttachments: [],
                imageAttachments: [],
                disableSearch: false,
                enableImageGeneration: true,
                returnImageBytes: false,
                returnRawGrokInXaiRequest: false,
                sendFinalMetadata: true,
                customInstructions: "",
                deepsearchPreset: "",
                isReasoning: false
            })
        });

        if (!response.ok) {
            const errorText = await response.text();
            console.error("[Grok Web] Error:", response.status, errorText);
            throw new Error(`Grok Web Error: ${response.status} - ${errorText.substring(0, 200)}`);
        }

        // Parse streaming response to find image URLs
        const text = await response.text();
        console.log("[Grok Web] Response length:", text.length);

        // Look for generated image URLs in the response
        const imageUrls = this.extractImageUrls(text);

        if (imageUrls.length === 0) {
            console.warn("[Grok Web] No image URLs found in response. Response preview:", text.substring(0, 500));
            throw new Error("Grok did not generate any images. Try a different prompt or check your cookies.");
        }

        return imageUrls.map(url => ({
            url,
            prompt,
            model: "grok-3"
        }));
    }

    /**
     * Extract image URLs from Grok's streaming response
     */
    private extractImageUrls(responseText: string): string[] {
        const urls: string[] = [];

        // Try to parse as JSON lines (NDJSON format)
        const lines = responseText.split('\n').filter(line => line.trim());

        for (const line of lines) {
            try {
                const data = JSON.parse(line);

                // Check for generatedImageUrls field
                if (data.generatedImageUrls && Array.isArray(data.generatedImageUrls)) {
                    urls.push(...data.generatedImageUrls);
                }

                // Check for imageUrls in result
                if (data.result?.imageUrls) {
                    urls.push(...data.result.imageUrls);
                }

                // Check for media attachments
                if (data.attachments) {
                    for (const attachment of data.attachments) {
                        if (attachment.type === 'image' && attachment.url) {
                            urls.push(attachment.url);
                        }
                    }
                }
            } catch {
                // Not JSON, try regex extraction
            }
        }

        // Fallback: regex for image URLs
        if (urls.length === 0) {
            const urlRegex = /https:\/\/[^"\s]+\.(png|jpg|jpeg|webp)/gi;
            const matches = responseText.match(urlRegex);
            if (matches) {
                urls.push(...matches);
            }
        }

        // Deduplicate
        return [...new Set(urls)];
    }

    /**
     * Download image from URL and convert to base64
     */
    async downloadAsBase64(url: string): Promise<string> {
        const response = await fetch(url);
        const buffer = await response.arrayBuffer();
        const base64 = Buffer.from(buffer).toString('base64');
        return base64;
    }
}