import { describe, test, expect } from 'bun:test';
import {
  toBase64, fromBase64,
  serializeIdentityKeyPair, deserializeIdentityKeyPair,
  serializeSignedPreKey, deserializeSignedPreKey,
  serializeOneTimePreKey, deserializeOneTimePreKey,
  serializeSessionState, deserializeSessionState,
} from '../src/serialization.js';
import type { IdentityKeyPair, SignedPreKey, OneTimePreKey, SessionState } from '../src/types.js';

function randBytes(n: number): Uint8Array {
  const buf = new Uint8Array(n);
  crypto.getRandomValues(buf);
  return buf;
}

describe('Serialization', () => {
  describe('base64', () => {
    test('roundtrip for various lengths', () => {
      for (const len of [0, 1, 12, 32, 64, 256]) {
        const buf = randBytes(len);
        expect(fromBase64(toBase64(buf))).toEqual(buf);
      }
    });
  });

  describe('IdentityKeyPair', () => {
    test('roundtrip', () => {
      const ikp: IdentityKeyPair = {
        signingPublicKey: randBytes(32),
        signingPrivateKey: randBytes(32),
        dhPublicKey: randBytes(32),
        dhPrivateKey: randBytes(32),
      };
      const json = serializeIdentityKeyPair(ikp);
      const restored = deserializeIdentityKeyPair(json);
      expect(restored.signingPublicKey).toEqual(ikp.signingPublicKey);
      expect(restored.signingPrivateKey).toEqual(ikp.signingPrivateKey);
      expect(restored.dhPublicKey).toEqual(ikp.dhPublicKey);
      expect(restored.dhPrivateKey).toEqual(ikp.dhPrivateKey);
    });
  });

  describe('SignedPreKey', () => {
    test('roundtrip', () => {
      const spk: SignedPreKey = {
        keyId: 42,
        keyPair: { publicKey: randBytes(32), privateKey: randBytes(32) },
        signature: randBytes(64),
        timestamp: Date.now(),
      };
      const restored = deserializeSignedPreKey(serializeSignedPreKey(spk));
      expect(restored.keyId).toBe(42);
      expect(restored.keyPair.publicKey).toEqual(spk.keyPair.publicKey);
      expect(restored.signature).toEqual(spk.signature);
      expect(restored.timestamp).toBe(spk.timestamp);
    });
  });

  describe('OneTimePreKey', () => {
    test('roundtrip', () => {
      const otpk: OneTimePreKey = {
        keyId: 100,
        keyPair: { publicKey: randBytes(32), privateKey: randBytes(32) },
      };
      const restored = deserializeOneTimePreKey(serializeOneTimePreKey(otpk));
      expect(restored.keyId).toBe(100);
      expect(restored.keyPair.publicKey).toEqual(otpk.keyPair.publicKey);
      expect(restored.keyPair.privateKey).toEqual(otpk.keyPair.privateKey);
    });
  });

  describe('SessionState', () => {
    test('roundtrip with all fields', () => {
      const state: SessionState = {
        remoteIdentityKey: randBytes(32),
        rootKey: randBytes(32),
        sendChain: { chainKey: randBytes(32), counter: 5 },
        receiveChain: { chainKey: randBytes(32), counter: 3 },
        dhSend: { publicKey: randBytes(32), privateKey: randBytes(32) },
        dhReceive: randBytes(32),
        previousSendCounter: 4,
        skippedKeys: new Map([
          ['abc:1', randBytes(32)],
          ['def:2', randBytes(32)],
        ]),
      };

      const restored = deserializeSessionState(serializeSessionState(state));
      expect(restored.remoteIdentityKey).toEqual(state.remoteIdentityKey);
      expect(restored.rootKey).toEqual(state.rootKey);
      expect(restored.sendChain.chainKey).toEqual(state.sendChain.chainKey);
      expect(restored.sendChain.counter).toBe(5);
      expect(restored.receiveChain!.chainKey).toEqual(state.receiveChain!.chainKey);
      expect(restored.receiveChain!.counter).toBe(3);
      expect(restored.dhSend.publicKey).toEqual(state.dhSend.publicKey);
      expect(restored.dhReceive).toEqual(state.dhReceive);
      expect(restored.previousSendCounter).toBe(4);
      expect(restored.skippedKeys.size).toBe(2);
      expect(restored.skippedKeys.get('abc:1')).toEqual(state.skippedKeys.get('abc:1'));
      expect(restored.skippedKeys.get('def:2')).toEqual(state.skippedKeys.get('def:2'));
    });

    test('roundtrip with null receiveChain and dhReceive', () => {
      const state: SessionState = {
        remoteIdentityKey: randBytes(32),
        rootKey: randBytes(32),
        sendChain: { chainKey: randBytes(32), counter: 0 },
        receiveChain: null,
        dhSend: { publicKey: randBytes(32), privateKey: randBytes(32) },
        dhReceive: null,
        previousSendCounter: 0,
        skippedKeys: new Map(),
      };

      const restored = deserializeSessionState(serializeSessionState(state));
      expect(restored.receiveChain).toBeNull();
      expect(restored.dhReceive).toBeNull();
      expect(restored.skippedKeys.size).toBe(0);
    });

    test('roundtrip preserves __x3dh metadata', () => {
      const state: SessionState = {
        remoteIdentityKey: randBytes(32),
        rootKey: randBytes(32),
        sendChain: { chainKey: randBytes(32), counter: 0 },
        receiveChain: null,
        dhSend: { publicKey: randBytes(32), privateKey: randBytes(32) },
        dhReceive: null,
        previousSendCounter: 0,
        skippedKeys: new Map(),
      };

      const ephKey = randBytes(32);
      const idKey = randBytes(32);
      (state as any).__x3dh = {
        ephemeralPublicKey: ephKey,
        signedPreKeyId: 1,
        preKeyId: 100,
        identityDHKey: idKey,
        registrationId: 42,
      };

      const restored = deserializeSessionState(serializeSessionState(state));
      const x3dh = (restored as any).__x3dh;
      expect(x3dh).toBeDefined();
      expect(x3dh.ephemeralPublicKey).toEqual(ephKey);
      expect(x3dh.signedPreKeyId).toBe(1);
      expect(x3dh.preKeyId).toBe(100);
      expect(x3dh.identityDHKey).toEqual(idKey);
      expect(x3dh.registrationId).toBe(42);
    });

    test('roundtrip without __x3dh', () => {
      const state: SessionState = {
        remoteIdentityKey: randBytes(32),
        rootKey: randBytes(32),
        sendChain: { chainKey: randBytes(32), counter: 0 },
        receiveChain: null,
        dhSend: { publicKey: randBytes(32), privateKey: randBytes(32) },
        dhReceive: null,
        previousSendCounter: 0,
        skippedKeys: new Map(),
      };

      const restored = deserializeSessionState(serializeSessionState(state));
      expect((restored as any).__x3dh).toBeUndefined();
    });
  });
});