server.js

import express from 'express';
// Import the 'fal' object directly for configuration within the retry loop
import { fal } from '@fal-ai/client';

// --- Key Management Setup ---
// Read comma-separated keys from the SINGLE environment variable FAL_KEY
const FAL_KEY_STRING = process.env.FAL_KEY;
// Read the custom API Key for proxy authentication
const API_KEY = process.env.API_KEY;

// --- Initial Environment Variable Checks ---
if (!FAL_KEY_STRING) {
    console.error("ERROR: FAL_KEY environment variable is not set.");
    console.error("Ensure FAL_KEY contains a comma-separated list of your Fal AI keys.");
    process.exit(1); // Exit if no Fal keys are provided
}

// Parse the comma-separated keys from FAL_KEY_STRING
const falKeys = FAL_KEY_STRING.split(',')
    .map(key => key.trim()) // Remove leading/trailing whitespace
    .filter(key => key.length > 0); // Remove any empty strings resulting from extra commas

if (falKeys.length === 0) {
    console.error("ERROR: No valid Fal keys found in the FAL_KEY environment variable after parsing.");
    console.error("Ensure FAL_KEY is a comma-separated list, e.g., 'key1,key2,key3'.");
    process.exit(1); // Exit if parsing results in zero valid keys
}

if (!API_KEY) {
    console.error("ERROR: API_KEY environment variable is not set.");
    process.exit(1); // Exit if the proxy auth key is missing
}
// --- End Initial Checks ---


let currentKeyIndex = 0;
// Keep track of keys that failed persistently during runtime
const invalidKeys = new Set();

console.log(`Loaded ${falKeys.length} Fal AI Key(s) from the FAL_KEY environment variable.`);

/**
 * Gets the next available valid Fal AI key in a round-robin fashion.
 * Skips keys that have been marked as invalid.
 * @returns {object|null} An object containing the key and its original index { key, index }, or null if no valid keys remain.
 */
function getNextValidKey() {
    // Check if all keys have been marked as invalid
    if (invalidKeys.size >= falKeys.length) {
        console.error("All Fal AI keys are marked as invalid.");
        return null; // No valid keys left
    }

    const initialIndex = currentKeyIndex;
    let attempts = 0; // Prevent infinite loops in edge cases
    while (attempts < falKeys.length) {
        const keyIndex = currentKeyIndex % falKeys.length;
        const key = falKeys[keyIndex];

        // Move to the next index for the *next* call, regardless of validity
        currentKeyIndex = (keyIndex + 1) % falKeys.length;

        // Check if the current key is NOT in the invalid set
        if (!invalidKeys.has(key)) {
            // Found a valid key
            console.log(`Using Fal Key index: ${keyIndex} (from FAL_KEY list)`);
            return { key, index: keyIndex }; // Return the key and its original index
        } else {
            console.log(`Skipping invalid Fal Key index: ${keyIndex}`);
        }

        attempts++;
        // Safety check: If we've looped back to the start after trying, break.
        if (currentKeyIndex === initialIndex && attempts > 0) {
            console.warn("Looped through all keys, potentially all are invalid.");
            break;
        }
    }

    // If we exit the loop, it means no valid key was found
    console.error("Could not find a valid Fal AI key after checking all potentially available keys.");
    return null;
}

/**
 * Checks if an error object likely indicates an issue with the Fal AI API key.
 * This is heuristic-based and may need refinement based on observed Fal errors.
 * @param {Error|object|null} error - The error object caught.
 * @returns {boolean} True if the error seems key-related, false otherwise.
 */
function isKeyRelatedError(error) {
    if (!error) return false; // Handle null/undefined errors

    const message = error.message?.toLowerCase() || '';
    // Check common HTTP status properties
    const status = error.status || error.statusCode;

    // Check for specific HTTP status codes (401: Unauthorized, 403: Forbidden, 429: Too Many Requests/Quota)
    if (status === 401 || status === 403 || status === 429) {
        console.warn(`Detected potential key-related error (HTTP Status: ${status}).`);
        return true;
    }

    // Check for common error message patterns (case-insensitive)
    const keyErrorPatterns = [
        'invalid api key', 'authentication failed', 'permission denied',
        'quota exceeded', 'forbidden', 'unauthorized', 'rate limit',
        'credentials', 'api key missing', 'invalid credential',
        'exhausted balance', 'user is locked' // Added based on observed logs
    ];
    if (keyErrorPatterns.some(pattern => message.includes(pattern))) {
        console.warn(`Detected potential key-related error (message contains relevant pattern: "${message}")`);
        return true;
    }

    // Check the body.detail if status is 403 (as seen in Fal errors)
    if (status === 403 && error.body?.detail) {
        const detailMessage = String(error.body.detail).toLowerCase();
         if (keyErrorPatterns.some(pattern => detailMessage.includes(pattern))) {
            console.warn(`Detected potential key-related error (body.detail contains relevant pattern: "${detailMessage}")`);
            return true;
         }
    }

    // Add more specific checks based on observed Fal AI errors if needed
    // e.g., if (error.code === 'FAL_AUTH_FAILURE') return true;

    return false;
}
// --- End Key Management Setup ---

const app = express();
// Increase payload size limits if needed
app.use(express.json({ limit: '50mb' }));
app.use(express.urlencoded({ extended: true, limit: '50mb' }));

const PORT = process.env.PORT || 3000;

// API Key Authentication Middleware
const apiKeyAuth = (req, res, next) => {
    const authHeader = req.headers['authorization'];

    if (!authHeader) {
        console.warn('Unauthorized: No Authorization header provided');
        return res.status(401).json({ error: 'Unauthorized: No API Key provided' });
    }

    // Expecting "Bearer YOUR_API_KEY"
    const authParts = authHeader.split(' ');
    if (authParts.length !== 2 || authParts[0].toLowerCase() !== 'bearer') {
        console.warn('Unauthorized: Invalid Authorization header format. Expected "Bearer <key>".');
        return res.status(401).json({ error: 'Unauthorized: Invalid Authorization header format' });
    }

    const providedKey = authParts[1];
    if (providedKey !== API_KEY) {
        console.warn('Unauthorized: Invalid API Key provided.');
        return res.status(401).json({ error: 'Unauthorized: Invalid API Key' });
    }

    // Key is valid, proceed to the next middleware or route handler
    next();
};

// Apply API Key Authentication to relevant endpoints
app.use(['/v1/models', '/v1/chat/completions'], apiKeyAuth);

// === Global Limits Definition ===
const PROMPT_LIMIT = 4800; // Max length for the main 'prompt' field
const SYSTEM_PROMPT_LIMIT = 4800; // Max length for the 'system_prompt' field
// === End Limits Definition ===

// Define the list of models supported by fal-ai/any-llm (Update as needed)
const FAL_SUPPORTED_MODELS = [
    "anthropic/claude-3.7-sonnet",
    "anthropic/claude-3.5-sonnet",
    "anthropic/claude-3-5-haiku",
    "anthropic/claude-3-haiku",
    "google/gemini-pro-1.5",
    "google/gemini-flash-1.5",
    "google/gemini-flash-1.5-8b",
    "google/gemini-2.0-flash-001",
    "meta-llama/llama-3.2-1b-instruct",
    "meta-llama/llama-3.2-3b-instruct",
    "meta-llama/llama-3.1-8b-instruct",
    "meta-llama/llama-3.1-70b-instruct",
    "openai/gpt-4o-mini",
    "openai/gpt-4o",
    "deepseek/deepseek-r1",
    "meta-llama/llama-4-maverick",
    "meta-llama/llama-4-scout"
    // Add or remove models here
];

// Helper function to extract the owner/organization from a model ID string
const getOwner = (modelId) => {
    if (modelId && typeof modelId === 'string' && modelId.includes('/')) {
        return modelId.split('/')[0];
    }
    // Default owner if format is unexpected or missing
    return 'fal-ai';
}

// GET /v1/models endpoint - Returns the list of supported models
app.get('/v1/models', (req, res) => {
    console.log("Received request for GET /v1/models");
    try {
        const modelsData = FAL_SUPPORTED_MODELS.map(modelId => ({
            id: modelId,
            object: "model",
            created: Math.floor(Date.now() / 1000), // Use current timestamp
            owned_by: getOwner(modelId)
        }));
        res.json({ object: "list", data: modelsData });
        console.log("Successfully returned model list.");
    } catch (error) {
        console.error("Error processing GET /v1/models:", error);
        res.status(500).json({ error: "Failed to retrieve model list." });
    }
});


/**
 * Converts OpenAI-style messages array to Fal AI's prompt and system_prompt format.
 * Implements System prompt top-priority, separator, and recency-based history filling.
 * Includes robustness checks for input validation and internal errors.
 * @param {Array<object>} messages - Array of message objects ({ role: string, content: string|null }).
 * @returns {object} An object containing { system_prompt: string, prompt: string }.
 * @throws {Error} If input is fundamentally invalid (e.g., not an array) or an unexpected internal processing error occurs.
 */
function convertMessagesToFalPrompt(messages) {
    // --- Optional Debug Log: Uncomment to see the exact input causing issues ---
    // console.log(">>> Entering convertMessagesToFalPrompt with messages:", JSON.stringify(messages, null, 2));

    // --- Input Validation ---
    if (!Array.isArray(messages)) {
        console.error("!!! ERROR in convertMessagesToFalPrompt: Input 'messages' is not an array.");
        // Throw an error here because this is a fundamental type mismatch from the caller.
        throw new Error("Invalid input: 'messages' must be an array.");
    }
    if (messages.length === 0) {
        console.warn("Warning in convertMessagesToFalPrompt: Input 'messages' array is empty.");
        // Return empty strings if no messages, this is valid input.
        return { system_prompt: "", prompt: "" };
    }
    // --- End Input Validation ---

    // --- Main Processing Logic ---
    try { // *** Wrap core logic in try...catch to handle internal errors ***
        let fixed_system_prompt_content = "";
        const conversation_message_blocks = [];
        // console.log(`Original messages count: ${messages.length}`);

        // 1. Separate System messages, format User/Assistant messages
        for (const message of messages) {
            // ** Validate individual message structure **
            if (!message || typeof message !== 'object' || typeof message.role !== 'string') {
                console.warn(`--> Skipping invalid message object in convertMessagesToFalPrompt: ${JSON.stringify(message)}`);
                continue; // Skip this malformed message, proceed with others
            }

            // ** Safely handle content (null/undefined/non-string become empty string) **
            let content = (message.content === null || message.content === undefined) ? "" : String(message.content);

            switch (message.role) {
                case 'system':
                    fixed_system_prompt_content += `System: ${content}\n\n`;
                    break;
                case 'user':
                    conversation_message_blocks.push(`Human: ${content}\n\n`);
                    break;
                case 'assistant':
                    conversation_message_blocks.push(`Assistant: ${content}\n\n`);
                    break;
                default:
                    console.warn(`--> Unsupported role encountered in convertMessagesToFalPrompt: '${message.role}'. Skipping message.`);
                    continue; // Skip messages with unsupported roles
            }
        }

        // 2. Truncate combined system messages if they exceed the limit
        if (fixed_system_prompt_content.length > SYSTEM_PROMPT_LIMIT) {
            const originalLength = fixed_system_prompt_content.length;
            fixed_system_prompt_content = fixed_system_prompt_content.substring(0, SYSTEM_PROMPT_LIMIT);
            console.warn(`Combined system messages truncated from ${originalLength} to ${SYSTEM_PROMPT_LIMIT} characters.`);
        }
        fixed_system_prompt_content = fixed_system_prompt_content.trim();

        // 3. Calculate remaining space in system_prompt for history
        let space_occupied_by_fixed_system = 0;
        if (fixed_system_prompt_content.length > 0) {
            space_occupied_by_fixed_system = fixed_system_prompt_content.length + 4; // Approx for spacing
        }
        const remaining_system_limit = Math.max(0, SYSTEM_PROMPT_LIMIT - space_occupied_by_fixed_system);

        // 4. Fill history backwards (recency): Prioritize 'prompt', then 'system_prompt' overflow
        const prompt_history_blocks = [];
        const system_prompt_history_blocks = [];
        let current_prompt_length = 0;
        let current_system_history_length = 0;
        let promptFull = (PROMPT_LIMIT <= 0);
        let systemHistoryFull = (remaining_system_limit <= 0);

        for (let i = conversation_message_blocks.length - 1; i >= 0; i--) {
            const message_block = conversation_message_blocks[i];
            const block_length = (typeof message_block === 'string') ? message_block.length : 0; // Ensure it's a string

            if (block_length === 0) continue;
            if (promptFull && systemHistoryFull) break;

            // Try fitting into the main 'prompt' first
            if (!promptFull) {
                if (current_prompt_length + block_length <= PROMPT_LIMIT) {
                    prompt_history_blocks.unshift(message_block);
                    current_prompt_length += block_length;
                    continue;
                } else {
                    promptFull = true;
                }
            }

            // If prompt is full, try fitting into the 'system_prompt' remaining space
            if (!systemHistoryFull) {
                if (current_system_history_length + block_length <= remaining_system_limit) {
                    system_prompt_history_blocks.unshift(message_block);
                    current_system_history_length += block_length;
                    continue;
                } else {
                    systemHistoryFull = true;
                }
            }
        }

        // 5. Combine the final prompt and system_prompt parts
        const system_prompt_history_content = system_prompt_history_blocks.join('').trim();
        const final_prompt = prompt_history_blocks.join('').trim();
        const SEPARATOR = "\n\n------- Earlier Conversation History -------\n\n";
        let final_system_prompt = "";
        const hasFixedSystem = fixed_system_prompt_content.length > 0;
        const hasSystemHistory = system_prompt_history_content.length > 0;

        if (hasFixedSystem && hasSystemHistory) {
            final_system_prompt = fixed_system_prompt_content + SEPARATOR + system_prompt_history_content;
        } else if (hasFixedSystem) {
            final_system_prompt = fixed_system_prompt_content;
        } else if (hasSystemHistory) {
            final_system_prompt = system_prompt_history_content;
        }

        // 6. ** Crucially, always return an object **
        const result = {
            system_prompt: final_system_prompt,
            prompt: final_prompt
        };
        // console.log("<<< Exiting convertMessagesToFalPrompt successfully."); // Optional success log
        return result;

    } catch (internalError) { // *** Catch any unexpected errors during processing ***
        console.error("!!! CRITICAL ERROR inside convertMessagesToFalPrompt processing:", internalError);
        // Log the input that caused the error for debugging
        console.error("!!! Failing messages input was:", JSON.stringify(messages, null, 2));
        // Re-throw the error so the main handler knows something went wrong during setup.
        // This will be caught by the outer try...catch in the route handler.
        throw new Error(`Failed to process messages internally: ${internalError.message}`);
    }
}
// === End convertMessagesToFalPrompt function ===


/**
 * Makes a request to the Fal AI API, handling key rotation and retries on key-related errors.
 * For stream requests, returns the stream AND the key info used.
 * @param {object} falInput - The input object for the Fal AI API call.
 * @param {boolean} [stream=false] - Whether to make a streaming request.
 * @returns {Promise<object|{stream: AsyncIterable<object>, keyUsed: string, indexUsed: number}>}
 *          The result object for non-stream, or an object containing the stream and key info for stream.
 * @throws {Error} If the request fails after trying all valid keys, or if a non-key-related error occurs during *initiation*.
 */
async function makeFalRequestWithRetry(falInput, stream = false) {
    let attempts = 0;
    const maxAttempts = falKeys.length;
    const attemptedKeysInThisRequest = new Set();

    while (attempts < maxAttempts) {
        const keyInfo = getNextValidKey();
        if (!keyInfo) {
            console.error("makeFalRequestWithRetry: No valid Fal AI keys remaining.");
            throw new Error("No valid Fal AI keys available (all marked as invalid).");
        }
        if (attemptedKeysInThisRequest.has(keyInfo.key)) {
            console.warn(`Key at index ${keyInfo.index} was already attempted for this request. Skipping to find next different key.`);
            continue;
        }
        attemptedKeysInThisRequest.add(keyInfo.key);
        attempts++;

        try {
            console.log(`Attempt ${attempts}/${maxAttempts}: Trying Fal Key index ${keyInfo.index}...`);
            console.warn(`Configuring GLOBAL fal client with key index ${keyInfo.index}. Review concurrency implications.`);
            fal.config({ credentials: keyInfo.key });

            if (stream) {
                const falStream = await fal.stream("fal-ai/any-llm", { input: falInput });
                console.log(`Successfully initiated stream with key index ${keyInfo.index}.`);
                // Return stream AND key info
                return { stream: falStream, keyUsed: keyInfo.key, indexUsed: keyInfo.index };
            } else {
                // Non-stream logic
                console.log(`Executing non-stream request with key index ${keyInfo.index}...`);
                const result = await fal.subscribe("fal-ai/any-llm", { input: falInput, logs: true });
                console.log(`Successfully received non-stream result with key index ${keyInfo.index}.`);
                if (result && result.error) {
                     console.error(`Fal AI returned an error object within the non-stream result payload (Key Index ${keyInfo.index}):`, result.error);
                     if (isKeyRelatedError(result.error)) {
                         console.warn(`Marking Fal Key index ${keyInfo.index} as invalid due to error in response payload.`);
                         invalidKeys.add(keyInfo.key);
                         continue; // Try next key
                     } else {
                         throw new Error(`Fal AI error reported in result payload: ${JSON.stringify(result.error)}`);
                     }
                 }
                return result; // Return only result for non-stream
            }
        } catch (error) {
            // This catch block handles errors during *request initiation*
            console.error(`Error caught during request initiation using Fal Key index ${keyInfo.index}:`, error.message || error);
            if (isKeyRelatedError(error)) {
                console.warn(`Marking Fal Key index ${keyInfo.index} as invalid due to caught initiation error.`);
                invalidKeys.add(keyInfo.key);
                // Continue loop to try the next key
            } else {
                console.error("Initiation error does not appear to be key-related. Failing request without further key retries.");
                throw error; // Re-throw non-key-related initiation error
            }
        }
    } // End while loop

    // If loop finishes, all keys failed during initiation
    throw new Error(`Request initiation failed after trying ${attempts} unique Fal key(s). All failed with key-related errors or were already marked invalid.`);
}


// POST /v1/chat/completions endpoint - Handles chat requests
app.post('/v1/chat/completions', async (req, res) => {
    const { model, messages, stream = false, reasoning = false, ...restOpenAIParams } = req.body;
    console.log(`--> POST /v1/chat/completions | Model: ${model} | Stream: ${stream}`);

    // --- Input validation for model and messages ---
     if (!model || !messages || !Array.isArray(messages) || messages.length === 0) {
        console.error("Invalid request: Missing 'model' or 'messages' array is empty/invalid.");
        return res.status(400).json({ error: 'Bad Request: `model` and a non-empty `messages` array are required.' });
    }
    if (!FAL_SUPPORTED_MODELS.includes(model)) {
        console.warn(`Warning: Requested model '${model}' is not in the explicitly supported list. Proxy will still attempt the request.`);
    }
    // --- End Input Validation ---

    let keyUsedForRequest = null; // Variable to store the key used if initiation succeeds
    let indexUsedForRequest = null;

    try {
        // --- Prepare Fal AI Input ---
        // This might throw an error if convertMessagesToFalPrompt fails internally
        const { prompt, system_prompt } = convertMessagesToFalPrompt(messages);

        const falInput = {
            model: model,
            prompt: prompt,
            reasoning: !!reasoning,
        };
        if (system_prompt && system_prompt.length > 0) {
             falInput.system_prompt = system_prompt;
        }
        // --- End Prepare Input ---

        console.log("Attempting Fal request with key rotation/retry logic...");
        console.log(`Prepared Input Lengths - System Prompt: ${system_prompt?.length || 0}, Prompt: ${prompt?.length || 0}`);

        // --- Handle Stream vs Non-Stream ---
        if (stream) {
            // Set headers for Server-Sent Events (SSE)
            res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
            res.setHeader('Cache-Control', 'no-cache');
            res.setHeader('Connection', 'keep-alive');
            res.setHeader('Access-Control-Allow-Origin', '*'); // Adjust CORS for production if needed
            res.flushHeaders(); // Send headers immediately

            let previousOutput = ''; // Track previous output for delta calculation
            let streamResult; // To hold the object { stream, keyUsed, indexUsed }

            try {
                // **Initiate the stream using the retry helper**
                // This can throw if initiation fails after all retries
                streamResult = await makeFalRequestWithRetry(falInput, true);
                const falStream = streamResult.stream;
                keyUsedForRequest = streamResult.keyUsed; // Store the key used for this stream
                indexUsedForRequest = streamResult.indexUsed;

                // Process the stream events asynchronously
                for await (const event of falStream) {
                    // Safely extract data from the event
                    const currentOutput = (event && typeof event.output === 'string') ? event.output : '';
                    const isPartial = (event && typeof event.partial === 'boolean') ? event.partial : true;
                    const errorInfo = (event && event.error) ? event.error : null;

                    // Handle errors reported *within* a stream event payload
                    if (errorInfo) {
                        console.error("Error received *within* fal stream event payload:", errorInfo);
                        const errorChunk = {
                            id: `chatcmpl-${Date.now()}-error`, object: "chat.completion.chunk", created: Math.floor(Date.now() / 1000), model: model,
                            choices: [{ index: 0, delta: {}, finish_reason: "error", message: { role: 'assistant', content: `Fal Stream Event Error: ${JSON.stringify(errorInfo)}` } }]
                        };
                        if (!res.writableEnded) { res.write(`data: ${JSON.stringify(errorChunk)}\n\n`); }
                        else { console.warn("Stream ended before writing event error."); }
                        // Potentially break or add logic based on error type here
                    }

                    // Calculate the delta (new content)
                    let deltaContent = '';
                    if (currentOutput.startsWith(previousOutput)) {
                        deltaContent = currentOutput.substring(previousOutput.length);
                    } else if (currentOutput.length > 0) {
                         console.warn("Fal stream output mismatch/reset. Sending full current output as delta.");
                         deltaContent = currentOutput;
                         previousOutput = '';
                    }
                    previousOutput = currentOutput;

                    // Send OpenAI-compatible SSE chunk
                    if (deltaContent || !isPartial) {
                        const openAIChunk = {
                            id: `chatcmpl-${Date.now()}`,
                            object: "chat.completion.chunk",
                            created: Math.floor(Date.now() / 1000),
                            model: model,
                            choices: [{
                                index: 0,
                                delta: { content: deltaContent },
                                finish_reason: isPartial === false ? "stop" : null
                             }]
                        };
                        if (!res.writableEnded) { res.write(`data: ${JSON.stringify(openAIChunk)}\n\n`); }
                        else { console.warn("Stream ended before writing data chunk."); }
                    }
                } // End for-await loop

                // Send the final [DONE] marker
                if (!res.writableEnded) {
                    res.write(`data: [DONE]\n\n`);
                    res.end(); // Close the connection
                    console.log("<-- Stream finished successfully and [DONE] sent.");
                } else {
                    console.log("<-- Stream processing finished, but connection was already ended before [DONE].");
                }

            } catch (streamError) {
                // **Catch block for errors during stream processing OR initiation failure**
                console.error('Error during stream request processing/initiation:', streamError.message || streamError);

                // **Check if the error is key-related AND if initiation succeeded**
                // This ensures we only invalidate the key if the error happened *during* processing
                // using a key that successfully initiated the stream.
                if (keyUsedForRequest && isKeyRelatedError(streamError)) {
                    console.warn(`Marking Fal Key index ${indexUsedForRequest} as invalid due to error during stream processing.`);
                    invalidKeys.add(keyUsedForRequest);
                }
                // else: Error was either non-key-related or happened during initiation (already handled/logged in makeFalRequestWithRetry).

                // --- Report error back to the client ---
                 try {
                     if (!res.headersSent) {
                         // Error likely during initiation (caught from makeFalRequestWithRetry)
                         const errorMessage = (streamError instanceof Error) ? streamError.message : JSON.stringify(streamError);
                         res.status(502).json({ error: 'Failed to initiate Fal stream', details: errorMessage });
                         console.log("<-- Stream initiation failed response sent (502).");
                     } else if (!res.writableEnded) {
                         // Error during stream processing after headers sent
                         const errorDetails = (streamError instanceof Error) ? streamError.message : JSON.stringify(streamError);
                         res.write(`data: ${JSON.stringify({ error: { message: "Stream processing error after initiation", type: "proxy_error", details: errorDetails } })}\n\n`);
                         res.write(`data: [DONE]\n\n`); // Still send DONE for robust client handling
                         res.end();
                         console.log("<-- Stream error sent within stream, stream ended.");
                     } else {
                        console.log("<-- Stream error occurred, but connection was already ended.");
                     }
                 } catch (finalError) {
                    console.error('Error sending stream error message to client:', finalError);
                    if (!res.writableEnded) { res.end(); }
                 }
                 // --- End error reporting ---
            }

        } else {
            // --- Non-Stream Request ---
            try {
                // Get the result using the retry helper
                const result = await makeFalRequestWithRetry(falInput, false);

                // Construct OpenAI compatible response
                const openAIResponse = {
                    id: `chatcmpl-${result.requestId || Date.now()}`,
                    object: "chat.completion",
                    created: Math.floor(Date.now() / 1000),
                    model: model,
                    choices: [{
                        index: 0,
                        message: {
                            role: "assistant",
                            content: result.output || "" // Ensure content is string
                         },
                        finish_reason: "stop"
                    }],
                    usage: { prompt_tokens: null, completion_tokens: null, total_tokens: null },
                    system_fingerprint: null,
                    ...(result.reasoning && { fal_reasoning: result.reasoning }),
                };

                res.json(openAIResponse);
                console.log("<-- Non-stream response sent successfully.");

            } catch (error) {
                // Catches errors from makeFalRequestWithRetry (e.g., all keys failed or non-key error)
                console.error('Error during non-stream request processing:', error.message || error);
                if (!res.headersSent) {
                    const errorMessage = (error instanceof Error) ? error.message : JSON.stringify(error);
                    const finalMessage = errorMessage.includes("No valid Fal AI keys available") || errorMessage.includes("Request failed after trying")
                        ? `Fal request failed: ${errorMessage}`
                        : `Internal Server Error processing Fal request: ${errorMessage}`;
                    res.status(502).json({ error: 'Fal Request Failed', details: finalMessage });
                    console.log("<-- Non-stream error response sent (502).");
                } else {
                    console.error("Headers already sent for non-stream error response? This is unexpected.");
                    if (!res.writableEnded) { res.end(); }
                }
            }
        } // --- End Stream/Non-Stream Logic ---

    } catch (error) {
        // **Catch block for errors BEFORE Fal request attempt**
        // (e.g., errors from convertMessagesToFalPrompt, JSON parsing errors)
        console.error('Unhandled error before initiating Fal request (likely setup or input conversion):', error.message || error);
        if (!res.headersSent) {
            const errorMessage = (error instanceof Error) ? error.message : JSON.stringify(error);
            // Use 500 Internal Server Error for issues within the proxy itself during setup
            res.status(500).json({ error: 'Internal Server Error in Proxy Setup', details: errorMessage });
            console.log("<-- Proxy setup error response sent (500).");
        } else {
             console.error("Headers already sent when catching setup error. Ending response.");
             if (!res.writableEnded) { res.end(); }
        }
    }
});

// Start the Express server
app.listen(PORT, () => {
    console.log(`=====================================================================`);
    console.log(` Fal OpenAI Proxy Server (Multi-Key Rotation & Failover)`);
    console.log(`---------------------------------------------------------------------`);
    console.log(` Listening on port   : ${PORT}`);
    console.log(` Reading Fal Keys from : FAL_KEY environment variable (comma-separated)`);
    console.log(` Loaded Keys Count   : ${falKeys.length}`);
    console.log(` Invalid Keys Set    : Initialized (size: ${invalidKeys.size})`);
    console.log(` Proxy API Key Auth  : ${API_KEY ? 'Enabled (using API_KEY env var)' : 'DISABLED'}`);
    console.log(` Input Limits        : System Prompt=${SYSTEM_PROMPT_LIMIT}, Prompt=${PROMPT_LIMIT}`);
    console.log(` Concurrency Warning : Global Fal client reconfigured per request attempt!`);
    console.log(`---------------------------------------------------------------------`);
    console.log(` Endpoints Available:`);
    console.log(`   POST http://localhost:${PORT}/v1/chat/completions`);
    console.log(`   GET  http://localhost:${PORT}/v1/models`);
    console.log(`=====================================================================`);
});

// Root path handler for basic health check / info
app.get('/', (req, res) => {
    res.send(`Fal OpenAI Proxy (Multi-Key Rotation/Failover from FAL_KEY) is running. Loaded ${falKeys.length} key(s). Currently ${invalidKeys.size} key(s) marked as invalid.`);
});