import assert from './_assert.js';
import { CHash, Input, toBytes } from './utils.js';
import { hmac } from './hmac.js';

// HKDF (RFC 5869)
// https://soatok.blog/2021/11/17/understanding-hkdf/

/**
 * HKDF-Extract(IKM, salt) -> PRK
 * Arguments position differs from spec (IKM is first one, since it is not optional)
 * @param hash
 * @param ikm
 * @param salt
 * @returns
 */
export function extract(hash: CHash, ikm: Input, salt?: Input) {
  assert.hash(hash);
  // NOTE: some libraries treat zero-length array as 'not provided';
  // we don't, since we have undefined as 'not provided'
  // https://github.com/RustCrypto/KDFs/issues/15
  if (salt === undefined) salt = new Uint8Array(hash.outputLen); // if not provided, it is set to a string of HashLen zeros
  return hmac(hash, toBytes(salt), toBytes(ikm));
}

// HKDF-Expand(PRK, info, L) -> OKM
const HKDF_COUNTER = new Uint8Array([0]);
const EMPTY_BUFFER = new Uint8Array();

/**
 * HKDF-expand from the spec.
 * @param prk - a pseudorandom key of at least HashLen octets (usually, the output from the extract step)
 * @param info - optional context and application specific information (can be a zero-length string)
 * @param length - length of output keying material in octets
 */
export function expand(hash: CHash, prk: Input, info?: Input, length: number = 32) {
  assert.hash(hash);
  assert.number(length);
  if (length > 255 * hash.outputLen) throw new Error('Length should be <= 255*HashLen');
  const blocks = Math.ceil(length / hash.outputLen);
  if (info === undefined) info = EMPTY_BUFFER;
  // first L(ength) octets of T
  const okm = new Uint8Array(blocks * hash.outputLen);
  // Re-use HMAC instance between blocks
  const HMAC = hmac.create(hash, prk);
  const HMACTmp = HMAC._cloneInto();
  const T = new Uint8Array(HMAC.outputLen);
  for (let counter = 0; counter < blocks; counter++) {
    HKDF_COUNTER[0] = counter + 1;
    // T(0) = empty string (zero length)
    // T(N) = HMAC-Hash(PRK, T(N-1) | info | N)
    HMACTmp.update(counter === 0 ? EMPTY_BUFFER : T)
      .update(info)
      .update(HKDF_COUNTER)
      .digestInto(T);
    okm.set(T, hash.outputLen * counter);
    HMAC._cloneInto(HMACTmp);
  }
  HMAC.destroy();
  HMACTmp.destroy();
  T.fill(0);
  HKDF_COUNTER.fill(0);
  return okm.slice(0, length);
}

/**
 * HKDF (RFC 5869): extract + expand in one step.
 * @param hash - hash function that would be used (e.g. sha256)
 * @param ikm - input keying material, the initial key
 * @param salt - optional salt value (a non-secret random value)
 * @param info - optional context and application specific information
 * @param length - length of output keying material in octets
 */
export const hkdf = (
  hash: CHash,
  ikm: Input,
  salt: Input | undefined,
  info: Input | undefined,
  length: number
) => expand(hash, extract(hash, ikm, salt), info, length);
