import type { IdentityKeyPair, KeyPair, SignedPreKey, OneTimePreKey, SessionState, ChainState } from './types.js';

// ─── Base64 Helpers ──────────────────────────────────────────

export function toBase64(buf: Uint8Array): string {
  return Buffer.from(buf).toString('base64');
}

export function fromBase64(str: string): Uint8Array {
  return new Uint8Array(Buffer.from(str, 'base64'));
}

// ─── KeyPair ─────────────────────────────────────────────────

interface SerializedKeyPair { publicKey: string; privateKey: string }

function serializeKeyPair(kp: KeyPair): SerializedKeyPair {
  return { publicKey: toBase64(kp.publicKey), privateKey: toBase64(kp.privateKey) };
}

function deserializeKeyPair(s: SerializedKeyPair): KeyPair {
  return { publicKey: fromBase64(s.publicKey), privateKey: fromBase64(s.privateKey) };
}

// ─── IdentityKeyPair ─────────────────────────────────────────

export function serializeIdentityKeyPair(ikp: IdentityKeyPair): string {
  return JSON.stringify({
    signingPublicKey: toBase64(ikp.signingPublicKey),
    signingPrivateKey: toBase64(ikp.signingPrivateKey),
    dhPublicKey: toBase64(ikp.dhPublicKey),
    dhPrivateKey: toBase64(ikp.dhPrivateKey),
  });
}

export function deserializeIdentityKeyPair(json: string): IdentityKeyPair {
  const o = JSON.parse(json);
  return {
    signingPublicKey: fromBase64(o.signingPublicKey),
    signingPrivateKey: fromBase64(o.signingPrivateKey),
    dhPublicKey: fromBase64(o.dhPublicKey),
    dhPrivateKey: fromBase64(o.dhPrivateKey),
  };
}

// ─── SignedPreKey ────────────────────────────────────────────

export function serializeSignedPreKey(spk: SignedPreKey): string {
  return JSON.stringify({
    keyId: spk.keyId,
    keyPair: serializeKeyPair(spk.keyPair),
    signature: toBase64(spk.signature),
    timestamp: spk.timestamp,
  });
}

export function deserializeSignedPreKey(json: string): SignedPreKey {
  const o = JSON.parse(json);
  return {
    keyId: o.keyId,
    keyPair: deserializeKeyPair(o.keyPair),
    signature: fromBase64(o.signature),
    timestamp: o.timestamp,
  };
}

// ─── OneTimePreKey ──────────────────────────────────────────

export function serializeOneTimePreKey(otpk: OneTimePreKey): string {
  return JSON.stringify({
    keyId: otpk.keyId,
    keyPair: serializeKeyPair(otpk.keyPair),
  });
}

export function deserializeOneTimePreKey(json: string): OneTimePreKey {
  const o = JSON.parse(json);
  return { keyId: o.keyId, keyPair: deserializeKeyPair(o.keyPair) };
}

// ─── ChainState ─────────────────────────────────────────────

function serializeChain(c: ChainState): { chainKey: string; counter: number } {
  return { chainKey: toBase64(c.chainKey), counter: c.counter };
}

function deserializeChain(o: { chainKey: string; counter: number }): ChainState {
  return { chainKey: fromBase64(o.chainKey), counter: o.counter };
}

// ─── SessionState ───────────────────────────────────────────

export function serializeSessionState(state: SessionState): string {
  const skipped: Record<string, string> = {};
  for (const [k, v] of state.skippedKeys) {
    skipped[k] = toBase64(v);
  }

  const obj: any = {
    remoteIdentityKey: toBase64(state.remoteIdentityKey),
    rootKey: toBase64(state.rootKey),
    sendChain: serializeChain(state.sendChain),
    receiveChain: state.receiveChain ? serializeChain(state.receiveChain) : null,
    dhSend: serializeKeyPair(state.dhSend),
    dhReceive: state.dhReceive ? toBase64(state.dhReceive) : null,
    previousSendCounter: state.previousSendCounter,
    skippedKeys: skipped,
  };

  // Preserve __x3dh metadata if present (used for first PreKeyMessage)
  const x3dh = (state as any).__x3dh;
  if (x3dh) {
    obj.__x3dh = {
      ephemeralPublicKey: toBase64(x3dh.ephemeralPublicKey),
      signedPreKeyId: x3dh.signedPreKeyId,
      preKeyId: x3dh.preKeyId,
      identityDHKey: toBase64(x3dh.identityDHKey),
      registrationId: x3dh.registrationId,
    };
  }

  return JSON.stringify(obj);
}

export function deserializeSessionState(json: string): SessionState {
  const o = JSON.parse(json);

  const skippedKeys = new Map<string, Uint8Array>();
  if (o.skippedKeys) {
    for (const [k, v] of Object.entries(o.skippedKeys)) {
      skippedKeys.set(k, fromBase64(v as string));
    }
  }

  const state: SessionState = {
    remoteIdentityKey: fromBase64(o.remoteIdentityKey),
    rootKey: fromBase64(o.rootKey),
    sendChain: deserializeChain(o.sendChain),
    receiveChain: o.receiveChain ? deserializeChain(o.receiveChain) : null,
    dhSend: deserializeKeyPair(o.dhSend),
    dhReceive: o.dhReceive ? fromBase64(o.dhReceive) : null,
    previousSendCounter: o.previousSendCounter,
    skippedKeys,
  };

  // Restore __x3dh metadata if present
  if (o.__x3dh) {
    (state as any).__x3dh = {
      ephemeralPublicKey: fromBase64(o.__x3dh.ephemeralPublicKey),
      signedPreKeyId: o.__x3dh.signedPreKeyId,
      preKeyId: o.__x3dh.preKeyId,
      identityDHKey: fromBase64(o.__x3dh.identityDHKey),
      registrationId: o.__x3dh.registrationId,
    };
  }

  return state;
}
