import { ec as EC } from 'elliptic';
import { randomBytes } from 'crypto';
import BN from 'bn.js';

export interface Point {
  getY(): BN;
  getX(): BN;
  add(point: Point): Point;
  mul(n: BN | bigint | number): Point;
  neg(): Point;
}

export default class ElGamal {
  curve: EC;

  constructor(curve: EC) {
    this.curve = curve;
  }

  static toBytes(n: BN): Uint8Array {
    const byteArray = n.toString(16).padStart(2, '0');
    if (byteArray.length % 2 !== 0) {
      return Uint8Array.from(Buffer.from('0' + byteArray, 'hex'));
    }
    return Uint8Array.from(Buffer.from(byteArray, 'hex'));
  }

  encrypt(messagePoint: Point, publicKey: Point): [Point, Point] {
    const ephemeralKey = new BN(randomBytes(32).toString('hex'), 16).mod(
      this.curve.n!
    );
    const ephemeralKeyBigInt = BigInt(ephemeralKey.toString(10));
    const point1 = this.curve.g.mul(ephemeralKeyBigInt);
    const point2 = messagePoint.add(publicKey.mul(ephemeralKeyBigInt));

    return [point1, point2];
  }

  static decrypt(encrypted: [Point, Point], privateKey: BN): Point {
    const [point1, point2] = encrypted;
    const sharedSecret = point1.mul(privateKey);
    const decryptedMessage = point2.add(sharedSecret.neg());
    return decryptedMessage;
  }
}