import dotenv from 'dotenv';
dotenv.config();

import helmet from 'helmet';
import express from 'express';
import puppeteer from 'puppeteer-core';
import { logError, logInfo } from './logger.mjs';
import 'express-async-errors';
import rateLimit from 'express-rate-limit';
import cors from 'cors';

const app = express();

// Middleware for API key validation
const validateApiKey = (req, res, next) => {
    const apiKey = req.header('x-api-key');

    if (!apiKey || apiKey !== process.env.API_KEY) {
        return res.status(401).json({
            success: false,
            message: 'Invalid or missing API key'
        });
    }
    next();
};

app.use(helmet());
app.use(cors());
app.use(express.json({ limit: process.env.BODY_SIZE_LIMIT || '10mb' }));
app.use(validateApiKey);

// Rate limit setup
const limiter = rateLimit({
    windowMs: parseInt(process.env.RATE_LIMIT_WINDOW_MS, 10) || 60 * 60 * 1000,
    max: parseInt(process.env.RATE_LIMIT_MAX, 10) || 500,
    message: {
        success: false,
        message: 'Too many requests from this IP, please try again after an hour.'
    },
    keyGenerator: (req) => req.socket.remoteAddress,
    standardHeaders: true,
    legacyHeaders: false,
});

app.use(limiter);

// Handle large input errors
app.use((err, req, res, next) => {
    if (err.type === 'entity.too.large') {
        return res.status(413).json({
            success: false,
            message: 'Payload too large. Please reduce the size of your input.',
        });
    }
    next(err);
});

// Usage endpoint
app.get('/', (req, res) => {
    res.status(200).json({
        success: true,
        message: 'Welcome to the Image Generation API!',
        usage: {
            endpoint: '/generate-image',
            method: 'POST',
            description: 'Generate an image from HTML template.',
            parameters: {
                htmlTemplate: 'HTML template as a string (required)',
                data: 'An object containing dynamic data to replace in the template (optional)',
            },
            exampleRequest: {
                htmlTemplate: '<h1>{{title}}</h1>',
                data: { title: 'Hello World' },
            },
            exampleResponse: {
                success: true,
                message: 'Image generated successfully',
                image: '<base64-image>',
            },
        },
    });
});

// Generate image endpoint
app.post('/generate-image', async (req, res) => {
    const { htmlTemplate, data = {} } = req.body;

    if (!htmlTemplate) {
        return res.status(400).json({
            success: false,
            message: 'htmlTemplate is required',
        });
    }

    try {
        const base64Data = await generateImageFromHtml({ htmlTemplate, data });

        res.status(200).json({
            success: true,
            message: 'Image generated successfully',
            image: base64Data,
        });
    } catch (error) {
        logError('Error generating image:', error);
        res.status(500).json({
            success: false,
            message: 'An error occurred while generating the image',
            error: error.message,
        });
    }
});

let browser;

async function initBrowser() {
    try {
        if (!browser || !(await browser.isConnected())) {
            browser = await puppeteer.launch({
                args: [
                    '--no-sandbox',
                    '--disable-setuid-sandbox',
                    '--disable-dev-shm-usage',
                    '--disable-cache',
                    '--disable-accelerated-2d-canvas',
                    '--disable-gpu',
                    '--disable-extensions',
                ],
                executablePath: process.env.PUPPETEER_EXECUTABLE_PATH,
                headless: process.env.PUPPETEER_HEADLESS === 'false' ? false : "new",
                ignoreDefaultArgs: ['--enable-automation'],
            });
            logInfo("Browser initialized successfully.");
        }
    } catch (error) {
        logError("Error initializing the browser:", error);
        throw new Error('Failed to initialize the browser: ' + error.message);
    }
}

async function closeBrowser() {
    if (browser) {
        try {
            await browser.close();
            logInfo("Browser closed successfully.");
        } catch (error) {
            logError("Error closing the browser:", error);
        } finally {
            browser = null;
        }
    }
}

async function generateImageFromHtml({ htmlTemplate, data = {}, retryCount = 3 }) {
    let page;
    for (let attempt = 0; attempt < retryCount; attempt++) {
        try {
            await initBrowser();

            page = await browser.newPage();
            await page.setCacheEnabled(false);
            await page.setRequestInterception(true);

            page.on('request', (request) => {
                if (['redirect'].includes(request.redirectChain().length)) {
                    request.abort();
                } else {
                    request.continue();
                }
            });

            const html = Object.keys(data).reduce(
                (acc, key) => acc.replace(new RegExp(`{{${key}}}`, 'g'), data[key]),
                htmlTemplate
            );

            await page.setViewport({
                width: 1920,
                height: 1080,
                deviceScaleFactor: 2
            });

            await page.setContent(html, { waitUntil: 'networkidle0' });
            await new Promise(r => setTimeout(r, 2000));

            const screenshotOptions = {
                type: process.env.IMAGE_TYPE === "png" ? 'png' : "jpeg",
                encoding: 'base64',
                fullPage: true
            };

            if (process.env.IMAGE_TYPE !== "png" && process.env.IMAGE_QUALITY) {
                screenshotOptions.quality = parseInt(process.env.IMAGE_QUALITY, 10);
            }

            const base64Data = await page.screenshot(screenshotOptions);
            return base64Data;

        } catch (error) {
            logError(`Attempt ${attempt + 1} failed:`, error);
            if (attempt === retryCount - 1) {
                throw new Error('Error generating image: ' + error.message);
            }
            await closeBrowser();
            await new Promise(r => setTimeout(r, 2000));
        } finally {
            if (page) {
                await page.close();
            }
        }
    }
}

function errorHandler(err, req, res, next) {
    logError('Error:', err);
    res.status(err.status || 500).json({
        success: false,
        message: err.message || 'Internal Server Error',
        stack: process.env.NODE_ENV === 'development' ? err.stack : undefined,
    });
}

app.use(errorHandler);

const PORT = process.env.PORT || 3000;
app.listen(PORT, () => {
    logInfo(`Server is running on port ${PORT}`);
});

process.on('exit', closeBrowser);
process.on('SIGINT', closeBrowser);
process.on('SIGTERM', closeBrowser);
process.on('uncaughtException', async (error) => {
    console.error('Uncaught exception:', error);
    await closeBrowser();
    process.exit(1);
});