import { describe, test, expect } from 'bun:test';
import { readFileSync } from 'fs';
import { join } from 'path';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import {
  computeFingerprint,
  kdfRootKey,
  kdfChainKey,
  deriveInitialRootKey,
} from '../src/index.js';
import { encodeEnvelope, decodeEnvelope, encodeStreamChunk, decodeStreamChunk } from '@shade/proto';
import type { PreKeyMessage, RatchetMessage, ShadeEnvelope } from '../src/index.js';
// Imported via relative path: shade-streams depends on shade-core, so adding
// it to shade-core's dependencies would create a workspace cycle.
import {
  deriveStreamKey,
  deriveLaneKey,
  buildChunkNonce,
  buildChunkAad,
  aesGcmEncryptWithNonce,
  aesGcmDecryptWithNonce,
} from '../../shade-streams/src/index.js';

const crypto = new SubtleCryptoProvider();

const VECTORS_DIR = join(import.meta.dir, '..', '..', '..', 'test-vectors');
const EXPECTED_VECTOR_VERSION = 2;

function hex(bytes: Uint8Array): string {
  return Array.from(bytes, (b) => b.toString(16).padStart(2, '0')).join('');
}

function fromHex(str: string): Uint8Array {
  const bytes = new Uint8Array(str.length / 2);
  for (let i = 0; i < bytes.length; i++) {
    bytes[i] = parseInt(str.substring(i * 2, i * 2 + 2), 16);
  }
  return bytes;
}

function loadVectors(name: string): { version: number; vectors: any[] } {
  const parsed = JSON.parse(readFileSync(join(VECTORS_DIR, name), 'utf-8'));
  expect(parsed.version).toBe(EXPECTED_VECTOR_VERSION);
  return parsed;
}

async function aesGcmEncryptDeterministic(
  key: Uint8Array,
  nonce: Uint8Array,
  plaintext: Uint8Array,
  aad: Uint8Array,
): Promise<Uint8Array> {
  const subtle = globalThis.crypto.subtle;
  const aesKey = await subtle.importKey(
    'raw',
    key as unknown as ArrayBuffer,
    'AES-GCM',
    false,
    ['encrypt', 'decrypt'],
  );
  const out = await subtle.encrypt(
    {
      name: 'AES-GCM',
      iv: nonce as unknown as ArrayBuffer,
      additionalData: aad as unknown as ArrayBuffer,
    },
    aesKey,
    plaintext as unknown as ArrayBuffer,
  );
  return new Uint8Array(out);
}

function encodeRatchetHeader(
  dhPublicKey: Uint8Array,
  previousCounter: number,
  counter: number,
): Uint8Array {
  const buf = new Uint8Array(40);
  buf.set(dhPublicKey, 0);
  const view = new DataView(buf.buffer);
  view.setUint32(32, previousCounter, false);
  view.setUint32(36, counter, false);
  return buf;
}

describe('Cross-platform test vectors', () => {
  test('HKDF vectors match', async () => {
    const { vectors } = loadVectors('hkdf.json');
    for (const v of vectors) {
      const out = await crypto.hkdf(
        fromHex(v.ikm),
        fromHex(v.salt),
        new TextEncoder().encode(v.info),
        v.length,
      );
      expect(hex(out)).toBe(v.output);
    }
  });

  test('KDF chain vectors match', async () => {
    const { vectors } = loadVectors('kdf-chain.json');

    const rootVec = vectors[0];
    const rootResult = await kdfRootKey(
      crypto,
      fromHex(rootVec.rootKey),
      fromHex(rootVec.dhOutput),
    );
    expect(hex(rootResult.newRootKey)).toBe(rootVec.newRootKey);
    expect(hex(rootResult.chainKey)).toBe(rootVec.chainKey);

    const chainVec = vectors[1];
    const chainResult = await kdfChainKey(crypto, fromHex(chainVec.chainKey));
    expect(hex(chainResult.newChainKey)).toBe(chainVec.newChainKey);
    expect(hex(chainResult.messageKey)).toBe(chainVec.messageKey);
  });

  test('X3DH initial root key vectors match', async () => {
    const { vectors } = loadVectors('x3dh.json');
    for (const v of vectors) {
      const rootKey = await deriveInitialRootKey(
        crypto,
        v.secrets.map((s: string) => fromHex(s)),
      );
      expect(hex(rootKey)).toBe(v.rootKey);
    }
  });

  test('Fingerprint vectors match', async () => {
    const { vectors } = loadVectors('fingerprint.json');
    for (const v of vectors) {
      const fp = await computeFingerprint(crypto, fromHex(v.signingKey), fromHex(v.dhKey));
      expect(fp).toBe(v.fingerprint);
    }
  });

  test('Wire format: RatchetMessage encode + decode', () => {
    const { vectors } = loadVectors('wire-format.json');
    const v = vectors.find((x: any) => x.kind === 'ratchet');
    expect(v).toBeDefined();

    const msg: RatchetMessage = {
      dhPublicKey: fromHex(v.message.dhPublicKey),
      previousCounter: v.message.previousCounter,
      counter: v.message.counter,
      ciphertext: fromHex(v.message.ciphertext),
      nonce: fromHex(v.message.nonce),
    };
    const envelope: ShadeEnvelope = {
      type: 'ratchet',
      content: msg,
      timestamp: 0,
      senderAddress: '',
    };
    const encoded = encodeEnvelope(envelope);
    expect(hex(encoded)).toBe(v.encoded);

    const decoded = decodeEnvelope(encoded);
    expect(decoded.type).toBe('ratchet');
    const rm = decoded.content as RatchetMessage;
    expect(rm.counter).toBe(msg.counter);
    expect(hex(rm.ciphertext)).toBe(hex(msg.ciphertext));
  });

  test('Wire format: PreKeyMessage encode + decode (with and without OTPK)', () => {
    const { vectors } = loadVectors('wire-format.json');
    const preKeyVectors = vectors.filter((x: any) => x.kind === 'prekey');
    expect(preKeyVectors.length).toBeGreaterThanOrEqual(2);

    for (const v of preKeyVectors) {
      const inner: RatchetMessage = {
        dhPublicKey: fromHex(v.message.inner.dhPublicKey),
        previousCounter: v.message.inner.previousCounter,
        counter: v.message.inner.counter,
        ciphertext: fromHex(v.message.inner.ciphertext),
        nonce: fromHex(v.message.inner.nonce),
      };
      const pre: PreKeyMessage = {
        registrationId: v.message.registrationId,
        preKeyId: v.message.preKeyId === null ? undefined : v.message.preKeyId,
        signedPreKeyId: v.message.signedPreKeyId,
        ephemeralKey: fromHex(v.message.ephemeralKey),
        identityDHKey: fromHex(v.message.identityDHKey),
        message: inner,
      };
      const envelope: ShadeEnvelope = {
        type: 'prekey',
        content: pre,
        timestamp: 0,
        senderAddress: '',
      };
      const encoded = encodeEnvelope(envelope);
      expect(hex(encoded)).toBe(v.encoded);

      const decoded = decodeEnvelope(encoded);
      expect(decoded.type).toBe('prekey');
      const dm = decoded.content as PreKeyMessage;
      expect(dm.registrationId).toBe(pre.registrationId);
      expect(dm.preKeyId).toBe(pre.preKeyId);
      expect(dm.signedPreKeyId).toBe(pre.signedPreKeyId);
      expect(hex(dm.ephemeralKey)).toBe(hex(pre.ephemeralKey));
      expect(hex(dm.message.ciphertext)).toBe(hex(inner.ciphertext));
    }
  });

  test('Streams 0x11: deriveStreamKey + deriveLaneKey + nonce/AAD + chunk encrypt + wire roundtrip', async () => {
    const { vectors } = loadVectors('streams.json');

    // 1. deriveStreamKey
    const sk = vectors.find((v: any) => v.description.startsWith('deriveStreamKey'));
    expect(sk).toBeDefined();
    const streamKey = await deriveStreamKey(crypto, fromHex(sk.streamSecret), fromHex(sk.streamId));
    expect(hex(streamKey)).toBe(sk.streamKey);

    // 2. deriveLaneKey for each laneId
    const lk = vectors.find((v: any) => v.description.startsWith('deriveLaneKey'));
    expect(lk).toBeDefined();
    for (const lane of lk.lanes) {
      const k = await deriveLaneKey(crypto, fromHex(lk.streamKey), fromHex(lk.streamId), lane.laneId);
      expect(hex(k)).toBe(lane.laneKey);
    }

    // 3. buildChunkNonce
    const nv = vectors.find((v: any) => v.description.startsWith('buildChunkNonce'));
    expect(nv).toBeDefined();
    for (const n of nv.nonces) {
      const seq = BigInt(n.seq);
      const out = buildChunkNonce(n.laneId, seq);
      expect(hex(out)).toBe(n.nonce);
    }

    // 4. buildChunkAad
    const av = vectors.find((v: any) => v.description.startsWith('buildChunkAad'));
    expect(av).toBeDefined();
    for (const c of av.cases) {
      const seq = BigInt(c.seq);
      const out = buildChunkAad(fromHex(av.streamId), c.laneId, seq, c.isLast);
      expect(hex(out)).toBe(c.aad);
    }

    // 5. End-to-end chunk encrypt + decrypt
    const ev = vectors.find((v: any) => v.description.startsWith('End-to-end chunk encrypt'));
    expect(ev).toBeDefined();
    const ct = await aesGcmEncryptWithNonce(
      fromHex(ev.laneKey),
      fromHex(ev.nonce),
      fromHex(ev.plaintext),
      fromHex(ev.aad),
    );
    expect(hex(ct)).toBe(ev.ciphertext);
    const pt = await aesGcmDecryptWithNonce(
      fromHex(ev.laneKey),
      fromHex(ev.nonce),
      fromHex(ev.ciphertext),
      fromHex(ev.aad),
    );
    expect(hex(pt)).toBe(ev.plaintext);

    // 6. Wire 0x11 envelope encode/decode
    const wv = vectors.find((v: any) => v.description.startsWith('Wire 0x11'));
    expect(wv).toBeDefined();
    const encoded = encodeStreamChunk({
      streamId: fromHex(wv.streamId),
      laneId: wv.laneId,
      seq: BigInt(wv.seq),
      isLast: wv.isLast,
      nonce: fromHex(wv.nonce),
      aad: fromHex(wv.extraAad),
      ciphertext: fromHex(wv.ciphertext),
    });
    expect(hex(encoded)).toBe(wv.encoded);
    const decoded = decodeStreamChunk(encoded);
    expect(hex(decoded.streamId)).toBe(wv.streamId);
    expect(decoded.laneId).toBe(wv.laneId);
    expect(decoded.seq.toString()).toBe(wv.seq);
    expect(decoded.isLast).toBe(wv.isLast);
    expect(hex(decoded.nonce)).toBe(wv.nonce);
    expect(hex(decoded.ciphertext)).toBe(wv.ciphertext);
  });

  test('Backup v1: HKDF backupKey + AES-GCM roundtrip', async () => {
    const { vectors } = loadVectors('backup.json');

    const kv = vectors.find((v: any) => v.description.startsWith('Backup v1: HKDF'));
    expect(kv).toBeDefined();
    const backupKey = await crypto.hkdf(
      new TextEncoder().encode(kv.passphrase),
      fromHex(kv.salt),
      new TextEncoder().encode(kv.info),
      32,
    );
    expect(hex(backupKey)).toBe(kv.backupKey);

    const ev = vectors.find((v: any) => v.description.startsWith('Backup v1: AES-256-GCM'));
    expect(ev).toBeDefined();
    const ct = await aesGcmEncryptDeterministic(
      fromHex(ev.backupKey),
      fromHex(ev.nonce),
      fromHex(ev.plaintext),
      new Uint8Array(0),
    );
    expect(hex(ct)).toBe(ev.ciphertext);

    const pt = await crypto.aesGcmDecrypt(
      fromHex(ev.backupKey),
      fromHex(ev.ciphertext),
      fromHex(ev.nonce),
    );
    expect(hex(pt)).toBe(ev.plaintext);
  });

  test('Group sender-keys: header AAD + step + Ed25519 signature', async () => {
    const { vectors } = loadVectors('group.json');

    // 1. Header AAD encoding
    const hv = vectors.find((v: any) => v.description.startsWith('Sender header AAD'));
    expect(hv).toBeDefined();
    const enc = new TextEncoder();
    const gBytes = enc.encode(hv.groupId);
    const sBytes = enc.encode(hv.senderAddress);
    const aad = new Uint8Array(2 + gBytes.length + 2 + sBytes.length + 4);
    const view = new DataView(aad.buffer);
    let off = 0;
    view.setUint16(off, gBytes.length, false); off += 2;
    aad.set(gBytes, off); off += gBytes.length;
    view.setUint16(off, sBytes.length, false); off += 2;
    aad.set(sBytes, off); off += sBytes.length;
    view.setUint32(off, hv.iteration, false);
    expect(hex(aad)).toBe(hv.aad);

    // 2. Sender-key step
    const sv = vectors.find((v: any) => v.description.startsWith('Sender-key step'));
    expect(sv).toBeDefined();
    const chain = await kdfChainKey(crypto, fromHex(sv.chainKey));
    expect(hex(chain.newChainKey)).toBe(sv.newChainKey);
    expect(hex(chain.messageKey)).toBe(sv.messageKey);

    const ct = await aesGcmEncryptDeterministic(
      chain.messageKey,
      fromHex(sv.nonce),
      fromHex(sv.plaintext),
      fromHex(sv.aad),
    );
    expect(hex(ct)).toBe(sv.ciphertext);

    // Ed25519 sign(aad || ct) — verify signature is valid for the recorded keys
    const signed = new Uint8Array(fromHex(sv.aad).length + ct.length);
    signed.set(fromHex(sv.aad), 0);
    signed.set(ct, fromHex(sv.aad).length);
    const ok = await crypto.verify(
      fromHex(sv.signingPublicKey),
      signed,
      fromHex(sv.signature),
    );
    expect(ok).toBe(true);

    // Decrypt roundtrip
    const pt = await crypto.aesGcmDecrypt(
      chain.messageKey,
      fromHex(sv.ciphertext),
      fromHex(sv.nonce),
      fromHex(sv.aad),
    );
    expect(hex(pt)).toBe(sv.plaintext);
  });

  test('Storage encryption HKDF subset: storageKey + fieldKey + rowNonce', async () => {
    const { vectors } = loadVectors('storage-hkdf.json');

    const sv = vectors.find((v: any) => v.description.startsWith('Storage HKDF: storageKey'));
    expect(sv).toBeDefined();
    const storageKey = await crypto.hkdf(
      fromHex(sv.masterKey),
      new Uint8Array(0),
      new TextEncoder().encode('shade-storage-v1'),
      32,
    );
    expect(hex(storageKey)).toBe(sv.storageKey);

    const fv = vectors.find((v: any) => v.description.startsWith('Storage HKDF: fieldKey'));
    expect(fv).toBeDefined();
    for (const f of fv.fields) {
      const k = await crypto.hkdf(
        fromHex(fv.storageKey),
        new Uint8Array(0),
        new TextEncoder().encode(`shade-field-v1:${f.table}:${f.column}`),
        32,
      );
      expect(hex(k)).toBe(f.fieldKey);
    }

    const nv = vectors.find((v: any) => v.description.startsWith('Storage HKDF: rowNonce'));
    expect(nv).toBeDefined();
    for (const n of nv.nonces) {
      const out = await crypto.hkdf(
        fromHex(nv.rowKey),
        new Uint8Array(0),
        new TextEncoder().encode(`shade-row-nonce-v1:${n.table}:${n.pk}`),
        12,
      );
      expect(hex(out)).toBe(n.nonce);
    }
  });

  test('Ratchet step: deterministic encrypt + decrypt roundtrip', async () => {
    const { vectors } = loadVectors('ratchet-step.json');
    for (const v of vectors) {
      const rootKey = fromHex(v.inputs.rootKey);
      const dhSendPriv = fromHex(v.inputs.dhSendPrivateKey);
      const dhSendPub = fromHex(v.inputs.dhSendPublicKey);
      const dhRemotePub = fromHex(v.inputs.dhRemotePublicKey);
      const plaintext = fromHex(v.inputs.plaintext);
      const nonce = fromHex(v.inputs.nonce);
      const previousCounter: number = v.inputs.previousCounter;
      const counter: number = v.inputs.counter;

      // 1. DH
      const dhOutput = await crypto.x25519(dhSendPriv, dhRemotePub);
      expect(hex(dhOutput)).toBe(v.derived.dhOutput);

      // 2. kdfRootKey
      const root = await kdfRootKey(crypto, rootKey, dhOutput);
      expect(hex(root.newRootKey)).toBe(v.derived.newRootKey);
      expect(hex(root.chainKey)).toBe(v.derived.chainKey);

      // 3. kdfChainKey
      const chain = await kdfChainKey(crypto, root.chainKey);
      expect(hex(chain.newChainKey)).toBe(v.derived.newChainKey);
      expect(hex(chain.messageKey)).toBe(v.derived.messageKey);

      // 4. Header AAD
      const aad = encodeRatchetHeader(dhSendPub, previousCounter, counter);
      expect(hex(aad)).toBe(v.derived.aad);

      // 5. AES-GCM encrypt with fixed nonce
      const ciphertext = await aesGcmEncryptDeterministic(chain.messageKey, nonce, plaintext, aad);
      expect(hex(ciphertext)).toBe(v.ciphertext);

      // 6. Roundtrip decrypt — verify the recorded ciphertext recovers the plaintext
      const recovered = await crypto.aesGcmDecrypt(
        chain.messageKey,
        fromHex(v.ciphertext),
        nonce,
        aad,
      );
      expect(hex(recovered)).toBe(v.inputs.plaintext);
    }
  });
});