import { Worker } from 'worker_threads';
import { fileURLToPath } from 'url';
import env from '../config/env.js';

// The sandbox worker pays a large fixed startup cost: importing vm2, wasm-image-optimization, exceljs, pdfkit,
// react-canvas-label-editor AND `vfs.js` (which pulls in libdashstore) costs ~700 ms in a fresh
// thread, while the snippet itself runs in ~2 ms. ESM modules are never shared between worker
// threads, so spawning one thread per applet run re-paid that cost on every request.
//
// This pool keeps threads alive and hands them one job at a time. Threads are keyed by VFS root
// (`data/vfs/<companyId>`), so a reused thread always has the right `vfsDir` baked in and two
// companies can never share a thread.
const WORKER_PATH = fileURLToPath(new URL('../tools/javascript-worker.js', import.meta.url));
const ROOT_DIR = fileURLToPath(new URL('../..', import.meta.url));
const WORKER_BOOT_TIMEOUT_MS = 15000;

// Two threads per company rides out the common "run it, then run it again" loop without queueing.
const MAX_WORKERS_PER_COMPANY = 2;
// ...but the process is shared by every company, so the total stays bounded. Each thread holds
// WASM modules plus a NodeVM, so 5 is a deliberate memory ceiling rather than a core count.
const MAX_WORKERS_TOTAL = 5;

const DEFAULT_BUFFER_LIMIT = 32 * 1024 * 1024;

/**
 * @typedef {object} PoolEntry
 * @property {Worker} worker
 * @property {number} id        Monotonic thread id, used to tag job ids.
 * @property {string} poolKey   VFS root this thread was spawned for.
 * @property {boolean} busy
 * @property {number} lastUsed  `Date.now()` of the last job handoff (eviction order).
 */

/** @type {Array<PoolEntry>} */
const pool = [];
/** @type {Array<{ resolve: Function, reject: Function, timer: NodeJS.Timeout, poolKey: string, job: object, timeoutMs: number }>} */
const queue = [];
let workerSeq = 0;
let jobSeq = 0;
let shuttingDown = false;
/** Serializes pool growth so concurrent callers cannot overshoot either cap. */
let growing = Promise.resolve();

function countFor(poolKey) {
    return pool.filter((entry) => entry.poolKey === poolKey).length;
}

function findIdle(poolKey) {
    return pool.find((entry) => !entry.busy && entry.poolKey === poolKey);
}

function dropWorker(entry) {
    const index = pool.indexOf(entry);
    if (index !== -1) pool.splice(index, 1);
}

/**
 * Reclaims an idle thread belonging to a DIFFERENT company so a queued job can be served.
 *
 * Without this the total cap would deadlock: with 5 threads spread over three companies, a
 * fourth company's job could never be dispatched (no thread for it) and never spawn (cap hit).
 * Only genuinely idle threads are evicted, and preferentially ones whose own company has
 * nothing waiting.
 *
 * @param {string} poolKey
 * @returns {boolean} whether a thread was reclaimed
 */
function evictIdleOtherCompany(poolKey) {
    const candidates = pool.filter((entry) => !entry.busy && entry.poolKey !== poolKey);
    if (candidates.length === 0) return false;

    const idleOthers = new Set(
        queue.filter((item) => item.poolKey !== poolKey).map((item) => item.poolKey)
    );
    candidates.sort((a, b) => {
        const aWanted = idleOthers.has(a.poolKey) ? 1 : 0;
        const bWanted = idleOthers.has(b.poolKey) ? 1 : 0;
        return aWanted - bWanted || a.lastUsed - b.lastUsed;
    });

    const victim = candidates[0];
    dropWorker(victim);
    try { victim.worker.terminate(); } catch { /* already closed */ }
    return true;
}

/**
 * Spawns a thread for `vfsRoot` and resolves once it has reported `{ ready: true }`.
 *
 * The worker signals readiness only after its module body (and therefore every import) has
 * finished, so a thread is never handed a job while it is still booting — which is what makes
 * the reuse actually cheap.
 *
 * @param {string} vfsRoot  VFS root of the company this thread serves (the pool key).
 * @param {number} googleAuthId
 * @returns {Promise<PoolEntry>}
 */
function spawnWorker(vfsRoot, googleAuthId) {
    const entry = { worker: null, id: ++workerSeq, poolKey: vfsRoot, busy: false, lastUsed: Date.now() };

    return new Promise((resolve, reject) => {
        const bootTimer = setTimeout(() => {
            cleanup();
            try { entry.worker?.terminate(); } catch { /* already closed */ }
            reject(new Error(`Sandbox worker failed to start within ${WORKER_BOOT_TIMEOUT_MS} ms`));
        }, WORKER_BOOT_TIMEOUT_MS);

        function cleanup() {
            clearTimeout(bootTimer);
            entry.worker?.off('message', onBooted);
            entry.worker?.off('error', onError);
            entry.worker?.off('exit', onExit);
        }

        function onBooted(msg) {
            if (!msg || msg.ready !== true) return;
            cleanup();
            resolve(entry);
        }

        function onError(err) {
            cleanup();
            reject(err instanceof Error ? err : new Error(err?.message || 'Sandbox worker failed to start'));
        }

        function onExit(code) {
            cleanup();
            reject(new Error(`Sandbox worker exited before becoming ready (code ${code})`));
        }

        try {
            entry.worker = new Worker(WORKER_PATH, {
                workerData: {
                    root: ROOT_DIR,
                    bufferLimit: DEFAULT_BUFFER_LIMIT,
                    dashbackUrl: `http://127.0.0.1:${env.port}/api`,
                    vfsDir: vfsRoot,
                    googleAuthId,
                },
            });
        } catch (err) {
            cleanup();
            reject(err);
            return;
        }

        // Permanent listener: without it an error emitted while the thread is idle (no job
        // listener attached) is unhandled and takes down the whole process.
        entry.worker.on('error', (err) => {
            console.error('[sandbox] worker error:', err?.message || err);
            dropWorker(entry);
        });
        entry.worker.on('message', onBooted);
        entry.worker.on('error', onError);
        entry.worker.on('exit', onExit);
    });
}

/**
 * Returns an idle thread for `poolKey`, spawning (or evicting) to grow the pool when allowed.
 * Resolves to `null` when every permitted thread is busy — the caller then queues.
 *
 * @param {string} poolKey
 * @param {number} googleAuthId
 * @returns {Promise<PoolEntry | null>}
 */
async function acquireWorker(poolKey, googleAuthId) {
    const idle = findIdle(poolKey);
    // Reserve SYNCHRONOUSLY, before any `await`. Otherwise two concurrent callers both pass
    // this check while the thread is still unmarked, and both post to the same worker.
    if (idle) {
        idle.busy = true;
        idle.lastUsed = Date.now();
        return idle;
    }

    if (countFor(poolKey) < MAX_WORKERS_PER_COMPANY) {
        growing = growing.then(async () => {
            // Re-check inside the serialized block: another caller may have grown the pool.
            if (countFor(poolKey) >= MAX_WORKERS_PER_COMPANY) return;
            if (pool.length >= MAX_WORKERS_TOTAL && !evictIdleOtherCompany(poolKey)) return;
            try {
                pool.push(await spawnWorker(poolKey, googleAuthId));
            } catch (err) {
                // Fail the oldest queued job for this company rather than leaving it waiting on
                // a thread that will never arrive, which would deadlock the queue instead of
                // surfacing the boot failure to the caller.
                const index = queue.findIndex((item) => item.poolKey === poolKey);
                if (index === -1) throw err;
                const waiting = queue.splice(index, 1)[0];
                clearTimeout(waiting.timer);
                waiting.reject(err);
            }
        });
        try {
            await growing;
        } catch {
            // Already reported to the queued job; continue so this caller still queues.
        }
        // Reserve the freshly booted thread atomically; a concurrent waiter that wakes after us
        // finds none and queues instead.
        const after = findIdle(poolKey);
        if (after) {
            after.busy = true;
            after.lastUsed = Date.now();
            return after;
        }
    }

    return null;
}

/**
 * Hands one job to one thread and resolves with the worker's report.
 *
 * Always RESOLVES with a report-shaped object — it never rejects. The tool layer distinguishes
 * outcomes by shape (`timedOut`, `error`, or a `result` key), matching what the old one-shot
 * worker returned, so timeout/error mapping in the routes keeps working unchanged.
 *
 * @param {PoolEntry} entry
 * @param {object} job        Forwarded to the worker verbatim.
 * @param {number} timeoutMs
 * @returns {Promise<object>}
 */
function runJobOnWorker(entry, job, timeoutMs) {
    return new Promise((resolve) => {
        const jobId = `${entry.id}:${++jobSeq}`;
        let settled = false;

        const timer = setTimeout(() => {
            // Terminate first, then report. A snippet stuck in a WASM or I/O call (optimizeImage, fetch)
            // would otherwise keep burning the thread after the caller has moved on, and its
            // eventual reply would arrive against an already-settled job.
            try { entry.worker.terminate(); } catch { /* already closed */ }
            dropWorker(entry);
            finish({ result: null, stdout: '', consoleLineCount: 0, timedOut: true, timeoutMs });
        }, timeoutMs);

        function cleanup() {
            clearTimeout(timer);
            entry.worker.off('message', onMessage);
            entry.worker.off('error', onError);
            entry.worker.off('exit', onExit);
        }

        function finish(payload) {
            if (settled) return;
            settled = true;
            cleanup();
            resolve(payload);
        }

        function onMessage(msg) {
            // Ignore anything that is not THIS job's report. Under `node --watch` the dev
            // server's module-loader hook posts `watch:import` notices on the same channel, so
            // the first message is not necessarily our reply; job ids are matched, not order.
            if (!msg || msg.__sandboxReport !== true || msg.id !== jobId) return;
            finish(msg);
        }

        function onError(err) {
            dropWorker(entry);
            finish({
                result: null,
                stdout: '',
                consoleLineCount: 0,
                error: (err && err.message) ? err.message : `Sandbox worker failed: ${err?.stack || err?.name || 'unknown error'}`,
                errorName: (err && err.name) || 'Error',
                stack: err && typeof err.stack === 'string' ? err.stack : undefined,
            });
        }

        function onExit(code) {
            dropWorker(entry);
            // Reaching here means the thread died WITHOUT reporting — a hard failure, never a
            // silent empty success.
            console.error(`[sandbox] worker exited without reporting (code ${code})`);
            finish({
                result: null,
                stdout: '',
                consoleLineCount: 0,
                error: `Sandbox worker exited unexpectedly without reporting a result (code ${code}).`,
                errorName: 'WorkerExit',
            });
        }

        entry.worker.on('message', onMessage);
        entry.worker.on('error', onError);
        entry.worker.on('exit', onExit);

        try {
            entry.worker.postMessage({ ...job, id: jobId });
        } catch (err) {
            finish({
                result: null,
                stdout: '',
                consoleLineCount: 0,
                error: (err && err.message) || 'Failed to post job to sandbox worker',
                errorName: 'PostMessageError',
            });
        }
    });
}

/**
 * Starts as many queued jobs as there are idle threads for their company.
 */
function drainQueue() {
    for (;;) {
        const index = queue.findIndex((item) => findIdle(item.poolKey));
        if (index === -1) return;

        const item = queue.splice(index, 1)[0];
        clearTimeout(item.timer);
        const entry = findIdle(item.poolKey);
        entry.busy = true;
        entry.lastUsed = Date.now();

        runJobOnWorker(entry, item.job, item.timeoutMs)
            .then(item.resolve, item.reject)
            .finally(() => {
                entry.busy = false;
                drainQueue();
            });
    }
}

/**
 * Runs one sandbox job on a pooled thread, spawning a thread if the caps allow and otherwise
 * queueing until one frees up.
 *
 * @param {object} options
 * @param {string} options.poolKey      VFS root; identifies the company (and the thread).
 * @param {number} options.googleAuthId
 * @param {string} options.code         Source to run.
 * @param {number} options.timeoutMs
 * @param {string} options.dashbackToken Per-run JWT for the sandbox's `fetch()`.
 * @param {object} [options.appletInput]
 * @param {string} [options.appletFile]
 * @returns {Promise<object>} Worker report (`result` / `timedOut` / `error` shape).
 */
export async function runSandboxJob(options) {
    if (shuttingDown) {
        throw new Error('Sandbox worker pool is shutting down');
    }

    const { poolKey, googleAuthId, code, timeoutMs, dashbackToken, appletInput, appletFile } = options;
    const job = { code, timeoutMs, dashbackToken, appletInput, appletFile, poolKey, googleAuthId };

    const entry = await acquireWorker(poolKey, googleAuthId);
    if (entry) {
        try {
            return await runJobOnWorker(entry, job, timeoutMs);
        } finally {
            entry.busy = false;
            drainQueue();
        }
    }

    return new Promise((resolve, reject) => {
        // The queue wait is bounded by the same timeout budget as execution, so a caller cannot
        // wait forever behind other companies' jobs.
        const timer = setTimeout(() => {
            const index = queue.findIndex((item) => item.timer === timer);
            if (index !== -1) queue.splice(index, 1);
            resolve({
                result: null,
                stdout: '',
                consoleLineCount: 0,
                error: `Sandbox queue timed out after ${timeoutMs} ms (all workers busy).`,
                errorName: 'SandboxQueueTimeout',
            });
        }, timeoutMs);

        queue.push({ resolve, reject, timer, poolKey, job, timeoutMs });
    });
}

/**
 * Terminates every pooled thread and fails anything still queued.
 *
 * Mirrors `shutdownImageCropWorkers`: live worker threads keep the process alive after the last
 * request, so `node --watch` hangs on "Waiting for graceful termination..." without this.
 *
 * @returns {Promise<void>}
 */
export async function shutdownJavascriptSandboxWorkers() {
    shuttingDown = true;

    while (queue.length > 0) {
        const pending = queue.shift();
        clearTimeout(pending.timer);
        pending.resolve({
            result: null,
            stdout: '',
            consoleLineCount: 0,
            error: 'Sandbox worker pool is shutting down',
            errorName: 'SandboxShutdown',
        });
    }

    const entries = pool.splice(0, pool.length);
    await Promise.all(
        entries.map((entry) =>
            Promise.resolve()
                .then(() => entry.worker.terminate())
                .catch(() => { /* already closed */ })
        )
    );

    shuttingDown = false;
}