import crypto from 'node:crypto';
import { readFileSync, existsSync } from 'node:fs';

// ==================== Minimal Protobuf Codec ====================

function encodeVarint(value) {
  const bytes = [];
  value = value >>> 0;
  do {
    let b = value & 0x7f;
    value >>>= 7;
    if (value > 0) b |= 0x80;
    bytes.push(b);
  } while (value > 0);
  return Buffer.from(bytes);
}

function encodeVarintField(fieldNumber, value) {
  const tag = encodeVarint((fieldNumber << 3) | 0);
  const val = encodeVarint(value);
  return Buffer.concat([tag, val]);
}

function encodeBytesField(fieldNumber, data) {
  if (!Buffer.isBuffer(data)) data = Buffer.from(data);
  const tag = encodeVarint((fieldNumber << 3) | 2);
  const len = encodeVarint(data.length);
  return Buffer.concat([tag, len, data]);
}

function decodeVarintAt(buf, offset) {
  let value = 0, shift = 0;
  while (offset < buf.length) {
    const b = buf[offset++];
    value |= (b & 0x7f) << shift;
    if (!(b & 0x80)) break;
    shift += 7;
    if (shift > 35) throw new Error('Varint too long');
  }
  return [value >>> 0, offset];
}

function decodeProtobuf(buf) {
  const fields = [];
  let offset = 0;
  while (offset < buf.length) {
    const [tag, off1] = decodeVarintAt(buf, offset);
    offset = off1;
    const fieldNum = tag >>> 3;
    const wireType = tag & 0x7;

    if (wireType === 0) {
      const [value, off2] = decodeVarintAt(buf, offset);
      offset = off2;
      fields.push({ field: fieldNum, type: wireType, value });
    } else if (wireType === 2) {
      const [length, off2] = decodeVarintAt(buf, offset);
      offset = off2;
      fields.push({ field: fieldNum, type: wireType, data: Buffer.from(buf.subarray(offset, offset + length)) });
      offset += length;
    } else if (wireType === 1) {
      offset += 8;
    } else if (wireType === 5) {
      offset += 4;
    } else {
      break;
    }
  }
  return fields;
}

// ==================== PSSH Box Parser ====================

export function parsePsshBox(base64Data) {
  const box = Buffer.from(base64Data, 'base64');
  let offset = 0;

  const size = box.readUInt32BE(offset); offset += 4;
  const type = box.subarray(offset, offset + 4).toString('ascii'); offset += 4;
  if (type !== 'pssh') throw new Error('Not a PSSH box');

  const version = box[offset]; offset += 4; // version (1) + flags (3)
  const systemId = box.subarray(offset, offset + 16); offset += 16;

  if (version === 1) {
    const kidCount = box.readUInt32BE(offset); offset += 4;
    offset += kidCount * 16;
  }

  const dataSize = box.readUInt32BE(offset); offset += 4;
  const initData = Buffer.from(box.subarray(offset, offset + dataSize));

  return { version, systemId: systemId.toString('hex'), initData };
}

// ==================== WVD File Parser ====================

function parseWvdFile(path) {
  const data = readFileSync(path);
  let offset = 0;

  const magic = data.subarray(0, 3).toString('ascii');
  if (magic !== 'WVD') throw new Error('Invalid WVD file');
  offset += 3;

  const version = data[offset++];
  if (version !== 2) throw new Error(`Unsupported WVD version: ${version}`);

  const deviceType = data[offset++];
  const securityLevel = data[offset++];
  const flags = data[offset++];

  const pkLen = data.readUInt16BE(offset); offset += 2;
  const privateKeyDer = Buffer.from(data.subarray(offset, offset + pkLen)); offset += pkLen;

  const cidLen = data.readUInt16BE(offset); offset += 2;
  const clientId = Buffer.from(data.subarray(offset, offset + cidLen));

  // Try PKCS8 first, fall back to PKCS1 (RSA)
  let privateKey;
  try {
    privateKey = crypto.createPrivateKey({ key: privateKeyDer, format: 'der', type: 'pkcs8' });
  } catch {
    privateKey = crypto.createPrivateKey({ key: privateKeyDer, format: 'der', type: 'pkcs1' });
  }

  return { deviceType, securityLevel, flags, privateKey, clientId };
}

// ==================== Service Certificate Parser ====================

function parseServiceCertResponse(responseBuffer) {
  // Response is a SignedMessage: type(1), msg(2), signature(3)
  const signed = decodeProtobuf(responseBuffer);
  let msgBytes = null;
  for (const f of signed) {
    if (f.field === 2 && f.type === 2) msgBytes = f.data;
  }
  if (!msgBytes) throw new Error('No msg in service cert response');

  // msg is a SignedDrmCertificate: drm_certificate(1), signature(2)
  const signedCert = decodeProtobuf(msgBytes);
  let certBytes = null;
  for (const f of signedCert) {
    if (f.field === 1 && f.type === 2) certBytes = f.data;
  }
  if (!certBytes) throw new Error('No certificate in service cert response');

  // DrmCertificate: type(1), serial_number(2), creation_time(3), public_key(4),
  //                 system_id(7), provider_id(6 or bytes field)
  const cert = decodeProtobuf(certBytes);
  let publicKeyDer = null;
  let serialNumber = null;
  let providerId = null;

  for (const f of cert) {
    if (f.field === 4 && f.type === 2) publicKeyDer = f.data;
    if (f.field === 2 && f.type === 2) serialNumber = f.data;
    if (f.field === 6 && f.type === 2) providerId = f.data;
  }

  if (!publicKeyDer) throw new Error('No public key in service certificate');

  // Try SPKI first, fall back to PKCS1
  let publicKey;
  try {
    publicKey = crypto.createPublicKey({ key: publicKeyDer, format: 'der', type: 'spki' });
  } catch {
    publicKey = crypto.createPublicKey({ key: publicKeyDer, format: 'der', type: 'pkcs1' });
  }

  return { publicKey, serialNumber, providerId };
}

// ==================== Widevine CDM ====================

export class WidevineCDM {
  constructor(wvdPath) {
    if (!existsSync(wvdPath)) throw new Error(`WVD file not found: ${wvdPath}`);
    const device = parseWvdFile(wvdPath);
    this.privateKey = device.privateKey;
    this.clientId = device.clientId;
    this.securityLevel = device.securityLevel;
    console.log(`[widevine] CDM initialized (L${device.securityLevel})`);
  }

  static get SERVICE_CERTIFICATE_CHALLENGE() {
    return Buffer.from([0x08, 0x04]);
  }

  parseServiceCertificate(responseBuffer) {
    return parseServiceCertResponse(responseBuffer);
  }

  generateChallenge(psshInitData, serviceCert) {
    // WidevinePsshContentId: pssh_data(1), license_type(2), request_id(3)
    const requestId = crypto.randomBytes(16);
    const wvPsshData = Buffer.concat([
      encodeBytesField(1, psshInitData),
      encodeVarintField(2, 1), // STREAMING
      encodeBytesField(3, requestId),
    ]);

    // ContentIdentification: widevine_pssh_data(1)
    const contentId = encodeBytesField(1, wvPsshData);

    // Build client identification (privacy mode if service cert available)
    let clientIdField;
    if (serviceCert) {
      // Privacy mode: encrypt client ID with service certificate's public key
      const privacyKey = crypto.randomBytes(16);
      const privacyIv = crypto.randomBytes(16);

      // AES-128-CBC encrypt the client ID
      const cipher = crypto.createCipheriv('aes-128-cbc', privacyKey, privacyIv);
      const encryptedClientId = Buffer.concat([cipher.update(this.clientId), cipher.final()]);

      // RSA-OAEP encrypt the AES key with service cert's public key
      const encryptedPrivacyKey = crypto.publicEncrypt(
        {
          key: serviceCert.publicKey,
          padding: crypto.constants.RSA_PKCS1_OAEP_PADDING,
          oaepHash: 'sha1',
        },
        privacyKey,
      );

      // EncryptedClientIdentification:
      //   provider_id(1), service_certificate_serial_number(2),
      //   encrypted_client_id(3), encrypted_client_id_iv(4),
      //   encrypted_privacy_key(5)
      const encClientId = Buffer.concat([
        serviceCert.providerId ? encodeBytesField(1, serviceCert.providerId) : Buffer.alloc(0),
        serviceCert.serialNumber ? encodeBytesField(2, serviceCert.serialNumber) : Buffer.alloc(0),
        encodeBytesField(3, encryptedClientId),
        encodeBytesField(4, privacyIv),
        encodeBytesField(5, encryptedPrivacyKey),
      ]);

      // LicenseRequest field 8 = encrypted_client_id
      clientIdField = encodeBytesField(8, encClientId);
    } else {
      // No privacy mode: send raw client ID
      // LicenseRequest field 1 = client_id
      clientIdField = encodeBytesField(1, this.clientId);
    }

    // LicenseRequest: content_id(2), type(3), request_time(4),
    //                 protocol_version(6), key_control_nonce(7)
    const licenseRequest = Buffer.concat([
      clientIdField,
      encodeBytesField(2, contentId),
      encodeVarintField(3, 1), // NEW
      encodeVarintField(4, Math.floor(Date.now() / 1000)),
      encodeVarintField(6, 21),
      encodeVarintField(7, crypto.randomInt(1, 2 ** 31)),
    ]);

    // Sign with RSA PKCS1v15 SHA1
    const signature = crypto.sign('sha1', licenseRequest, {
      key: this.privateKey,
      padding: crypto.constants.RSA_PKCS1_PADDING,
    });

    // SignedMessage: type(1)=LICENSE_REQUEST, msg(2), signature(3)
    return Buffer.concat([
      encodeVarintField(1, 1),
      encodeBytesField(2, licenseRequest),
      encodeBytesField(3, signature),
    ]);
  }

  parseLicenseResponse(responseBuffer) {
    const signed = decodeProtobuf(responseBuffer);

    let encSessionKey = null;
    let licenseMsg = null;

    for (const f of signed) {
      if (f.field === 4 && f.type === 2) encSessionKey = f.data;
      if (f.field === 2 && f.type === 2) licenseMsg = f.data;
    }

    if (!encSessionKey) throw new Error('No session key in license response');
    if (!licenseMsg) throw new Error('No license message in response');

    // Decrypt session key with RSA-OAEP SHA1
    const sessionKey = crypto.privateDecrypt(
      {
        key: this.privateKey,
        padding: crypto.constants.RSA_PKCS1_OAEP_PADDING,
        oaepHash: 'sha1',
      },
      encSessionKey,
    );

    // Parse License — KeyContainer is at field 3
    const license = decodeProtobuf(licenseMsg);
    const keys = [];

    for (const f of license) {
      if (f.field === 3 && f.type === 2) {
        const kc = decodeProtobuf(f.data);
        let kid = null, iv = null, encKey = null, keyType = 0;

        for (const kf of kc) {
          if (kf.field === 1 && kf.type === 2) kid = kf.data;
          if (kf.field === 2 && kf.type === 2) iv = kf.data;
          if (kf.field === 3 && kf.type === 2) encKey = kf.data;
          if (kf.field === 4 && kf.type === 0) keyType = kf.value;
        }

        if (encKey && iv) {
          const algo = sessionKey.length === 16 ? 'aes-128-cbc' : 'aes-256-cbc';
          const decipher = crypto.createDecipheriv(algo, sessionKey, iv);
          const decrypted = Buffer.concat([decipher.update(encKey), decipher.final()]);

          keys.push({
            kid: kid ? kid.toString('hex') : null,
            key: decrypted.toString('hex'),
            type: keyType, // 2 = CONTENT
          });
        }
      }
    }

    const contentKeys = keys.filter(k => k.type === 2);
    return contentKeys.length > 0 ? contentKeys : keys;
  }
}