import fs from "fs/promises";
import path from "path";
import express from "express";
import cors from "cors";
import morgan from "morgan";
import puppeteer from "puppeteer-extra";
import StealthPlugin from "puppeteer-extra-plugin-stealth";

puppeteer.use(StealthPlugin());

const PORT = process.env.PORT ? Number(process.env.PORT) : 4000;
const API_KEY = process.env.API_KEY || ""; // nếu đặt, yêu cầu query ?key=... hoặc header x-api-key
const COOKIES_DIR = path.resolve("./cookies");
await fs.mkdir(COOKIES_DIR, { recursive: true });

const app = express();
app.use(cors());
app.use(express.json({ limit: "2mb" }));
if (process.env.NODE_ENV !== "production") {
  app.use(morgan("combined"))
}

function requireApiKey(req, res, next) {
    if (!API_KEY) return next();
    const key = req.query.key || req.headers["x-api-key"] || req.headers["authorization"];
    if (!key || !(key === API_KEY || key === `Bearer ${API_KEY}`)) {
        return res.status(403).json({ error: "Forbidden - invalid API key" });
    }
    next();
}

// single browser instance reused across requests
let browser;
async function ensureBrowser() {
    if (browser) return browser;
    browser = await puppeteer.launch({
        headless: "new",
        args: [
            "--no-sandbox",
            "--disable-setuid-sandbox",
            "--disable-dev-shm-usage",
            "--disable-blink-features=AutomationControlled",
            "--window-size=1200,900",
            "--disable-features=IsolateOrigins,site-per-process",
        ],

    });
    // close on exit
    const closeHandler = async () => {
        try { if (browser) await browser.close(); } catch (e) { }
        process.exit();
    };
    process.on("SIGINT", closeHandler);
    process.on("SIGTERM", closeHandler);
    return browser;
}

async function cookieFileFor(url) {
    try {
        const u = new URL(url);
        const file = path.join(COOKIES_DIR, `${u.hostname}.json`);
        return file;
    } catch {
        return null;
    }
}

async function loadCookies(page, url) {
    const file = await cookieFileFor(url);
    if (!file) return;
    try {
        const raw = await fs.readFile(file, "utf-8");
        const cookies = JSON.parse(raw);
        if (Array.isArray(cookies) && cookies.length > 0) {
            await page.setCookie(...cookies);
        }
    } catch (e) {
        // ignore if no file
    }
}

async function saveCookies(page, url) {
    const file = await cookieFileFor(url);
    if (!file) return;
    try {
        const cookies = await page.cookies();
        await fs.writeFile(file, JSON.stringify(cookies, null, 2), "utf-8");
    } catch (e) {
        console.warn("saveCookies error", e.message);
    }
}

// forwarded headers we allow the client to set
const ALLOWED_FORWARD = ["authorization", "content-type", "accept", "cookie", "user-agent"];

function buildForwardHeaders(req) {
    const h = {};
    for (const k of Object.keys(req.headers)) {
        if (ALLOWED_FORWARD.includes(k.toLowerCase())) {
            h[k] = req.headers[k];
        }
    }
    return h;
}

app.get("/", (req, res) => res.json({ ok: true, msg: "puppeteer-stealth-proxy" }));

// proxy GET/HEAD (returns raw response from target, using page.goto)
app.all("/proxy", requireApiKey, async (req, res) => {
    const target = req.query.url;
    if (!target) return res.status(400).json({ error: "Missing url query param" });
    let page;
    try {
        const b = await ensureBrowser();
        page = await b.newPage();

        // set UA (override headless UA)
        const ua = req.headers["user-agent"] ||
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36";
        await page.setUserAgent(ua);

        // load cookies if any
        await loadCookies(page, target);

        // forward headers
        const extraHeaders = buildForwardHeaders(req);
        if (Object.keys(extraHeaders).length > 0) {
            await page.setExtraHTTPHeaders(extraHeaders);
        }

        // if non-GET, we intercept request to inject body
        if (req.method !== "GET" && req.method !== "HEAD") {
            await page.setRequestInterception(true);
            page.on("request", intercepted => {
                if (intercepted.isNavigationRequest() && intercepted.url() === target) {
                    // set postData only when we have a body
                    const opts = {
                        method: req.method,
                        headers: { ...intercepted.headers(), ...extraHeaders },
                    };
                    if (req.body && Object.keys(req.body).length > 0) {
                        const ct = req.headers["content-type"] || "application/json";
                        opts.postData = ct.includes("application/json") ? JSON.stringify(req.body) : req.rawBody || JSON.stringify(req.body);
                        opts.headers["content-type"] = ct;
                    }
                    intercepted.continue(opts);
                } else {
                    intercepted.continue();
                }
            });
        }

        const response = await page.goto(target, { waitUntil: "networkidle2", timeout: 45000 });
        if (!response) throw new Error("No response from page");

        // get headers & body
        const rHeaders = response.headers();
        const buf = await response.buffer();

        // persist cookies
        await saveCookies(page, target);

        // set response headers (filter some hop-by-hop headers)
        const skip = new Set(["transfer-encoding", "content-encoding", "content-length", "connection"]);
        for (const [k, v] of Object.entries(rHeaders)) {
        if (skip.has(k.toLowerCase())) continue;
        try {
            const cleanValue = Array.isArray(v) ? v.join(", ") : String(v).replace(/[\r\n]+/g, " ").trim();
            res.setHeader(k, cleanValue);
        } catch (e) {
            console.warn(`⚠️ Skipped invalid header ${k}: ${e.message}`);
        }
        }

        res.status(response.status()).send(buf);
    } catch (err) {
        console.error("proxy error", err);
        res.status(500).json({ error: "proxy failed", details: err.message });
    } finally {
        try { if (page) await page.close(); } catch (e) { }
    }
});

// screenshot endpoint
app.get("/screenshot", requireApiKey, async (req, res) => {
    const target = req.query.url;
    if (!target) return res.status(400).json({ error: "Missing url" });
    let page;
    try {
        const b = await ensureBrowser();
        page = await b.newPage();
        await page.setViewport({ width: 1280, height: 800 });

        const ua = req.headers["user-agent"] ||
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36";
        await page.setUserAgent(ua);

        await loadCookies(page, target);

        // optionally set extra headers
        const extraHeaders = buildForwardHeaders(req);
        if (Object.keys(extraHeaders).length > 0) await page.setExtraHTTPHeaders(extraHeaders);

        await page.goto(target, { waitUntil: "networkidle2", timeout: 45000 });

        const buffer = await page.screenshot({ fullPage: true, type: "png" });

        await saveCookies(page, target);

        res.type("image/png").send(buffer);
    } catch (err) {
        console.error("screenshot error", err);
        res.status(500).json({ error: "screenshot failed", details: err.message });
    } finally {
        try { if (page) await page.close(); } catch (e) { }
    }
});

// graceful shutdown
process.on("SIGINT", async () => {
    try { if (browser) await browser.close(); } catch (e) { }
    process.exit(0);
});
process.on("SIGTERM", async () => {
    try { if (browser) await browser.close(); } catch (e) { }
    process.exit(0);
});

app.listen(PORT, () => console.log(`puppeteer-stealth-proxy listening on http://localhost:${PORT}`));