/**
 * Meta AI Client for Image Generation
 * 
 * Uses the Meta AI web interface via GraphQL
 * Requires cookies from a logged-in meta.ai session
 * 
 * Image Model: Imagine (Emu)
 * 
 * Based on: https://github.com/Strvm/meta-ai-api
 */

const META_AI_BASE = "https://www.meta.ai";
const GRAPHQL_ENDPOINT = `${META_AI_BASE}/api/graphql/`;

interface MetaAIOptions {
    cookies: string;
}

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

interface MetaSession {
    lsd?: string;
    fb_dtsg?: string;
    accessToken?: string;
    externalConversationId?: string;
}

// Aspect ratio types for Meta AI
export type AspectRatio = 'portrait' | 'landscape' | 'square';
const ORIENTATION_MAP: Record<AspectRatio, string> = {
    'portrait': 'VERTICAL',    // 9:16
    'landscape': 'HORIZONTAL', // 16:9
    'square': 'SQUARE'         // 1:1
};

export class MetaAIClient {
    private cookies: string;
    private session: MetaSession = {};
    private useFreeWrapper: boolean = true;
    private freeWrapperUrl: string = 'http://localhost:8000';

    constructor(options: MetaAIOptions & { useFreeWrapper?: boolean; freeWrapperUrl?: string }) {
        this.cookies = this.normalizeCookies(options.cookies);
        this.useFreeWrapper = options.useFreeWrapper !== undefined ? options.useFreeWrapper : true;
        this.freeWrapperUrl = options.freeWrapperUrl || 'http://localhost:8000';
        console.log("[Meta AI] Cookie string length:", this.cookies.length);
        if (this.cookies) {
            this.parseSessionFromCookies();
        }
    }

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

        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;
    }

    /**
     * Parse session tokens from cookies
     */
    private parseSessionFromCookies(): void {
        // Extract lsd token if present in cookies
        const lsdMatch = this.cookies.match(/lsd=([^;]+)/);
        if (lsdMatch) {
            this.session.lsd = lsdMatch[1];
        }

        // Extract fb_dtsg if present
        const dtsgMatch = this.cookies.match(/fb_dtsg=([^;]+)/);
        if (dtsgMatch) {
            this.session.fb_dtsg = dtsgMatch[1];
        }
    }

    /**
     * Get the initialized session tokens (call initSession first if not done)
     */
    async getSession(): Promise<MetaSession & { externalConversationId?: string }> {
        if (!this.useFreeWrapper && !this.session.lsd && !this.session.fb_dtsg) {
            await this.initSession();
        }
        return this.session;
    }

    /**
     * Get the normalized cookie string
     */
    getCookies(): string {
        return this.cookies;
    }

    /**
     * Generate images using Meta AI's Imagine model
     */
    async generate(prompt: string, numImages: number = 4, aspectRatio: AspectRatio = 'portrait'): Promise<MetaImageResult[]> {
        console.log(`[Meta AI] Generating images for: "${prompt.substring(0, 50)}..." (${aspectRatio})`);

        if (this.useFreeWrapper) {
            return this.generateWithFreeWrapper(prompt, numImages);
        }

        // First, get the access token and session info if not already fetched
        if (!this.session.accessToken) {
            await this.initSession();
        }

        // Use "Imagine" prefix for image generation
        const imagePrompt = prompt.toLowerCase().startsWith('imagine')
            ? prompt
            : `Imagine ${prompt}`;

        // Send the prompt via GraphQL with aspect ratio
        const response = await this.sendPrompt(imagePrompt, aspectRatio);

        // Extract image URLs from response
        const images = this.extractImages(response, prompt);

        if (images.length === 0) {
            // Meta AI might return the prompt response without images immediately
            // We may need to poll for the result
            console.log("[Meta AI] No images in initial response, polling...");
            return this.pollForImages(response, prompt);
        }

        return images;
    }

    /**
     * Generate using free API wrapper (mir-ashiq/metaai-api)
     * Connects to local docker service
     */
    /**
     * Generate using free API wrapper (mir-ashiq/metaai-api)
     * Connects to local docker service
     */
    private async generateWithFreeWrapper(prompt: string, numImages: number): Promise<MetaImageResult[]> {
        console.log(`[Meta Wrapper] Generating image for: "${prompt.substring(0, 50)}..." via ${this.freeWrapperUrl}`);

        const cookieDict = this.parseCookiesToDict(this.cookies);

        const response = await fetch(`${this.freeWrapperUrl}/chat`, {
            method: "POST",
            headers: {
                "Content-Type": "application/json"
            },
            body: JSON.stringify({
                message: `Imagine ${prompt}`,
                stream: false,
                cookies: cookieDict
            })
        });

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

        const data = await response.json();
        console.log("[Meta Wrapper] Response:", JSON.stringify(data, null, 2).substring(0, 500));

        // Check for images in response
        const images: MetaImageResult[] = [];

        // mir-ashiq/metaai-api returns { media: [{ url: "...", ... }] }
        if (data.media && Array.isArray(data.media)) {
            data.media.forEach((m: any) => {
                if (m.url) images.push({ url: m.url, prompt, model: "meta-wrapper" });
            });
        }

        // Fallback checks
        if (images.length === 0 && data.images && Array.isArray(data.images)) {
            data.images.forEach((url: string) => {
                images.push({ url, prompt, model: "meta-wrapper" });
            });
        }

        if (images.length === 0 && data.sources && Array.isArray(data.sources)) {
            data.sources.forEach((s: any) => {
                if (s.url && (s.url.includes('.jpg') || s.url.includes('.png') || s.url.includes('.webp'))) {
                    images.push({ url: s.url, prompt, model: "meta-wrapper-source" });
                }
            });
        }

        if (images.length === 0) {
            console.warn("[Meta Wrapper] No images found via /chat endpoint", data);
            throw new Error("Meta Wrapper returned no images. Please check if the prompt triggered image generation.");
        }

        return images;
    }

    private parseCookiesToDict(cookieStr: string): Record<string, string> {
        const dict: Record<string, string> = {};
        if (!cookieStr) return dict;

        // Handle basic key=value; format
        cookieStr.split(';').forEach(pair => {
            const [key, ...rest] = pair.trim().split('=');
            if (key && rest.length > 0) {
                dict[key] = rest.join('=');
            }
        });
        return dict;
    }

    private async generateWithFreeWrapperFallback(prompt: string): Promise<MetaImageResult[]> {
        // Fallback logic not needed if /chat works
        throw new Error("Meta Wrapper endpoint /chat failed.");
    }

    /**
     * Initialize session - get access token from meta.ai page
     */
    private async initSession(): Promise<void> {
        console.log("[Meta AI] Initializing session...");

        const response = await fetch(META_AI_BASE, {
            headers: {
                "Cookie": this.cookies,
                "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": "navigate",
                "Sec-Fetch-Dest": "document",
                "Accept-Language": "en-US,en;q=0.9"
            }
        });

        const html = await response.text();

        // Extract access token from page HTML
        // Pattern 1: Simple JSON key
        let tokenMatch = html.match(/"accessToken":"([^"]+)"/);

        // Pattern 2: Inside a config object or different spacing
        if (!tokenMatch) {
            tokenMatch = html.match(/accessToken["']\s*:\s*["']([^"']+)["']/);
        }

        // Pattern 3: LSD token backup (sometimes needed)
        const lsdMatch = html.match(/"LSD",\[\],{"token":"([^"]+)"/) ||
            html.match(/"lsd":"([^"]+)"/) ||
            html.match(/name="lsd" value="([^"]+)"/);
        if (lsdMatch) {
            this.session.lsd = lsdMatch[1];
        }

        // Pattern 4: DTSG token (critical for some requests)
        const dtsgMatch = html.match(/"DTSGInitialData".*?"token":"([^"]+)"/) ||
            html.match(/"token":"([^"]+)"/); // Less specific fallback

        if (tokenMatch) {
            this.session.accessToken = tokenMatch[1];
            console.log("[Meta AI] Got access token");
        } else if (html.includes('login_form') || html.includes('login_page')) {
            throw new Error("Meta AI: Cookies expired or invalid (Login page detected). Please update cookies.");
        } else {
            console.warn("[Meta AI] Warning: Failed to extract access token. functionality may be limited.");
            console.log("HTML Preview:", html.substring(0, 200));
            // Don't throw here, try to proceed with just cookies/LSD
        }

        if (dtsgMatch) {
            this.session.fb_dtsg = dtsgMatch[1];
        }

        // Enhanced logging for debugging
        console.log("[Meta AI] Session tokens extracted:", {
            hasAccessToken: !!this.session.accessToken,
            hasLsd: !!this.session.lsd,
            hasDtsg: !!this.session.fb_dtsg
        });

        if (!this.session.accessToken && !this.session.lsd) {
            console.warn("[Meta AI] CRITICAL: No authentication tokens found. Check if cookies are valid.");
        }
    }

    /**
     * Send prompt via GraphQL mutation
     */
    /**
     * Send prompt via GraphQL mutation (Abra - for Image Generation)
     * Using EXACT variable structure from Python Strvm/meta-ai-api library
     */
    private async sendPrompt(prompt: string, aspectRatio: AspectRatio = 'portrait'): Promise<any> {
        // Generate external conversation ID (UUID) and offline threading ID (Snowflake-like)
        const externalConversationId = crypto.randomUUID();
        const timestamp = Date.now();
        const randomPart = Math.floor(Math.random() * 4194304); // 22 bits
        const offlineThreadingId = ((BigInt(timestamp) << 22n) | BigInt(randomPart)).toString();

        // Store for polling
        this.session.externalConversationId = externalConversationId;

        // Map aspect ratio to Meta AI orientation
        const orientation = ORIENTATION_MAP[aspectRatio];

        const variables = {
            message: {
                sensitive_string_value: prompt  // Python uses sensitive_string_value, not text!
            },
            externalConversationId: externalConversationId,
            offlineThreadingId: offlineThreadingId,
            suggestedPromptIndex: null,
            flashVideoRecapInput: { images: [] },
            flashPreviewInput: null,
            promptPrefix: null,
            entrypoint: "ABRA__CHAT__TEXT",
            icebreaker_type: "TEXT",
            imagineClientOptions: { orientation: orientation },  // Aspect ratio control
            __relay_internal__pv__AbraDebugDevOnlyrelayprovider: false,
            __relay_internal__pv__WebPixelRatiorelayprovider: 1
        };

        console.log("[Meta AI] Sending Variables:", JSON.stringify(variables, null, 2));

        const body = new URLSearchParams({
            fb_api_caller_class: "RelayModern",
            fb_api_req_friendly_name: "useAbraSendMessageMutation",
            variables: JSON.stringify(variables),
            server_timestamps: "true",
            doc_id: "7783822248314888", // Abra Mutation ID
            ...(this.session.lsd && { lsd: this.session.lsd }),
            ...(this.session.fb_dtsg && { fb_dtsg: this.session.fb_dtsg })
        });

        const response = await fetch(GRAPHQL_ENDPOINT, {
            method: "POST",
            headers: {
                "Content-Type": "application/x-www-form-urlencoded",
                "Cookie": this.cookies,
                "Origin": META_AI_BASE,
                "Referer": `${META_AI_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",
                "Accept-Language": "en-US,en;q=0.9",
                ...(this.session.accessToken && { "Authorization": `OAuth ${this.session.accessToken}` })
            },
            body: body.toString()
        });

        const rawText = await response.text();
        console.log("[Meta AI] Response received, parsing streaming data...");

        if (!response.ok) {
            throw new Error(`Meta AI Error: ${response.status} - ${rawText.substring(0, 500)}`);
        }

        // Meta AI returns streaming response (multiple JSON lines)
        // We need to find the final response where streaming_state === "OVERALL_DONE"
        let lastValidResponse: any = null;
        const lines = rawText.split('\n');

        for (const line of lines) {
            if (!line.trim()) continue;
            try {
                const parsed = JSON.parse(line);

                // Check for streaming state in both direct and nested paths
                const streamingState =
                    parsed?.data?.xfb_abra_send_message?.bot_response_message?.streaming_state ||
                    parsed?.data?.node?.bot_response_message?.streaming_state;

                if (streamingState === "OVERALL_DONE") {
                    console.log("[Meta AI] Found OVERALL_DONE response");
                    lastValidResponse = parsed;
                    break;
                }

                // Keep track of any valid response with imagine_card
                const imagineCard =
                    parsed?.data?.xfb_abra_send_message?.bot_response_message?.imagine_card ||
                    parsed?.data?.node?.bot_response_message?.imagine_card;
                if (imagineCard?.session?.media_sets) {
                    lastValidResponse = parsed;
                }
            } catch (e) {
                // Skip non-JSON lines
                continue;
            }
        }

        if (!lastValidResponse) {
            if (rawText.includes("login_form") || rawText.includes("facebook.com/login")) {
                throw new Error("Meta AI: Session expired. Please refresh your cookies.");
            }
            throw new Error("Meta AI: No valid response found in streaming data");
        }

        console.log("[Meta AI] Successfully parsed streaming response");
        return lastValidResponse;
    }

    /**
     * Extract image URLs from Meta AI response
     */
    private extractImages(response: any, originalPrompt: string): MetaImageResult[] {
        const images: MetaImageResult[] = [];

        // Navigate through the response structure
        // Abra streaming: xfb_abra_send_message.bot_response_message
        // Abra direct: node.bot_response_message  
        // Kadabra: useKadabraSendMessageMutation.node.bot_response_message
        const messageData = response?.data?.xfb_abra_send_message?.bot_response_message ||
            response?.data?.useKadabraSendMessageMutation?.node?.bot_response_message ||
            response?.data?.node?.bot_response_message ||
            response?.data?.xabraAIPreviewMessageSendMutation?.message;

        // For Polling/KadabraPromptRootQuery which also contains imagine_cards
        const pollMessages = response?.data?.kadabra_prompt?.messages?.edges || [];
        if (pollMessages.length > 0) {
            for (const edge of pollMessages) {
                const nodeImages = this.extractImagesFromMessage(edge?.node, originalPrompt);
                images.push(...nodeImages);
            }
        }

        if (messageData) {
            images.push(...this.extractImagesFromMessage(messageData, originalPrompt));
        }

        // --- STRENGTHENED FALLBACK ---
        // If still no images, but we have data, do a recursive search for any CDN URLs
        if (images.length === 0 && response?.data) {
            console.log("[Meta AI] Structured extraction failed, attempting recursive search...");
            const foundUrls = this.recursiveSearchForImages(response.data);
            for (const url of foundUrls) {
                images.push({
                    url: url,
                    prompt: originalPrompt,
                    model: "meta"
                });
            }
        }

        if (images.length === 0) {
            console.log("[Meta AI] Extraction failed. Response keys:", Object.keys(response || {}));
            if (response?.data) console.log("[Meta AI] Data keys:", Object.keys(response.data));
        }

        return images;
    }

    /**
     * Recursive search for image-like URLs in the JSON tree
     */
    private recursiveSearchForImages(obj: any, found: Set<string> = new Set()): string[] {
        if (!obj || typeof obj !== 'object') return [];

        for (const key in obj) {
            const val = obj[key];
            if (typeof val === 'string') {
                if ((val.includes('fbcdn.net') || val.includes('meta.ai')) &&
                    (val.includes('.jpg') || val.includes('.png') || val.includes('.webp') || val.includes('image_uri=') || val.includes('/imagine/'))) {
                    found.add(val);
                }
            } else if (typeof val === 'object') {
                this.recursiveSearchForImages(val, found);
            }
        }
        return Array.from(found);
    }

    /**
     * Helper to extract images from a single message node
     */
    private extractImagesFromMessage(messageData: any, originalPrompt: string): MetaImageResult[] {
        const images: MetaImageResult[] = [];
        if (!messageData) return images;

        // Check for imagine_card (image generation response)
        const imagineCard = messageData?.imagine_card;
        if (imagineCard?.session?.media_sets) {
            for (const mediaSet of imagineCard.session.media_sets) {
                if (mediaSet?.imagine_media) {
                    for (const media of mediaSet.imagine_media) {
                        const url = media?.uri || media?.image_uri || media?.image?.uri;
                        if (url) {
                            images.push({
                                url: url,
                                prompt: originalPrompt,
                                model: "meta"
                            });
                        }
                    }
                }
            }
        }

        // Check for attachments (alternative path)
        const attachments = messageData?.attachments;
        if (attachments) {
            for (const attachment of attachments) {
                const media = attachment?.media;
                const url = media?.image_uri || media?.uri || media?.image?.uri;
                if (url) {
                    images.push({
                        url: url,
                        prompt: originalPrompt,
                        model: "meta"
                    });
                }
            }
        }

        return images;
    }

    /**
     * Poll for image generation completion
     */
    private async pollForImages(initialResponse: any, prompt: string): Promise<MetaImageResult[]> {
        // Kadabra uses external_conversation_id for polling
        const conversationId = initialResponse?.data?.useKadabraSendMessageMutation?.node?.external_conversation_id ||
            initialResponse?.data?.node?.external_conversation_id;

        if (!conversationId) {
            console.warn("[Meta AI] No conversation ID found for polling");
            return [];
        }

        const maxAttempts = 30;
        const pollInterval = 2000;

        for (let attempt = 0; attempt < maxAttempts; attempt++) {
            console.log(`[Meta AI] Polling attempt ${attempt + 1}/${maxAttempts}...`);

            await new Promise(resolve => setTimeout(resolve, pollInterval));

            const variables = {
                external_conversation_id: conversationId
            };

            const body = new URLSearchParams({
                fb_api_caller_class: "RelayModern",
                fb_api_req_friendly_name: "KadabraPromptRootQuery",
                variables: JSON.stringify(variables),
                doc_id: "25290569913909283", // KadabraPromptRootQuery ID
                ...(this.session.lsd && { lsd: this.session.lsd }),
                ...(this.session.fb_dtsg && { fb_dtsg: this.session.fb_dtsg })
            });

            try {
                const response = await fetch(GRAPHQL_ENDPOINT, {
                    method: "POST",
                    headers: {
                        "Content-Type": "application/x-www-form-urlencoded",
                        "Cookie": this.cookies,
                        "Origin": META_AI_BASE,
                        ...(this.session.accessToken && { "Authorization": `OAuth ${this.session.accessToken}` })
                    },
                    body: body.toString()
                });

                const data = await response.json();
                const images = this.extractImages(data, prompt);

                if (images.length > 0) {
                    console.log(`[Meta AI] Got ${images.length} image(s) after polling!`);
                    return images;
                }

                // Check for failure status
                const status = data?.data?.kadabra_prompt?.status;
                if (status === "FAILED" || status === "ERROR") {
                    console.error("[Meta AI] Generation failed during polling");
                    break;
                }
            } catch (e: any) {
                console.error("[Meta AI] Poll error:", e.message);
            }
        }

        return [];
    }



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