import express from 'express';
// **CRITICAL: Import the 'fal' object directly for configuration**
import { fal } from '@fal-ai/client'; // Make sure 'fal' is the object you configure

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

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

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

if (!API_KEY) {
    console.error("Error: API_KEY environment variable is not set.");
    process.exit(1);
}

let currentKeyIndex = 0;
// **NEW: 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.`);

// **MODIFIED: Function to get the next *valid* key**
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++;
        // If the loop started at an invalid key and wrapped around,
        // check if we've checked all keys since the start
        if (currentKeyIndex === initialIndex && attempts > 0) {
            // This check might be redundant due to the invalidKeys.size check at the top,
            // but serves as an extra safeguard.
             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;
}

// **NEW/MODIFIED: Function to check if an error is likely related to a bad key**
// NOTE: This is a heuristic. You might need to adjust based on actual errors from Fal AI.
function isKeyRelatedError(error) {
    // Check if error itself is null/undefined
    if (!error) return false;

    const message = error.message?.toLowerCase() || '';
    // Check if the error object has a 'status' property (common in HTTP errors)
    const status = error.status || error.statusCode; // Check common status properties

    // Check for specific HTTP status codes indicative of auth/permission issues
    // 401: Unauthorized, 403: Forbidden, 429: Too Many Requests (often quota related)
    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)
    // Add more specific Fal AI error messages as you encounter them
    if (message.includes('invalid api key') ||
        message.includes('authentication failed') ||
        message.includes('permission denied') ||
        message.includes('quota exceeded') || // Treat quota errors as key-related for rotation
        message.includes('forbidden') ||
        message.includes('unauthorized') ||
        message.includes('rate limit') || // Often linked to key limits
        message.includes('credentials')) // Generic credential errors
        {
        console.warn(`Detected potential key-related error (message: "${message}")`);
        return true;
    }

    // Add more specific checks based on observed Fal AI errors if needed
    // Example: Check for specific error codes if Fal AI provides them
    // if (error.code === 'SOME_FAL_AUTH_ERROR_CODE') {
    //     return true;
    // }

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

const app = express();
app.use(express.json({ limit: '50mb' }));
app.use(express.urlencoded({ extended: true, limit: '50mb' }));

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

// API Key 鉴权中间件 (unchanged)
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' });
    }

    const authParts = authHeader.split(' ');
    if (authParts.length !== 2 || authParts[0].toLowerCase() !== 'bearer') {
        console.warn('Unauthorized: Invalid Authorization header format');
        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');
        return res.status(401).json({ error: 'Unauthorized: Invalid API Key' });
    }

    next();
};

app.use(['/v1/models', '/v1/chat/completions'], apiKeyAuth);

// === 全局定义限制 === (unchanged)
const PROMPT_LIMIT = 4800;
const SYSTEM_PROMPT_LIMIT = 4800;
// === 限制定义结束 ===

// 定义 fal-ai/any-llm 支持的模型列表 (unchanged)
const FAL_SUPPORTED_MODELS = [ /* ... model list ... */ ];
// Helper function getOwner (unchanged)
const getOwner = (modelId) => { /* ... */ };
// GET /v1/models endpoint (unchanged)
app.get('/v1/models', (req, res) => { /* ... */ });
// convertMessagesToFalPrompt 函数 (unchanged)
function convertMessagesToFalPrompt(messages) { /* ... */ }
// === convertMessagesToFalPrompt 函数结束 ===


// --- NEW: Helper function to make Fal AI request with retries ---
async function makeFalRequestWithRetry(falInput, stream = false) {
    let attempts = 0;
    // Max attempts should be the total number of keys initially available
    const maxAttempts = falKeys.length;
    // Keep track of keys tried *within this specific request attempt*
    // This prevents infinite loops if getNextValidKey had issues,
    // and ensures we try each *available* key at most once per request.
    const attemptedKeysInThisRequest = new Set();

    while (attempts < maxAttempts) {
        const keyInfo = getNextValidKey(); // Get the next *valid* key info { key, index }

        if (!keyInfo) {
            // This should only happen if all keys are in the invalidKeys set
            console.error("makeFalRequestWithRetry: No valid Fal AI keys remaining.");
            throw new Error("No valid Fal AI keys available (all marked as invalid).");
        }

        // Prevent retrying the exact same key if getNextValidKey logic had an issue
        // or if a key wasn't marked invalid correctly on a previous attempt within this request.
        if (attemptedKeysInThisRequest.has(keyInfo.key)) {
            console.warn(`Key at index ${keyInfo.index} was already attempted for this request. Skipping to find next.`);
            // Don't increment 'attempts' here as we didn't use the key. Let the loop find the next.
            // If all keys end up being skipped, the `!keyInfo` check or `attempts < maxAttempts` will eventually handle it.
             continue;
        }
        attemptedKeysInThisRequest.add(keyInfo.key);
        attempts++; // Count this as a distinct attempt with a unique key for this request

        try {
            console.log(`Attempt ${attempts}/${maxAttempts}: Trying Fal Key index ${keyInfo.index}...`);

            // *** CRITICAL: Reconfigure the global fal client with the selected key for this attempt ***
            // Warning: This reconfigures the GLOBAL client. If you have many concurrent requests,
            // this could lead to race conditions. Consider instance isolation or a pool manager
            // for high-concurrency scenarios if the library doesn't support per-request credentials easily.
            console.warn(`Configuring GLOBAL fal client with key index ${keyInfo.index}. Ensure this is safe for your concurrency model.`);
            fal.config({ credentials: keyInfo.key }); // Use the specific key for this attempt

            if (stream) {
                // Use the configured global 'fal' object for the stream request
                const falStream = await fal.stream("fal-ai/any-llm", { input: falInput });
                console.log(`Successfully initiated stream with key index ${keyInfo.index}.`);
                // Success! Return the stream iterator directly
                return falStream;
            } else {
                // Use the configured global 'fal' object for the non-stream request
                console.log(`Executing non-stream request with key index ${keyInfo.index}...`);
                // Assuming fal.subscribe or similar method for non-streaming
                // Adapt this line if your non-stream method is different
                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}.`);

                // Check for errors *within* the successful response structure (if applicable)
                if (result && result.error) {
                     console.error(`Fal-ai returned an error in the result payload (Key Index ${keyInfo.index}):`, result.error);
                     // Decide if this specific payload error should also invalidate the key
                     if (isKeyRelatedError(result.error)) { // Reuse the checker
                         console.warn(`Marking Fal Key index ${keyInfo.index} as invalid due to error in response payload.`);
                         invalidKeys.add(keyInfo.key);
                         // Continue the loop to try the next key
                         continue; // Go to the next iteration of the while loop
                     } else {
                         // Throw an error that will be caught by the outer handler, not retried
                         throw new Error(`Fal-ai error in result payload: ${JSON.stringify(result.error)}`);
                     }
                 }
                 // Success! Return the result
                return result;
            }
        } catch (error) {
            console.error(`Error using Fal Key index ${keyInfo.index}:`, error.message || error);

            // Check if the caught error indicates the key is invalid
            if (isKeyRelatedError(error)) {
                console.warn(`Marking Fal Key index ${keyInfo.index} as invalid due to caught error.`);
                // **ACTION: Add the failed key to the set of invalid keys**
                invalidKeys.add(keyInfo.key);
                // Continue to the next iteration of the while loop to try another key
            } else {
                // Error is not key-related (e.g., network issue, bad input, internal Fal error)
                // Don't retry with other keys for this type of error. Fail the request immediately.
                console.error("Error does not appear to be key-related. Failing request without further key retries.");
                throw error; // Re-throw the original error to be caught by the main endpoint handler
            }
        }
    } // End while loop

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


// POST /v1/chat/completions endpoint (MODIFIED to use retry logic)
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}`);

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

    try {
        // --- Prepare Input (unchanged) ---
        const { prompt, system_prompt } = convertMessagesToFalPrompt(messages);
        const falInput = {
            model: model,
            prompt: prompt,
            ...(system_prompt && { system_prompt: system_prompt }),
            reasoning: !!reasoning,
        };
        console.log("Attempting Fal request with key rotation/retry...");

        // --- Handle Stream vs Non-Stream using the retry helper ---
        if (stream) {
            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', '*');
            res.flushHeaders();

            let previousOutput = '';
            let falStream; // Declare falStream here

            try {
                // **MODIFIED: Initiate stream using the retry logic helper**
                falStream = await makeFalRequestWithRetry(falInput, true);

                // Process the stream events (logic mostly unchanged)
                for await (const event of falStream) {
                    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 *within* the stream payload
                    if (errorInfo) {
                        console.error("Error received *within* fal stream event:", errorInfo);
                        // Optionally send an error chunk (check if needed)
                        const errorChunk = { /* ... error chunk details ... */ };
                        if (!res.writableEnded) {
                             res.write(`data: ${JSON.stringify(errorChunk)}\n\n`);
                        } else { console.warn("Stream ended before writing event error."); }
                         // Decide if this stream error should also invalidate the key
                         // Note: The key might have already been marked invalid if the error happened during initial connection
                         // You might need more context to decide if an *in-stream* error means the key is bad.
                         // For now, we just log it. If it causes the stream to terminate prematurely,
                         // the outer catch might handle it, but key invalidation might not occur unless the
                         // error object passed to the catch block triggers isKeyRelatedError.
                    }

                    // Calculate delta (logic unchanged)
                    let deltaContent = '';
                    if (currentOutput.startsWith(previousOutput)) {
                         deltaContent = currentOutput.substring(previousOutput.length);
                    } else if (currentOutput.length > 0) {
                         deltaContent = currentOutput;
                         previousOutput = '';
                    }
                    previousOutput = currentOutput;

                    // Send OpenAI compatible chunk (logic unchanged)
                    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 (logic unchanged)
                if (!res.writableEnded) {
                    res.write(`data: [DONE]\n\n`);
                    res.end();
                    console.log("<-- Stream finished successfully.");
                } else {
                    console.log("<-- Stream finished, but connection was already ended.");
                }

            } catch (streamError) {
                // **MODIFIED: Catches errors from makeFalRequestWithRetry OR stream iteration**
                // This error could be a non-key-related error, or the "all keys failed" error.
                console.error('Error during stream request processing:', streamError.message || streamError);
                 try {
                     if (!res.headersSent) {
                         // Error likely occurred in makeFalRequestWithRetry before stream started
                         const errorMessage = (streamError instanceof Error) ? streamError.message : JSON.stringify(streamError);
                         res.status(502).json({ // 502 Bad Gateway is appropriate for upstream failure
                            error: 'Failed to initiate Fal stream',
                            details: errorMessage
                         });
                         console.log("<-- Stream initiation failed response sent.");
                     } else if (!res.writableEnded) {
                         // Stream started but failed during processing
                         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 client handling
                         res.end();
                         console.log("<-- Stream error sent, stream ended.");
                     } else {
                        console.log("<-- Stream error occurred, but connection already ended.");
                     }
                 } catch (finalError) {
                    console.error('Error sending stream error message to client:', finalError);
                    if (!res.writableEnded) { res.end(); }
                 }
            }

        } else {
            // --- Non-Stream ---
            try {
                // **MODIFIED: Get the result using the retry logic helper**
                const result = await makeFalRequestWithRetry(falInput, false);
                // console.log("Received non-stream result via retry function:", JSON.stringify(result, null, 2));

                // Construct OpenAI compatible response (logic unchanged)
                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 || "" }, 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) {
                 // **MODIFIED: Catches errors from makeFalRequestWithRetry**
                 // This error could be a non-key-related error, or the "all keys failed" 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);
                    // Check if it was the "all keys failed" error
                    const finalMessage = errorMessage.includes("No valid Fal AI keys available") || errorMessage.includes("Request failed after trying")
                        ? `Fal request failed: ${errorMessage}` // Simplified message
                        : `Internal Server Error processing Fal request: ${errorMessage}`;
                    res.status(502).json({ error: 'Fal Request Failed', details: finalMessage }); // 502 Bad Gateway
                    console.log("<-- Non-stream error response sent.");
                } else {
                    console.error("Headers already sent for non-stream error? This is unexpected.");
                    if (!res.writableEnded) { res.end(); }
                }
            }
        }

    } catch (error) {
        // Catch errors from parameter validation or prompt conversion *before* calling Fal
        console.error('Unhandled error before initiating Fal request:', error.message || error);
        if (!res.headersSent) {
            const errorMessage = (error instanceof Error) ? error.message : JSON.stringify(error);
            res.status(500).json({ error: 'Internal Server Error in Proxy Setup', details: errorMessage });
            console.log("<-- Proxy setup error response sent.");
        } else {
             console.error("Headers already sent when catching setup error. Ending response.");
             if (!res.writableEnded) { res.end(); }
        }
    }
});

// 启动服务器 (Updated startup message)
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})`); // Show invalid set size
    console.log(` 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:`);
    console.log(`   POST http://localhost:${PORT}/v1/chat/completions`);
    console.log(`   GET  http://localhost:${PORT}/v1/models`);
    console.log(`=====================================================================`);
});

// 根路径响应 (Updated message)
app.get('/', (req, res) => {
    res.send(`Fal OpenAI Proxy (Multi-Key Rotation from FAL_KEY) is running. Loaded ${falKeys.length} key(s). Currently ${invalidKeys.size} key(s) marked as invalid.`); // Show invalid count
});