import { bn254, htfDefaultsG1, mapToG1 } from "./bn254";
import { xor } from "./utils";
import {
  createHasher,
  expand_message_xmd,
  expand_message_xof,
  hash_to_field,
} from "@noble/curves/abstract/hash-to-curve";
import { Fp, Fp12, Fp2 } from "@noble/curves/abstract/tower";
import { CHash } from "@noble/curves/abstract/utils";
import { AffinePoint } from "@noble/curves/abstract/weierstrass";
import { keccak_256 } from "@noble/hashes/sha3";
import * as asn1js from "asn1js";
import { Buffer as BufferPolyfill } from "buffer";

declare let Buffer: typeof BufferPolyfill;
globalThis.Buffer = BufferPolyfill;

export type G1 = AffinePoint<Fp>;
export type G2 = AffinePoint<Fp2>;
export type GT = Fp12;

export interface Ciphertext {
  U: G2;
  V: Uint8Array;
  W: Uint8Array;
}

// Various options used to customize the IBE scheme
export type IbeOpts = {
  hash: CHash; // hash function
  k: number; // k-bit collision resistance of hash
  expand_fn: "xmd" | "xof"; // "xmd": expand_message_xmd, "xof": expand_message_xof, see RFC9380, Section 5.3.
  dsts: DstOpts;
};

// Various DSTs used throughout the IBE scheme
export type DstOpts = {
  H1_G1: Uint8Array;
  H2: Uint8Array;
  H3: Uint8Array;
  H4: Uint8Array;
};

// Default IBE options.
export const DEFAULT_OPTS: IbeOpts = {
  hash: keccak_256,
  k: 128,
  expand_fn: "xmd",
  dsts: {
    H1_G1: Buffer.from("IBE_BN254G1_XMD:KECCAK-256_SVDW_RO_H1_"),
    H2: Buffer.from("IBE_BN254_XMD:KECCAK-256_H2_"),
    H3: Buffer.from("IBE_BN254_XMD:KECCAK-256_H3_"),
    H4: Buffer.from("IBE_BN254_XMD:KECCAK-256_H4_"),
  },
};

// Our H4 hash function can output at most 2**16 - 1 = 65535 pseudorandom bytes.
const H4_MAX_OUTPUT_LEN: number = 65535;

/*
 * Convert the identity into a point on the curve.
 */
export function get_identity_g1(identity: Uint8Array, opts: IbeOpts = DEFAULT_OPTS): G1 {
  return hash_identity_to_point_g1(identity, opts);
}

/*
 * Encryption function for IBE based on https://www.iacr.org/archive/crypto2001/21390212.pdf Section 6 / https://eprint.iacr.org/2023/189.pdf, Algorithm 1
 * with the identity on G1, and the master public key on G2.
 */
export function encrypt_towards_identity_g1(
  m: Uint8Array,
  identity: Uint8Array,
  pk_g2: G2,
  opts: IbeOpts = DEFAULT_OPTS,
): Ciphertext {
  // We can encrypt at most 2**16 - 1 = 65535 bytes with our H4 hash function.
  const n_bytes = m.length;
  if (n_bytes > H4_MAX_OUTPUT_LEN) {
    throw new Error(`cannot encrypt messages larger than our hash output: ${H4_MAX_OUTPUT_LEN} bytes.`);
  }

  // Compute the identity's public key on G1
  // 3: PK_\rho \gets e(H_1(\rho), P)
  const identity_g1 = hash_identity_to_point_g1(identity, opts);
  const identity_g1p = bn254.G1.ProjectivePoint.fromAffine(identity_g1);
  const pk_g2p = bn254.G2.ProjectivePoint.fromAffine(pk_g2);
  const pk_rho = bn254.pairing(identity_g1p, pk_g2p);

  // Sample a one-time key
  // 4: \sigma \getsr \{0,1\}^\ell
  const sigma = new Uint8Array(32);
  crypto.getRandomValues(sigma);

  // Derive an ephemeral keypair
  // 5: r \gets H_3(\sigma, M)
  const r = hash_sigma_m_to_field(sigma, m, opts);
  // 6: U \gets [r]G_2
  const u_g2 = bn254.G2.ProjectivePoint.BASE.multiply(r).toAffine();

  // Hide the one-time key
  // 7: V \gets \sigma \xor H_2((PK_\rho)^r)
  const shared_key = bn254.fields.Fp12.pow(pk_rho, r);
  const v = xor(sigma, hash_shared_key_to_bytes(shared_key, sigma.length, opts));

  // Encrypt message m using a hash-based stream cipher with key \sigma
  // 8: W \gets M \xor H_4(\sigma)
  const w = xor(m, hash_sigma_to_bytes(sigma, n_bytes, opts));

  // 9: return ciphertext
  return {
    U: u_g2,
    V: v,
    W: w,
  };
}

/*
 * Decryption function for IBE based on https://www.iacr.org/archive/crypto2001/21390212.pdf Section 6 / https://eprint.iacr.org/2023/189.pdf, Algorithm 1
 * with the identity on G1, and the master public key on G2.
 */
export function decrypt_g1(ciphertext: Ciphertext, decryption_key_g1: G1, opts: IbeOpts = DEFAULT_OPTS): Uint8Array {
  // Get the one-time decryption key
  const key = preprocess_decryption_key_g1(ciphertext, decryption_key_g1, opts);
  return decrypt_g1_with_preprocess(ciphertext, key, opts);
}

/**
 * Decryption function for IBE based on https://www.iacr.org/archive/crypto2001/21390212.pdf Section 6 / https://eprint.iacr.org/2023/189.pdf, Algorithm 1
 * with the identity on G1, and the master public key on G2.
 */
export function decrypt_g1_with_preprocess(
  ciphertext: Ciphertext,
  preprocessed_decryption_key: Uint8Array,
  opts: IbeOpts = DEFAULT_OPTS,
): Uint8Array {
  // Check well-formedness of the ciphertext
  if (ciphertext.W.length > H4_MAX_OUTPUT_LEN) {
    throw new Error(`cannot decrypt messages larger than our hash output: ${H4_MAX_OUTPUT_LEN} bytes.`);
  }
  if (ciphertext.V.length !== opts.hash.outputLen) {
    throw new Error(`cannot decrypt encryption key of invalid length != ${opts.hash.outputLen} bytes.`);
  }
  if (ciphertext.V.length !== preprocessed_decryption_key.length) {
    throw new Error(`preprocessed decryption key of invalid length`);
  }

  // \ell = min(len(w), opts.hash.outputLen)
  const ell_bytes = ciphertext.W.length;

  // Get the one-time decryption key
  // 3: \sigma' \gets V \xor H_2(e(\pi_\rho, U))
  const sigma2 = xor(ciphertext.V, preprocessed_decryption_key);

  // Decrypt the message
  // 4: M' \gets W \xor H_4(\sigma')
  const m2 = xor(ciphertext.W, hash_sigma_to_bytes(sigma2, ell_bytes, opts));

  // Derive the ephemeral keypair with the candidate \sigma'
  // 5: r \gets H_3(\sigma, M)
  const r = hash_sigma_m_to_field(sigma2, m2, opts);

  // Verify that \sigma' is consistent with the message and ephemeral public key
  // 6: if U = [r]G_2 then return M' else return \bot
  const u_g2 = bn254.G2.ProjectivePoint.BASE.multiply(r);
  if (bn254.G2.ProjectivePoint.fromAffine(ciphertext.U).equals(u_g2)) {
    return m2;
  } else {
    throw new Error("invalid proof: rP check failed");
  }
}

/**
 * Preprocess a signature by computing the hash of the pairing, i.e.,
 * H_2(e(\pi_\rho, U)).
 * @param ciphertext ciphertext to preprocess the decryption key for
 * @param decryption_key_g1 decryption key on g1 for the ciphertext
 * @param opts IBE scheme options
 * @returns preprocessed decryption key
 */
export function preprocess_decryption_key_g1(
  ciphertext: Ciphertext,
  decryption_key_g1: G1,
  opts: IbeOpts = DEFAULT_OPTS,
): Uint8Array {
  const u_g2p = bn254.G2.ProjectivePoint.fromAffine(ciphertext.U);
  u_g2p.assertValidity(); // throws an error if point is invalid

  // Derive the shared key using the decryption key and the ciphertext's ephemeral public key
  const decryption_key_g1p = bn254.G1.ProjectivePoint.fromAffine(decryption_key_g1);
  const shared_key = bn254.pairing(decryption_key_g1p, u_g2p);

  // Return the mask H_2(e(\pi_\rho, U))
  return hash_shared_key_to_bytes(shared_key, ciphertext.V.length, opts);
}

/**
 * Serialize Ciphertext to ASN.1 structure
 * Ciphertext ::= SEQUENCE {
 *    u SEQUENCE {
 *        x SEQUENCE {
 *            c0 INTEGER,
 *            c1 INTEGER
 *        },
 *        y SEQUENCE {
 *            c0 INTEGER,
 *            c1 INTEGER
 *        }
 *    },
 *    v OCTET STRING,
 *    w OCTET STRING
 * }
 */
export function serializeCiphertext(ct: Ciphertext): Uint8Array {
  const sequence = new asn1js.Sequence({
    value: [
      new asn1js.Sequence({
        value: [
          new asn1js.Sequence({
            value: [asn1js.Integer.fromBigInt(ct.U.x.c0), asn1js.Integer.fromBigInt(ct.U.x.c1)],
          }),
          new asn1js.Sequence({
            value: [asn1js.Integer.fromBigInt(ct.U.y.c0), asn1js.Integer.fromBigInt(ct.U.y.c1)],
          }),
        ],
      }),
      new asn1js.OctetString({ valueHex: ct.V }),
      new asn1js.OctetString({ valueHex: ct.W }),
    ],
  });

  return new Uint8Array(sequence.toBER());
}

export function deserializeCiphertext(ct: Uint8Array): Ciphertext {
  const schema = new asn1js.Sequence({
    name: "ciphertext",
    value: [
      new asn1js.Sequence({
        name: "U",
        value: [
          new asn1js.Sequence({
            name: "x",
            value: [new asn1js.Integer(), new asn1js.Integer()],
          }),
          new asn1js.Sequence({
            name: "y",
            value: [new asn1js.Integer(), new asn1js.Integer()],
          }),
        ],
      }),
      new asn1js.OctetString({ name: "V" }),
      new asn1js.OctetString({ name: "W" }),
    ],
  });

  // Verify the validity of the schema
  const res = asn1js.verifySchema(ct, schema);
  if (!res.verified) {
    throw new Error("invalid ciphertext");
  }

  const V = new Uint8Array(res.result["V"].valueBlock.valueHex);
  const W = new Uint8Array(res.result["W"].valueBlock.valueHex);

  function bytesToBigInt(bytes: ArrayBuffer) {
    const byteArray = Array.from(new Uint8Array(bytes));
    const hex: string = byteArray.map((e) => e.toString(16).padStart(2, "0")).join("");
    return BigInt("0x" + hex);
  }
  const x = bn254.fields.Fp2.create({
    c0: bytesToBigInt(res.result["x"].valueBlock.value[0].valueBlock.valueHex),
    c1: bytesToBigInt(res.result["x"].valueBlock.value[1].valueBlock.valueHex),
  });
  const y = bn254.fields.Fp2.create({
    c0: bytesToBigInt(res.result["y"].valueBlock.value[0].valueBlock.valueHex),
    c1: bytesToBigInt(res.result["y"].valueBlock.value[1].valueBlock.valueHex),
  });
  const U = { x, y };

  return {
    U,
    V,
    W,
  };
}

// Concrete instantiation of H_1 that outputs a point on G1
// H_1: \{0, 1\}^\ast \rightarrow G_1
function hash_identity_to_point_g1(identity: Uint8Array, opts: IbeOpts): G1 {
  const hasher = createHasher(bn254.G1.ProjectivePoint, mapToG1, {
    p: htfDefaultsG1.p,
    m: htfDefaultsG1.m,
    hash: opts.hash,
    k: opts.k,
    DST: opts.dsts.H1_G1,
    expand: opts.expand_fn,
  });
  return hasher.hashToCurve(identity).toAffine();
}

// Concrete instantiation of H_2 that outputs a uniformly random byte string of length n
// H_2: G_T \rightarrow \{0, 1\}^\ell
function hash_shared_key_to_bytes(shared_key: GT, n: number, opts: IbeOpts): Uint8Array {
  // encode shared_key as BE(shared_key.c0.c0.c0) || BE(shared_key.c0.c0.c1) || BE(shared_key.c0.c1.c0) || ...
  if (opts.expand_fn == "xmd") {
    return expand_message_xmd(bn254.fields.Fp12.toBytes(shared_key), opts.dsts.H2, n, opts.hash);
  } else {
    return expand_message_xof(bn254.fields.Fp12.toBytes(shared_key), opts.dsts.H2, n, opts.k, opts.hash);
  }
}

// Concrete instantiation of H_3 that outputs a point in Fp
// H_3: \{0, 1\}^\ell \times \{0, 1\}^\ell \rightarrow Fp
function hash_sigma_m_to_field(sigma: Uint8Array, m: Uint8Array, opts: IbeOpts): bigint {
  // input = \sigma || m
  const input = new Uint8Array(sigma.length + m.length);
  input.set(sigma);
  input.set(m, sigma.length);

  // hash_to_field(\sigma || m)
  return hash_to_field(input, 1, {
    p: htfDefaultsG1.p,
    m: htfDefaultsG1.m,
    hash: opts.hash,
    k: opts.k,
    DST: opts.dsts.H3,
    expand: opts.expand_fn,
  })[0][0];
}

// Concrete instantiation of H_4 that outputs a uniformly random byte string of length n
// H_4: \{0, 1\}^\ell \rightarrow \{0, 1\}^\ell
function hash_sigma_to_bytes(sigma: Uint8Array, n: number, opts: IbeOpts): Uint8Array {
  if (opts.expand_fn == "xmd") {
    return expand_message_xmd(sigma, opts.dsts.H4, n, opts.hash);
  } else {
    return expand_message_xof(sigma, opts.dsts.H4, n, opts.k, opts.hash);
  }
}