const assert = require('node:assert/strict');
const { once } = require('node:events');
const fs = require('node:fs');
const http = require('node:http');
const net = require('node:net');
const path = require('node:path');
const { test } = require('node:test');
const tls = require('node:tls');
const vm = require('node:vm');

const source = fs.readFileSync(path.join(__dirname, 'bridge.js'), 'utf8');

async function listen(t, server) {
    const sockets = new Set();
    server.on('connection', socket => {
        sockets.add(socket);
        socket.on('close', () => sockets.delete(socket));
    });
    t.after(() => {
        for (const socket of sockets) socket.destroy();
        server.close();
    });
    server.listen(0, '127.0.0.1');
    await once(server, 'listening');
    return server.address().port;
}

async function bridge(t, upstreamPort, upstreamCfg) {
    const servers = [];
    const env = { ...process.env, ...(upstreamCfg && upstreamCfg.env ? upstreamCfg.env : {}) };
    // Isolate routing, credentials and listeners from the user's installed service.
    vm.runInNewContext(source, {
        Buffer, URL, console: { log() {} }, setTimeout, clearTimeout,
        process: { env },
        require(name) {
            if (name === 'fs') return {
                readFileSync: () => JSON.stringify(upstreamCfg
                    ? { enabled: true, host: upstreamCfg.host, port: upstreamCfg.port }
                    : upstreamPort
                        ? { enabled: true, host: '127.0.0.1', port: upstreamPort }
                        : { enabled: false }),
                existsSync: () => false,
                watchFile() {},
            };
            if (name === 'child_process') return { execFileSync: () => '' };
            if (name === 'http') return {
                ...http,
                createServer(...args) {
                    const server = http.createServer(...args);
                    server.listen = () => { servers.push(server); return server; };
                    return server;
                },
            };
            return require(name);
        },
    }, { filename: 'bridge.js' });
    const server = servers[0];
    server.listen = net.Server.prototype.listen;
    return listen(t, server);
}

// Start the bridge with env overrides and return the ephemeral port of a chosen
// listenerProfile ('dynamic', 'internet', or a PAC profile name). All listeners
// are tracked for cleanup; only the requested one is bound to a random port.
async function bridgeOnProfile(t, listenerProfile, { upstreamCfg, env = {}, withCreds = false } = {}) {
    const servers = [];
    const fullEnv = { ...process.env, ...env };
    vm.runInNewContext(source, {
        Buffer, URL, console: { log() {} }, setTimeout, clearTimeout,
        process: { env: fullEnv },
        require(name) {
            if (name === 'fs') return {
                readFileSync: (p) => {
                    // Bridge config file -> upstreamCfg
                    if (typeof p === 'string' && p.includes('config.json')) {
                        return JSON.stringify(upstreamCfg || { enabled: false });
                    }
                    if (typeof p === 'string' && p.includes('user.json')) {
                        if (!withCreds) throw new Error('ENOENT: no user.json in test');
                        return JSON.stringify({ username: 'test-user' });
                    }
                    // PAC files and everything else fall through to the REAL fs
                    return fs.readFileSync(p, 'utf8');
                },
                existsSync: (p) => {
                    if (typeof p === 'string') {
                        if (p.includes('user.json')) return withCreds;
                        // Isolate runtime state files from the user's live installation;
                        // everything else (PAC dir etc.) uses the real fs.
                        if (p.includes('profile.json') || p.includes('overrides.json') || p.includes('remotes.json')) return false;
                    }
                    return fs.existsSync(p);
                },
                readdirSync: fs.readdirSync,
                watchFile() {},
            };
            if (name === 'child_process') return {
                execFileSync: (cmd, args) => {
                    // secret-tool lookup -> fake standing creds when withCreds, else empty
                    if (Array.isArray(args) && args[0] === 'lookup' && withCreds) return 'test-password\n';
                    return '';
                },
            };
            if (name === 'http') return {
                ...http,
                createServer(...args) {
                    const server = http.createServer(...args);
                    server.listen = () => { servers.push(server); return server; };
                    return server;
                },
            };
            return require(name);
        },
    }, { filename: 'bridge.js' });

    // Listener creation order: index 0 = dynamic; then getAvailableProfiles()
    // order = internet (if enabled), direct, off, then PAC profiles (readdir order).
    const names = ['dynamic'];
    if (!upstreamCfg || upstreamCfg.enabled !== false) names.push('internet');
    names.push('direct', 'off');
    const pacDir = env.PROXY_BRIDGE_PAC_DIR || path.join(require('os').homedir(), '.mozilla');
    if (fs.existsSync(pacDir)) {
        for (const f of fs.readdirSync(pacDir)) if (f.endsWith('.pac')) names.push(f.slice(0, -4));
    }
    const idx = names.indexOf(listenerProfile);
    if (idx === -1 || !servers[idx]) throw new Error(`listener '${listenerProfile}' not created; have ${names.join(',')}`);
    for (let i = 0; i < servers.length; i++) {
        if (i !== idx && servers[i].listening === false) { /* never actually listened */ }
    }
    const server = servers[idx];
    server.listen = net.Server.prototype.listen;
    return listen(t, server);
}

async function connect(t, port, initial = '') {
    const socket = net.connect(port, '127.0.0.1');
    t.after(() => socket.destroy());
    await once(socket, 'connect');
    socket.write(`CONNECT example.test:443 HTTP/1.1\r\nHost: example.test:443\r\n\r\n${initial}`);
    return socket;
}

function readUntil(socket, predicate) {
    return new Promise((resolve, reject) => {
        let data = Buffer.alloc(0);
        const cleanup = () => {
            socket.off('data', onData);
            socket.off('end', onEnd);
            socket.off('error', onError);
        };
        const onData = chunk => {
            data = Buffer.concat([data, chunk]);
            if (predicate(data)) { cleanup(); resolve(data); }
        };
        const onEnd = () => { cleanup(); reject(new Error(`Unexpected EOF: ${data.toString()}`)); };
        const onError = error => { cleanup(); reject(error); };
        socket.on('data', onData);
        socket.on('end', onEnd);
        socket.on('error', onError);
    });
}

test('CONNECT buffers fragmented upstream headers and preserves tunnel bytes', { timeout: 3000 }, async t => {
    const upstreamPort = await listen(t, net.createServer(socket => {
        socket.once('data', () => {
            socket.write('HT');
            setTimeout(() => socket.write('TP/1.1 200 OK\r\nX-Proxy: test\r\n'), 30);
            setTimeout(() => socket.write('\r\nserver-data'), 60);
        });
    }));
    const socket = await connect(t, await bridge(t, upstreamPort));
    const data = await readUntil(socket, data => data.includes('server-data'));
    assert.equal(data.toString(), 'HTTP/1.1 200 Connection Established\r\n\r\nserver-data');
});

test('CONNECT forwards client head and a large payload without corruption', { timeout: 5000 }, async t => {
    const upstreamPort = await listen(t, net.createServer(socket => {
        socket.once('data', () => {
            socket.write('HTTP/1.1 200 OK\r\n\r\n');
            socket.pipe(socket);
        });
    }));
    const socket = await connect(t, await bridge(t, upstreamPort), 'early-data');
    const initial = await readUntil(socket, data => data.includes('early-data'));
    assert.equal(initial.toString(), 'HTTP/1.1 200 Connection Established\r\n\r\nearly-data');
    const payload = Buffer.alloc(4 * 1024 * 1024, 'x');
    const received = readUntil(socket, data => data.length >= payload.length);
    socket.write(payload);
    assert.deepEqual(await received, payload);
});

test('upstream reset closes an established tunnel without injecting HTTP', { timeout: 3000 }, async t => {
    const upstreamPort = await listen(t, net.createServer(socket => {
        socket.once('data', () => {
            socket.write('HTTP/1.1 200 OK\r\n\r\n');
            socket.once('data', () => socket.resetAndDestroy());
        });
    }));
    const socket = await connect(t, await bridge(t, upstreamPort));
    await readUntil(socket, data => data.includes('\r\n\r\n'));
    const chunks = [];
    socket.on('data', data => chunks.push(data));
    const ended = once(socket, 'end');
    socket.write('trigger-reset');
    await ended;
    assert.equal(Buffer.concat(chunks).length, 0);
});

test('upstream EOF before CONNECT completes does not leave client hanging', { timeout: 3000 }, async t => {
    const upstreamPort = await listen(t, net.createServer(socket => {
        socket.once('data', () => socket.end('HTTP/1.1 200'));
    }));
    const socket = await connect(t, await bridge(t, upstreamPort));
    const ended = once(socket, 'end');
    const data = await readUntil(socket, data => data.includes('\r\n\r\n'));
    assert.match(data.toString(), /^HTTP\/1.1 502 /);
    await ended;
});

test('HTTP streams large request bodies and delayed SSE responses', { timeout: 5000 }, async t => {
    const payload = Buffer.alloc(4 * 1024 * 1024, 'x');
    const targetPort = await listen(t, http.createServer(async (req, res) => {
        const chunks = [];
        for await (const chunk of req) chunks.push(chunk);
        assert.deepEqual(Buffer.concat(chunks), payload);
        res.writeHead(200, { 'Content-Type': 'text/event-stream' });
        res.write('data: first\n\n');
        setTimeout(() => res.end('data: [DONE]\n\n'), 100);
    }));
    const proxyPort = await bridge(t);
    const body = await new Promise((resolve, reject) => {
        const req = http.request({ host: '127.0.0.1', port: proxyPort, method: 'POST',
            path: `http://127.0.0.1:${targetPort}/chat`, headers: { 'Content-Length': payload.length } }, async res => {
            try {
                const chunks = [];
                for await (const chunk of res) chunks.push(chunk);
                resolve(Buffer.concat(chunks).toString());
            } catch (error) { reject(error); }
        });
        req.on('error', reject);
        req.end(payload);
    });
    assert.equal(body, 'data: first\n\ndata: [DONE]\n\n');
});

test('truncated HTTP responses abort the client stream', { timeout: 3000 }, async t => {
    const targetPort = await listen(t, http.createServer((req, res) => {
        res.writeHead(200, { 'Content-Type': 'text/event-stream' });
        res.write('data: partial\n\n');
        setTimeout(() => res.destroy(), 30);
    }));
    const proxyPort = await bridge(t);
    await new Promise((resolve, reject) => {
        const req = http.get({ host: '127.0.0.1', port: proxyPort,
            path: `http://127.0.0.1:${targetPort}/chat` }, res => {
            res.resume();
            res.on('aborted', resolve);
            res.on('error', () => {});
            res.on('end', () => reject(new Error('Truncated stream appeared successful')));
        });
        req.on('error', reject);
    });
});

test('DIRECT CONNECT tunnels bytes exactly once when upstream is disabled', { timeout: 3000 }, async t => {
    // Upstream disabled -> routing degrades to DIRECT to the CONNECT target.
    const target = net.createServer(socket => socket.pipe(socket));
    const targetPort = await listen(t, target);
    const proxyPort = await bridge(t); // no upstreamPort => { enabled: false }

    const socket = net.connect(proxyPort, '127.0.0.1');
    t.after(() => socket.destroy());
    await once(socket, 'connect');
    socket.write(`CONNECT 127.0.0.1:${targetPort} HTTP/1.1\r\nHost: 127.0.0.1:${targetPort}\r\n\r\nearly-bytes`);

    const chunks = [];
    socket.on('data', d => chunks.push(d));
    // Wait for establishment, then exercise the tunnel bidirectionally.
    const banner = await readUntil(socket, d => d.includes('early-bytes'));
    const occurrences = banner.toString().split('HTTP/1.1 200 Connection Established').length - 1;
    assert.equal(occurrences, 1, `expected exactly one 200 banner, got: ${JSON.stringify(banner.toString())}`);
    assert.equal(banner.toString(), 'HTTP/1.1 200 Connection Established\r\n\r\nearly-bytes');

    socket.write('ping');
    const echo = await readUntil(socket, d => d.length >= 4 && d.subarray(-4).toString() === 'ping');
    assert.ok(Buffer.concat(chunks).toString().includes('ping'));
    assert.equal(Buffer.concat([banner, echo]).includes('ping\r\nping'), false);
});

// Throwaway self-signed fixture (CN=127.0.0.1, generated for tests only, no real identity).
// Regenerate with: openssl req -x509 -newkey rsa:2048 -keyout k.pem -out c.pem -days 3650 -nodes -subj /CN=127.0.0.1
const TLS_FIXTURE = {
    key: fs.readFileSync(path.join(__dirname, 'fixtures', 'tls-key.pem'), 'utf8'),
    cert: fs.readFileSync(path.join(__dirname, 'fixtures', 'tls-cert.pem'), 'utf8'),
};

test('HTTPS upstream TLS-wraps the CONNECT handshake', { timeout: 5000 }, async t => {
    let sawTls = false;
    let sawConnect = false;
    const upstream = tls.createServer(TLS_FIXTURE, socket => {
        sawTls = true;
        let buf = Buffer.alloc(0);
        socket.on('data', d => {
            buf = Buffer.concat([buf, d]);
            if (buf.includes('\r\n\r\n')) {
                sawConnect = buf.toString().startsWith('CONNECT 127.0.0.1:');
                socket.write('HTTP/1.1 200 OK\r\n\r\n');
                socket.pipe(socket); // echo tunnel payload back
            }
        });
    });
    const upstreamPort = await listen(t, upstream);

    // Local PAC dir containing exactly one profile whose FindProxyForURL returns
    // an HTTPS directive — exercises PAC eval -> parseProxyConfig('HTTPS ...') -> tls.connect.
    const pacDir = fs.mkdtempSync(path.join(require('os').tmpdir(), 'bridge-pac-'));
    t.after(() => fs.rmSync(pacDir, { recursive: true, force: true }));
    fs.writeFileSync(path.join(pacDir, 'testhttps.pac'),
        `function FindProxyForURL() { return "HTTPS 127.0.0.1:${upstreamPort}"; }`);

    const proxyPort = await bridgeOnProfile(t, 'testhttps', {
        upstreamCfg: { enabled: true, host: '127.0.0.1', port: upstreamPort },
        env: { PROXY_BRIDGE_PAC_DIR: pacDir },
    });

    const echoTarget = await listen(t, net.createServer(s => s.pipe(s)));
    const socket = net.connect(proxyPort, '127.0.0.1');
    t.after(() => socket.destroy());
    await once(socket, 'connect');
    socket.write(`CONNECT 127.0.0.1:${echoTarget} HTTP/1.1\r\nHost: 127.0.0.1:${echoTarget}\r\n\r\n`);

    const banner = await readUntil(socket, d => d.includes('\r\n\r\n'));
    assert.equal(banner.toString(), 'HTTP/1.1 200 Connection Established\r\n\r\n');
    socket.write('tls-tunnel-probe');
    await readUntil(socket, d => d.includes('tls-tunnel-probe'));
    assert.ok(sawTls, 'bridge must speak TLS to an HTTPS upstream');
    assert.ok(sawConnect, 'bridge must send CONNECT inside the TLS tunnel');
});

// --- Credential-failure backoff (anti-lockout) ---
// A corp proxy that answers 407 to authed CONNECT. The bridge must give up
// sending Proxy-Authorization after consecutive failures instead of hammering.

function connectWithAuthCapture(t, port) {
    // One CONNECT through the bridge; returns the client socket (ended by upstream 407).
    const socket = net.connect(port, '127.0.0.1');
    t.after(() => socket.destroy());
    return once(socket, 'connect').then(() => {
        socket.write('CONNECT example.test:443 HTTP/1.1\r\nHost: example.test:443\r\n\r\n');
        return socket;
    });
}

// Wait until the client socket has seen one full proxy response head, then settle:
// covers both graceful end and destroy paths without depending on FIN timing.
function awaitProxyResponse(socket) {
    return new Promise((resolve) => {
        let buf = Buffer.alloc(0);
        const done = () => resolve(buf);
        const onData = (d) => {
            buf = Buffer.concat([buf, d]);
            if (buf.includes('\r\n\r\n')) {
                socket.off('data', onData);
                // give the close path one tick, but don't hang on it
                setTimeout(done, 25);
            }
        };
        socket.on('data', onData);
        socket.once('end', () => setTimeout(done, 25));
        socket.once('close', done);
        setTimeout(done, 2000); // absolute cap: never wedge the suite
    });
}

test('after 3 consecutive 407s the bridge stops sending Proxy-Authorization', { timeout: 5000 }, async t => {
    const authHeadersSeen = [];
    const upstream = net.createServer(socket => {
        socket.once('data', d => {
            authHeadersSeen.push(/^Proxy-Authorization:/mi.test(d.toString()));
            socket.end('HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic realm="corp"\r\n\r\n');
        });
    });
    const upstreamPort = await listen(t, upstream);
    const proxyPort = await bridgeOnProfile(t, 'internet', {
        upstreamCfg: { enabled: true, host: '127.0.0.1', port: upstreamPort },
        withCreds: true,
    });

    for (let i = 0; i < 4; i++) {
        const socket = await connectWithAuthCapture(t, proxyPort);
        await awaitProxyResponse(socket); // upstream 407 is relayed then closed
    }
    assert.deepEqual(authHeadersSeen, [true, true, true, false],
        'attempts 1-3 carry auth (optimistic), attempt 4 must not (lockout engaged)');
});

test('a 200 response after failures resets the lockout so auth resumes', { timeout: 5000 }, async t => {
    const authHeadersSeen = [];
    let calls = 0;
    const upstream = net.createServer(socket => {
        socket.once('data', d => {
            authHeadersSeen.push(/^Proxy-Authorization:/mi.test(d.toString()));
            calls++;
            // 407 twice (below threshold), then 200 once (clears), then 407 x3 (re-locks), then observe
            if (calls <= 2 || calls > 3) {
                socket.end('HTTP/1.1 407 Proxy Authentication Required\r\n\r\n');
            } else {
                socket.write('HTTP/1.1 200 OK\r\n\r\n');
                socket.pipe(socket);
            }
        });
    });
    const upstreamPort = await listen(t, upstream);
    const proxyPort = await bridgeOnProfile(t, 'internet', {
        upstreamCfg: { enabled: true, host: '127.0.0.1', port: upstreamPort },
        withCreds: true,
    });

    // 2 fail (count=2), 1 success (clears), 3 fail (re-lock at 3), final observe
    for (let i = 0; i < 7; i++) {
        const socket = await connectWithAuthCapture(t, proxyPort);
        await awaitProxyResponse(socket);
    }
    assert.deepEqual(authHeadersSeen, [true, true, true, true, true, true, false],
        'success mid-stream resets the counter; re-lock needs a fresh 3 consecutive 407s');
});

test('missing credentials never produce an auth header and do not crash the bridge', { timeout: 5000 }, async t => {
    const authHeadersSeen = [];
    const upstream = net.createServer(socket => {
        socket.once('data', d => {
            authHeadersSeen.push(/^Proxy-Authorization:/mi.test(d.toString()));
            socket.write('HTTP/1.1 200 OK\r\n\r\n');
            socket.pipe(socket);
        });
    });
    const upstreamPort = await listen(t, upstream);
    const proxyPort = await bridgeOnProfile(t, 'internet', {
        upstreamCfg: { enabled: true, host: '127.0.0.1', port: upstreamPort },
        withCreds: false, // no user.json, empty keyring
    });

    for (let i = 0; i < 3; i++) {
        const socket = await connectWithAuthCapture(t, proxyPort);
        await readUntil(socket, d => d.includes('200'));
        socket.destroy();
    }
    assert.deepEqual(authHeadersSeen, [false, false, false]);
});