/*
 * sabr-drive-test.mjs — end-to-end mock test of includes/sabr_downloader.js.
 *
 * Feeds driveSabr synthetic UMP responses through a mock fetchFn and verifies:
 *   1. Request bodies decode correctly via googlevideo's generated
 *      VideoPlaybackAbrRequest (oracle) — incl. the new fields:
 *      streamerContext.playbackCookie (f3), clientAbrState
 *      enabledTrackTypesBitfield (f40, audio-only), bufferedRange.timeRange (f6).
 *   2. MEDIA/MEDIA_END headerId is parsed as a UMP varint (test uses id 200,
 *      which a single-byte read would corrupt).
 *   3. Cumulative bufferedRanges grow across requests; playback cookie from
 *      NEXT_REQUEST_POLICY is echoed on the next request; rn increments.
 *   4. Out-of-order segment throws; consumed-segment replay is skipped.
 *   5. SABR_CONTEXT_UPDATE writePolicy KEEP_EXISTING(2) does not overwrite a
 *      context already held for the same type (yt-dlp + googlevideo parity).
 *   6. A 200 response with an HTML/JSON/text content-type (consent / captcha
 *      wall) is rejected as a dead session, not parsed as garbage UMP.
 *
 *   cd tools && node sabr-drive-test.mjs
 */
import { VideoPlaybackAbrRequest } from 'googlevideo/protos';
import { driveSabr } from '../includes/sabr_downloader.js';

let failures = 0;
const ok = (cond, name) => { console.log(`${cond ? 'PASS' : 'FAIL'}: ${name}`); if (!cond) failures++; };

// ---- minimal protobuf + UMP writers for synthetic server responses ---------
class W {
    constructor() { this.b = []; }
    varint(v) { do { let x = v & 0x7f; v = Math.floor(v / 128); if (v > 0) x |= 0x80; this.b.push(x); } while (v > 0); return this; }
    tag(f, wt) { return this.varint(f * 8 + wt); }
    int(f, v) { return this.tag(f, 0).varint(v); }
    bytes(f, u8) { this.tag(f, 2).varint(u8.length); this.b.push(...u8); return this; }
    msg(f, w) { return this.bytes(f, w.u8()); }
    u8() { return Uint8Array.from(this.b); }
}
function umpVarint(v) {
    if (v < 128) return [v];
    if (v < 16384) return [(v & 0x3f) | 0x80, v >> 6];
    throw new Error('test varint too big');
}
function ump(...parts) {
    const out = [];
    for (const [type, data] of parts) { out.push(...umpVarint(type), ...umpVarint(data.length), ...data); }
    return Uint8Array.from(out);
}
const fmtId = (itag, lmt) => new W().int(1, itag).int(2, lmt);
const mediaHeader = (hid, itag, lmt, { isInit = false, seq, startMs, durMs, len } = {}) => {
    const w = new W().int(1, hid).int(3, itag);
    if (isInit) w.int(8, 1); else w.int(9, seq).int(11, startMs).int(12, durMs);
    w.msg(13, fmtId(itag, lmt));
    if (len !== undefined) w.int(14, len);
    return w.u8();
};
const media = (hid, payload) => Uint8Array.from([...umpVarint(hid), ...payload]);
const mediaEnd = (hid) => Uint8Array.from(umpVarint(hid));
const formatInit = (itag, lmt, durUnits, timescale) => new W().msg(2, fmtId(itag, lmt)).int(9, durUnits).int(10, timescale).u8();
const nextReqPolicy = (cookie, backoffMs) => { const w = new W(); if (backoffMs) w.int(4, backoffMs); w.bytes(7, cookie); return w.u8(); };
// SabrContextUpdate (part 57): type(1), value(3), sendByDefault(4), writePolicy(5).
const sabrContextUpdate = (type, value, sendByDefault, writePolicy) => { const w = new W().int(1, type).bytes(3, value); if (sendByDefault) w.int(4, 1); if (writePolicy) w.int(5, writePolicy); return w.u8(); };

// ---- mock server -------------------------------------------------------------
const V = { itag: 248, lastModified: '111', xtags: undefined };
const A = { itag: 251, lastModified: '222', xtags: 'drc=1' };
const COOKIE = Uint8Array.from([9, 9, 9]);
const CTXVAL = Uint8Array.from([7, 7, 7, 7]); // SABR context blob the server hands us in resp1
const seg = (n) => Uint8Array.from(Array(50).fill(n)); // 50-byte payload per segment

const sentBodies = [], sentUrls = [];
function makeFetch(script) {
    let i = 0;
    return async (url, init) => {
        sentUrls.push(url); sentBodies.push(new Uint8Array(init.body));
        const buf = script[Math.min(i++, script.length - 1)];
        return { status: 200, arrayBuffer: async () => buf.buffer.slice(buf.byteOffset, buf.byteOffset + buf.byteLength) };
    };
}

// Response 1: FORMAT_INITs + init segments + first media segments.
// headerId 200 for the audio media segment — exercises the varint fix
// (single-byte parse reads 0xc8 and corrupts).
const resp1 = ump(
    [42, formatInit(V.itag, 111, 10000, 1000)],   // video total 10000ms
    [42, formatInit(A.itag, 222, 10000, 1000)],
    [35, nextReqPolicy(COOKIE)],
    [57, sabrContextUpdate(2, CTXVAL, true)],     // server hands us a SABR context to echo
    [20, mediaHeader(1, V.itag, 111, { isInit: true, len: 50 })], [21, media(1, seg(1))], [22, mediaEnd(1)],
    [20, mediaHeader(200, A.itag, 222, { isInit: true, len: 50 })], [21, media(200, seg(2))], [22, mediaEnd(200)],
    [20, mediaHeader(201, V.itag, 111, { seq: 1, startMs: 0, durMs: 5000, len: 50 })], [21, media(201, seg(3))], [22, mediaEnd(201)],
    [20, mediaHeader(202, A.itag, 222, { seq: 1, startMs: 0, durMs: 5000, len: 50 })], [21, media(202, seg(4))], [22, mediaEnd(202)],
);
// Response 2: replay of seg 1 (must be skipped) + seg 2 for both tracks → totals reach 10000ms → complete.
// Plus a foreign itag 888 (server-stitched ad format, SSAP) — must be ignored,
// never written to output: the totals assertions below catch it if it leaks.
const resp2 = ump(
    [20, mediaHeader(60, 888, 999, { seq: 1, startMs: 0, durMs: 6000, len: 50 })], [21, media(60, seg(9))], [22, mediaEnd(60)],
    [20, mediaHeader(1, V.itag, 111, { seq: 1, startMs: 0, durMs: 5000, len: 50 })], [21, media(1, seg(3))], [22, mediaEnd(1)],
    [20, mediaHeader(2, V.itag, 111, { seq: 2, startMs: 5000, durMs: 5000, len: 50 })], [21, media(2, seg(5))], [22, mediaEnd(2)],
    [20, mediaHeader(3, A.itag, 222, { seq: 2, startMs: 5000, durMs: 5000, len: 50 })], [21, media(3, seg(6))], [22, mediaEnd(3)],
);

const got = { video: 0, audio: 0 };
const res = await driveSabr({
    serverAbrStreamingUrl: 'https://rr1---sn-test.googlevideo.com/videoplayback?expire=99999999999&sabr=1&rn=42',
    ustreamerConfig: 'AAAA',
    poToken: 'cG90',
    clientInfo: { clientName: 1, clientVersion: '2.0' },
    video: V, audio: A,
    knownDurationMs: 10000,
    fetchFn: makeFetch([resp1, resp2]),
    onChunk: (k, c) => { got[k] += c.length; },
});

ok(res.video === 150 && res.audio === 150, `byte totals exact (video=${res.video} audio=${res.audio}; init+seg1+seg2 = 150 each)`);
ok(got.video === 150 && got.audio === 150, 'onChunk delivered all bytes (varint headerId 200/201/202 parsed)');

// rn: stale rn=42 stripped, fresh rn=1.. appended
ok(/[&?]rn=1$/.test(sentUrls[0]) && !/rn=42/.test(sentUrls[0]), `rn stripped+fresh (url1=${sentUrls[0].slice(-20)})`);
ok(/[&?]rn=2$/.test(sentUrls[1]), 'rn increments per request');

// Oracle-decode our request bodies with googlevideo generated protos
const req1 = VideoPlaybackAbrRequest.decode(sentBodies[0]);
const req2 = VideoPlaybackAbrRequest.decode(sentBodies[1]);

ok(req1.selectedFormatIds.length === 0, 'request 1: selectedFormatIds empty before FORMAT_INIT');
ok(req2.selectedFormatIds.length === 2, 'request 2: selectedFormatIds populated after FORMAT_INIT');
ok(!req1.streamerContext.playbackCookie?.length, 'request 1: no playback cookie yet');
ok(Buffer.from(req2.streamerContext.playbackCookie || []).equals(Buffer.from(COOKIE)), 'request 2: NEXT_REQUEST_POLICY cookie echoed in streamerContext.f3');
// SABR_CONTEXT_UPDATE (part 57) MUST be echoed in streamerContext.sabrContexts (f5)
// or the live server withholds all media — the root cause of the 0-byte stall.
ok((req1.streamerContext.sabrContexts || []).length === 0, 'request 1: no SABR contexts before update');
const sc = req2.streamerContext.sabrContexts || [];
ok(sc.length === 1 && sc[0].type === 2 && Buffer.from(sc[0].value).equals(Buffer.from(CTXVAL)), 'request 2: SABR_CONTEXT_UPDATE echoed in streamerContext.sabrContexts (f5)');
ok(req1.bufferedRanges.length === 0, 'request 1: no buffered ranges');
const br = req2.bufferedRanges;
ok(br.length === 2 && br.every(b => b.startSegmentIndex === 1 && b.endSegmentIndex === 1 && Number(b.durationMs) === 5000), 'request 2: cumulative buffered ranges (seg 1-1, 5000ms each)');
ok(br.every(b => b.timeRange && Number(b.timeRange.durationTicks) === 5000 && b.timeRange.timescale === 1000), 'request 2: bufferedRange.timeRange (f6) present');
ok(Number(req2.clientAbrState.playerTimeMs) === 5000, 'request 2: playerTimeMs = min(track durations) = 5000');

// ---- audio-only: enabledTrackTypesBitfield=1 -------------------------------
sentBodies.length = 0; sentUrls.length = 0;
const respA1 = ump(
    [42, formatInit(A.itag, 222, 5000, 1000)],
    [20, mediaHeader(1, A.itag, 222, { isInit: true, len: 50 })], [21, media(1, seg(1))], [22, mediaEnd(1)],
    [20, mediaHeader(2, A.itag, 222, { seq: 1, startMs: 0, durMs: 5000, len: 50 })], [21, media(2, seg(2))], [22, mediaEnd(2)],
);
const resA = await driveSabr({
    serverAbrStreamingUrl: 'https://rr1---sn-test.googlevideo.com/videoplayback?expire=99999999999&sabr=1',
    ustreamerConfig: 'AAAA', poToken: 'cG90', clientInfo: { clientName: 1, clientVersion: '2.0' },
    video: null, audio: A, knownDurationMs: 5000,
    fetchFn: makeFetch([respA1]), onChunk: () => { },
});
const reqA = VideoPlaybackAbrRequest.decode(sentBodies[0]);
ok(reqA.clientAbrState.enabledTrackTypesBitfield === 1, 'audio-only: enabledTrackTypesBitfield=1 (f40)');
ok(reqA.preferredVideoFormatIds.length === 0 && reqA.preferredAudioFormatIds.length === 1, 'audio-only: only audio preferred');
ok(resA.audio === 100 && resA.video === 0, 'audio-only: totals');

// ---- out-of-order segment must throw ----------------------------------------
const respBad = ump(
    [42, formatInit(A.itag, 222, 99000, 1000)],
    [20, mediaHeader(1, A.itag, 222, { seq: 1, startMs: 0, durMs: 1000, len: 50 })], [21, media(1, seg(1))], [22, mediaEnd(1)],
    [20, mediaHeader(2, A.itag, 222, { seq: 3, startMs: 2000, durMs: 1000, len: 50 })], [21, media(2, seg(2))], [22, mediaEnd(2)],
);
let threw = '';
try {
    await driveSabr({
        serverAbrStreamingUrl: 'https://x.googlevideo.com/videoplayback?expire=99999999999&sabr=1',
        ustreamerConfig: 'AAAA', poToken: 'cG90', clientInfo: { clientName: 1, clientVersion: '2.0' },
        video: null, audio: A, fetchFn: makeFetch([respBad]), onChunk: () => { },
    });
} catch (e) { threw = e.message; }
ok(/out of order/.test(threw), `out-of-order segment throws (${threw.slice(0, 60)})`);

// ---- expired URL fast-fails ---------------------------------------------------
threw = '';
try {
    await driveSabr({
        serverAbrStreamingUrl: 'https://x.googlevideo.com/videoplayback?expire=1000&sabr=1',
        ustreamerConfig: 'AAAA', poToken: 'cG90', clientInfo: { clientName: 1, clientVersion: '2.0' },
        video: null, audio: A, fetchFn: makeFetch([respA1]), onChunk: () => { },
    });
} catch (e) { threw = e.message; }
ok(/SABR session expired/.test(threw), `expired URL fast-fails (${threw.slice(0, 50)})`);

// ---- SABR_CONTEXT_UPDATE writePolicy KEEP_EXISTING(2) -----------------------
// Server hands context type 2 = KA (sendByDefault) in resp1, then RE-sends the
// same type with a different value KB + writePolicy=KEEP_EXISTING in resp2. The
// held KA must survive; request 3 must still echo KA, not KB.
sentBodies.length = 0; sentUrls.length = 0;
const KA = Uint8Array.from([7, 7, 7, 7]);
const KB = Uint8Array.from([8, 8, 8, 8]);
const k1 = ump(
    [42, formatInit(A.itag, 222, 15000, 1000)],
    [57, sabrContextUpdate(2, KA, true)],
    [20, mediaHeader(1, A.itag, 222, { isInit: true, len: 50 })], [21, media(1, seg(1))], [22, mediaEnd(1)],
    [20, mediaHeader(2, A.itag, 222, { seq: 1, startMs: 0, durMs: 5000, len: 50 })], [21, media(2, seg(2))], [22, mediaEnd(2)],
);
const k2 = ump(
    [57, sabrContextUpdate(2, KB, false, 2)], // writePolicy=KEEP_EXISTING — must be ignored
    [20, mediaHeader(3, A.itag, 222, { seq: 2, startMs: 5000, durMs: 5000, len: 50 })], [21, media(3, seg(3))], [22, mediaEnd(3)],
);
const k3 = ump(
    [20, mediaHeader(4, A.itag, 222, { seq: 3, startMs: 10000, durMs: 5000, len: 50 })], [21, media(4, seg(4))], [22, mediaEnd(4)],
);
const resK = await driveSabr({
    serverAbrStreamingUrl: 'https://rr1---sn-test.googlevideo.com/videoplayback?expire=99999999999&sabr=1',
    ustreamerConfig: 'AAAA', poToken: 'cG90', clientInfo: { clientName: 1, clientVersion: '2.0' },
    video: null, audio: A, knownDurationMs: 15000,
    fetchFn: makeFetch([k1, k2, k3]), onChunk: () => { },
});
const kc = (VideoPlaybackAbrRequest.decode(sentBodies[2]).streamerContext.sabrContexts) || [];
ok(kc.length === 1 && kc[0].type === 2 && Buffer.from(kc[0].value).equals(Buffer.from(KA)),
    `KEEP_EXISTING(2): re-sent context does not overwrite held one (value=${kc[0] ? Buffer.from(kc[0].value).join('') : '?'})`);
ok(resK.audio === 200, `KEEP_EXISTING test: audio fully drained (init+3 segs = 200, got ${resK.audio})`);

// ---- non-UMP 200 (consent / captcha wall) → session rejected ----------------
const htmlFetch = async (url, init) => {
    sentUrls.push(url);
    const body = new TextEncoder().encode('<!DOCTYPE html><html><body>Before you continue to YouTube</body></html>');
    return {
        status: 200,
        headers: { get: (k) => (String(k).toLowerCase() === 'content-type' ? 'text/html; charset=utf-8' : null) },
        arrayBuffer: async () => body.buffer.slice(body.byteOffset, body.byteOffset + body.byteLength),
    };
};
threw = '';
try {
    await driveSabr({
        serverAbrStreamingUrl: 'https://x.googlevideo.com/videoplayback?expire=99999999999&sabr=1',
        ustreamerConfig: 'AAAA', poToken: 'cG90', clientInfo: { clientName: 1, clientVersion: '2.0' },
        video: null, audio: A, fetchFn: htmlFetch, onChunk: () => { },
    });
} catch (e) { threw = e.message; }
ok(/SABR session rejected/.test(threw), `non-UMP HTML 200 → session rejected (${threw.slice(0, 50)})`);

// ---- server ad break (yt-dlp _check_vod_ad_wait parity) ---------------------
// Empty response with backoff + active SABR context = server withholding media
// for an ad's duration. The FULL backoff must be honored (old code capped at
// 8s), the ad context echoed throughout, and media resuming afterwards must
// complete the download.
const ADCTX = Uint8Array.from([5, 5, 5]);
const adResp = ump(
    [35, nextReqPolicy(COOKIE, 8500)],            // backoff ABOVE the old 8s cap
    [57, sabrContextUpdate(5, ADCTX, true)],      // type 5 = ad context
    [67, Uint8Array.from([1])],                   // SNACKBAR_MESSAGE, must be ignored
);
const adResp2 = ump(
    [42, formatInit(A.itag, 222, 10000, 1000)],
    [35, nextReqPolicy(COOKIE)],
    [20, mediaHeader(1, A.itag, 222, { isInit: true, len: 50 })], [21, media(1, seg(1))], [22, mediaEnd(1)],
    [20, mediaHeader(2, A.itag, 222, { seq: 1, startMs: 0, durMs: 5000, len: 50 })], [21, media(2, seg(2))], [22, mediaEnd(2)],
    [20, mediaHeader(3, A.itag, 222, { seq: 2, startMs: 5000, durMs: 5000, len: 50 })], [21, media(3, seg(3))], [22, mediaEnd(3)],
);
const adBase = sentBodies.length;
const adTimes = [];
const adFetchInner = makeFetch([adResp, adResp, adResp2]);
const resAd = await driveSabr({
    serverAbrStreamingUrl: 'https://rr1---sn-test.googlevideo.com/videoplayback?expire=99999999999&sabr=1',
    ustreamerConfig: 'AAAA', poToken: 'cG90', clientInfo: { clientName: 1, clientVersion: '2.0' },
    video: null, audio: A, knownDurationMs: 10000,
    fetchFn: (u, i) => { adTimes.push(Date.now()); return adFetchInner(u, i); }, onChunk: () => { },
});
ok(resAd.audio === 150, `ad break: media resumes after wait, download completes (got ${resAd.audio})`);
const adWaitMs = adTimes[2] - adTimes[1];
ok(adWaitMs >= 8300, `ad break: FULL server backoff honored, not 8s-capped (waited ${adWaitMs}ms >= 8300)`);
const adEcho = (VideoPlaybackAbrRequest.decode(sentBodies[adBase + 2]).streamerContext.sabrContexts) || [];
ok(adEcho.some((c) => c.type === 5), 'ad break: ad SABR context (type 5) still echoed after wait');

// Endless ad stall → retryable ad-break error, NOT non-retryable attestation.
// Small backoff keeps the test fast; classification is what's under test.
const adStallResp = ump(
    [35, nextReqPolicy(COOKIE, 30)],
    [57, sabrContextUpdate(5, ADCTX, true)],
    [67, Uint8Array.from([1])],
);
threw = '';
try {
    await driveSabr({
        serverAbrStreamingUrl: 'https://rr1---sn-test.googlevideo.com/videoplayback?expire=99999999999&sabr=1',
        ustreamerConfig: 'AAAA', poToken: 'cG90', clientInfo: { clientName: 1, clientVersion: '2.0' },
        video: null, audio: A, knownDurationMs: 10000,
        fetchFn: makeFetch([adStallResp]), onChunk: () => { },
    });
} catch (e) { threw = e.message; }
ok(/SABR stalled in server ad break/.test(threw) && !/attestation failed/.test(threw),
    `endless ad stall → retryable ad-break error, not attestation (${threw.slice(0, 60)})`);

// ---- per-format initialized_format_ids gating (yt-dlp 159ad90 + all-tracks completion) ----
// resp1 drives VIDEO to its full duration (init + 2 segs = 0..10000ms) but gives
// AUDIO only a FORMAT_INIT — no init segment, no media yet. The OLD single-gate
// code would (a) advertise BOTH formats in request 2 the instant any FORMAT_INIT
// arrived, and (b) break the loop the moment the video (gate track) hit 10000ms —
// never fetching audio at all. The hardened code lists ONLY video in request 2 and
// keeps looping until audio also drains.
sentBodies.length = 0; sentUrls.length = 0;
const gateResp1 = ump(
    [42, formatInit(V.itag, 111, 10000, 1000)],
    [42, formatInit(A.itag, 222, 10000, 1000)],
    [20, mediaHeader(1, V.itag, 111, { isInit: true, len: 50 })], [21, media(1, seg(1))], [22, mediaEnd(1)],
    [20, mediaHeader(2, V.itag, 111, { seq: 1, startMs: 0, durMs: 5000, len: 50 })], [21, media(2, seg(2))], [22, mediaEnd(2)],
    [20, mediaHeader(3, V.itag, 111, { seq: 2, startMs: 5000, durMs: 5000, len: 50 })], [21, media(3, seg(3))], [22, mediaEnd(3)],
);
const gateResp2 = ump(
    [20, mediaHeader(4, A.itag, 222, { isInit: true, len: 50 })], [21, media(4, seg(4))], [22, mediaEnd(4)],
    [20, mediaHeader(5, A.itag, 222, { seq: 1, startMs: 0, durMs: 5000, len: 50 })], [21, media(5, seg(5))], [22, mediaEnd(5)],
    [20, mediaHeader(6, A.itag, 222, { seq: 2, startMs: 5000, durMs: 5000, len: 50 })], [21, media(6, seg(6))], [22, mediaEnd(6)],
);
const gGot = { video: 0, audio: 0 };
const gRes = await driveSabr({
    serverAbrStreamingUrl: 'https://rr1---sn-test.googlevideo.com/videoplayback?expire=99999999999&sabr=1',
    ustreamerConfig: 'AAAA', poToken: 'cG90', clientInfo: { clientName: 1, clientVersion: '2.0' },
    video: V, audio: A, knownDurationMs: 10000,
    fetchFn: makeFetch([gateResp1, gateResp2]),
    onChunk: (k, c) => { gGot[k] += c.length; },
});
const gReq2 = VideoPlaybackAbrRequest.decode(sentBodies[1]);
ok(gReq2.selectedFormatIds.length === 1 && Number(gReq2.selectedFormatIds[0].itag) === V.itag,
   'per-format gate: request 2 advertises only the initialized video format, not the un-inited audio (yt-dlp 159ad90)');
ok(sentBodies.length >= 2 && gGot.audio === 150,
   'all-tracks completion: loop kept fetching until audio drained (audio=150, not truncated at video-complete)');

console.log(failures ? `\n${failures} FAILURE(S)` : '\nALL PASS');
process.exit(failures ? 1 : 0);
