#!/usr/bin/env node
/*
 * sabr-reload-test.mjs — validates the RELOAD_PLAYER_RESPONSE mitigation for
 * SABR-forced videos (e.g. kshobAibUr0) BEFORE wiring it into the add-on.
 *
 * The add-on harvests the web player's SABR POST (url + ustreamerConfig +
 * poToken) and replays it. For some videos the server answers the replay with
 * RELOAD_PLAYER_RESPONSE (UMP part 46) and zero media: the harvested
 * ustreamerConfig is stale and the server demands a fresh /player response.
 *
 * This harness checks whether honoring that reload fixes it:
 *   1. Replay the captured WEB SABR body verbatim → confirm reload reproduces.
 *   2. Decode the reload token (part 46 → ReloadPlaybackContext.reloadPlaybackParams.token).
 *   3. Re-fetch WEB /player with playbackContext.reloadPlaybackContext +
 *      the harvested poToken + visitorData → fresh serverAbrStreamingUrl +
 *      videoPlaybackUstreamerConfig.
 *   4. Swap field 5 (ustreamerConfig) of the captured body for the fresh one,
 *      POST to the fresh url, and report whether MEDIA (part 21) bytes return.
 *
 * Input: a fresh capture from harvest-pot-console.js, saved as JSON with
 *   { url, bodyB64, pot, visitor, videoId }.
 *
 * Usage:
 *   node tools/sabr-reload-test.mjs [tools/captured.json]
 */

import { readFile } from 'node:fs/promises';

const path = process.argv[2] || new URL('./captured.json', import.meta.url).pathname;
const cap = JSON.parse(await readFile(path, 'utf8'));
const origBody = Uint8Array.from(atob(cap.bodyB64.replace(/-/g, '+').replace(/_/g, '/')), c => c.charCodeAt(0));

const VIDEO_ID = cap.videoId || (cap.url.match(/[?&]id=([\w-]+)/) || [])[1] || '';
const UA = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36';

// ---- varint + protobuf field walk ------------------------------------------
const readVarint = (buf, pos) => { let r = 0, s = 0, b; do { b = buf[pos++]; r += (b & 0x7f) * Math.pow(2, s); s += 7; } while (b & 0x80); return [r, pos]; };

function* pbFields(buf) {
  let pos = 0;
  while (pos < buf.length) {
    let tag; [tag, pos] = readVarint(buf, pos);
    const f = tag >> 3, wt = tag & 7;
    if (wt === 0) { let v; [v, pos] = readVarint(buf, pos); yield [f, wt, v]; }
    else if (wt === 2) { let len; [len, pos] = readVarint(buf, pos); yield [f, wt, buf.subarray(pos, pos + len)]; pos += len; }
    else if (wt === 1) { yield [f, wt, buf.subarray(pos, pos + 8)]; pos += 8; }
    else if (wt === 5) { yield [f, wt, buf.subarray(pos, pos + 4)]; pos += 4; }
    else return;
  }
}

// varint encoder for rebuilding a length-delimited field
function writeVarint(arr, v) { v = BigInt(v); do { let b = Number(v & 0x7fn); v >>= 7n; if (v > 0n) b |= 0x80; arr.push(b); } while (v > 0n); }

// Replace the (single) length-delimited field `target` of a top-level protobuf
// message with `newVal` bytes, preserving every other field in original order.
function replaceField(buf, target, newVal) {
  const out = [];
  let replaced = false;
  for (const [f, wt, v] of pbFields(buf)) {
    const val = (f === target && wt === 2) ? (replaced = true, newVal) : v;
    writeVarint(out, f * 8 + wt);
    if (wt === 2) { writeVarint(out, val.length); for (const x of val) out.push(x); }
    else if (wt === 0) writeVarint(out, val);
    else if (wt === 1 || wt === 5) for (const x of val) out.push(x);
  }
  if (!replaced) { // field absent → append it
    writeVarint(out, target * 8 + 2); writeVarint(out, newVal.length); for (const x of newVal) out.push(x);
  }
  return Uint8Array.from(out);
}

// ---- UMP framing (leading-bit length varint) -------------------------------
function umpVarint(b, p) {
  const f = b[p];
  const n = f < 128 ? 1 : f < 192 ? 2 : f < 224 ? 3 : f < 240 ? 4 : 5;
  let v;
  if (n === 1) v = f;
  else if (n === 2) v = (f & 0x3f) + 64 * b[p + 1];
  else if (n === 3) v = (f & 0x1f) + 32 * (b[p + 1] + 256 * b[p + 2]);
  else if (n === 4) v = (f & 0x0f) + 16 * (b[p + 1] + 256 * (b[p + 2] + 256 * b[p + 3]));
  else v = b[p + 1] + 256 * (b[p + 2] + 256 * (b[p + 3] + 256 * b[p + 4]));
  return [v, p + n];
}
function parseUmp(buf) {
  const parts = []; let pos = 0;
  while (pos < buf.length) {
    let type, size; [type, pos] = umpVarint(buf, pos); if (pos > buf.length) break;
    [size, pos] = umpVarint(buf, pos);
    parts.push({ type, data: buf.subarray(pos, pos + size) }); pos += size;
  }
  return parts;
}

const PART = { MEDIA: 21, STREAM_PROTECTION: 58, RELOAD_PLAYER_RESPONSE: 46, SABR_REDIRECT: 43, SABR_ERROR: 44 };

// part 46 = ReloadPlaybackContext { reloadPlaybackParams=1 { token=1 (string) } }
function decodeReloadToken(data) {
  for (const [f, , v] of pbFields(data)) {
    if (f === 1 && v instanceof Uint8Array) {
      for (const [f2, , v2] of pbFields(v)) {
        if (f2 === 1 && v2 instanceof Uint8Array) return new TextDecoder().decode(v2);
      }
    }
  }
  return null;
}

function summarizeUmp(buf) {
  const parts = parseUmp(buf);
  const hist = {}; let mediaBytes = 0, protection = null, reloadToken = null, redirect = null, sabrErr = null;
  for (const { type, data } of parts) {
    hist[type] = (hist[type] || 0) + 1;
    if (type === PART.MEDIA) mediaBytes += data.length - 1; // minus headerId prefix
    else if (type === PART.STREAM_PROTECTION) { for (const [f, , v] of pbFields(data)) if (f === 1) protection = v; }
    else if (type === PART.RELOAD_PLAYER_RESPONSE) reloadToken = decodeReloadToken(data);
    else if (type === PART.SABR_REDIRECT) { for (const [f, , v] of pbFields(data)) if (f === 1) redirect = new TextDecoder().decode(v); }
    else if (type === PART.SABR_ERROR) sabrErr = Array.from(data.subarray(0, 16)).map(x => x.toString(16).padStart(2, '0')).join('');
  }
  const histStr = Object.keys(hist).map(Number).sort((a, b) => a - b).map(t => `${t}:${hist[t]}`).join(',');
  return { histStr, mediaBytes, protection, reloadToken, redirect, sabrErr };
}

async function postSabr(url, body) {
  const resp = await fetch(url, {
    method: 'POST',
    headers: { 'Referer': 'https://www.youtube.com/', 'Origin': 'https://www.youtube.com', 'User-Agent': UA },
    body
  });
  const buf = new Uint8Array(await resp.arrayBuffer());
  return { status: resp.status, buf };
}

// ---- WEB /player with reload token -----------------------------------------
async function reloadPlayer(token) {
  const body = {
    videoId: VIDEO_ID,
    racyCheckOk: true,
    contentCheckOk: true,
    playbackContext: {
      contentPlaybackContext: { signatureTimestamp: 20563, vis: 0, splay: false, lactMilliseconds: '-1' },
      reloadPlaybackContext: { reloadPlaybackParams: { token } }
    },
    serviceIntegrityDimensions: cap.pot ? { poToken: cap.pot } : undefined,
    context: {
      client: { clientName: 'WEB', clientVersion: '2.20260606.02.00', osName: 'Windows', osVersion: '10.0', hl: 'en', gl: 'US', visitorData: cap.visitor || '' },
      user: { enableSafetyMode: false, lockedSafetyMode: false },
      request: { useSsl: true, internalExperimentFlags: [] }
    }
  };
  const resp = await fetch('https://www.youtube.com/youtubei/v1/player?prettyPrint=false&alt=json', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'X-Youtube-Client-Name': '1',
      'X-Youtube-Client-Version': '2.20260606.02.00',
      'X-Goog-Visitor-Id': cap.visitor || '',
      'User-Agent': UA
    },
    body: JSON.stringify(body)
  });
  if (!resp.ok) throw new Error(`/player HTTP ${resp.status}: ${(await resp.text()).slice(0, 160)}`);
  const j = await resp.json();
  const sd = j.streamingData || {};
  const ust = j.playerConfig?.mediaCommonConfig?.mediaUstreamerRequestConfig?.videoPlaybackUstreamerConfig;
  return { status: j.playabilityStatus?.status, reason: j.playabilityStatus?.reason, url: sd.serverAbrStreamingUrl, ust };
}

// ---- run --------------------------------------------------------------------
console.log(`\n=== sabr-reload-test  video=${VIDEO_ID}  body=${origBody.length}B  pot=${cap.pot ? cap.pot.slice(0, 12) + '…' : '(none)'} ===`);

console.log('\n[1] Replay captured WEB SABR body verbatim …');
let r = await postSabr(cap.url, origBody);
let s = summarizeUmp(r.buf);
console.log(`    HTTP ${r.status}  parts={${s.histStr}}  media=${s.mediaBytes}B  sps=${s.protection ?? '-'}  reload=${s.reloadToken ? 'YES' : 'no'}  sabrErr=${s.sabrErr || '-'}`);

if (s.mediaBytes > 0 && !s.reloadToken) {
  console.log('\n=> Replay already streams media — no reload needed for this capture.');
  process.exit(0);
}
if (!s.reloadToken) {
  console.log('\n=> No reload token and no media — different failure (attestation/stale pot). Capture a fresher session.');
  process.exit(1);
}

console.log(`\n[2] Reload token: ${s.reloadToken.slice(0, 40)}… (len ${s.reloadToken.length})`);
console.log('[3] Re-fetching WEB /player with reloadPlaybackContext + harvested pot …');
let fresh;
try { fresh = await reloadPlayer(s.reloadToken); }
catch (e) { console.log(`    FAILED: ${e.message}`); process.exit(1); }
console.log(`    playabilityStatus=${fresh.status}${fresh.reason ? ' reason="' + fresh.reason + '"' : ''}  url=${fresh.url ? 'present' : 'MISSING'}  ustreamerConfig=${fresh.ust ? 'present (len ' + fresh.ust.length + ')' : 'MISSING'}`);

if (!fresh.url || !fresh.ust) {
  console.log('\n=> Reload /player did not yield fresh url+ustreamerConfig (likely needs a content-binding poToken we do not have).');
  process.exit(1);
}

console.log('[4] Swapping fresh ustreamerConfig into the body, POSTing to fresh url …');
const freshUst = Uint8Array.from(atob(fresh.ust.replace(/-/g, '+').replace(/_/g, '/')), c => c.charCodeAt(0));
const newBody = replaceField(origBody, 5, freshUst);
r = await postSabr(fresh.url, newBody);
s = summarizeUmp(r.buf);
console.log(`    HTTP ${r.status}  parts={${s.histStr}}  media=${s.mediaBytes}B  sps=${s.protection ?? '-'}  reload=${s.reloadToken ? 'YES' : 'no'}  sabrErr=${s.sabrErr || '-'}`);

console.log('\n=== VERDICT ===');
if (s.mediaBytes > 0) console.log('✓ RELOAD MITIGATION WORKS — fresh ustreamerConfig from reload /player streams media. Wire it into the add-on.');
else if (s.reloadToken) console.log('✗ Still reloading — server rejects even the fresh config (token may be single-use; try driving more iterations).');
else console.log(`✗ No media after reload (sps=${s.protection ?? '-'}) — fresh config insufficient with the harvested pot.`);
