/**
 * Encrypting JSON Web Encryption (JWE) in Flattened JSON Serialization
 *
 * @module
 */

import { encode as b64u } from '../../util/base64url.js'
import { unprotected, assertNotSet } from '../../lib/helpers.js'
import { encrypt } from '../../lib/content_encryption.js'
import type * as types from '../../types.d.ts'
import { encryptKeyManagement } from '../../lib/key_management.js'
import { JOSENotSupported, JWEInvalid } from '../../util/errors.js'
import { isDisjoint } from '../../lib/type_checks.js'
import { concat, encode } from '../../lib/buffer_utils.js'
import { validateCrit } from '../../lib/validate_crit.js'
import { normalizeKey } from '../../lib/normalize_key.js'
import { checkKeyType } from '../../lib/check_key_type.js'
import { compress } from '../../lib/deflate.js'

/**
 * The FlattenedEncrypt class is used to build and encrypt Flattened JWE objects.
 *
 * This class is exported (as a named export) from the main `'jose'` module entry point as well as
 * from its subpath export `'jose/jwe/flattened/encrypt'`.
 *
 * @example
 *
 * ```js
 * const jwe = await new jose.FlattenedEncrypt(
 *   new TextEncoder().encode('It’s a dangerous business, Frodo, going out your door.'),
 * )
 *   .setProtectedHeader({ alg: 'RSA-OAEP-256', enc: 'A256GCM' })
 *   .setAdditionalAuthenticatedData(encoder.encode('The Fellowship of the Ring'))
 *   .encrypt(publicKey)
 *
 * console.log(jwe)
 * ```
 */
export class FlattenedEncrypt {
  #plaintext: Uint8Array

  #protectedHeader!: types.JWEHeaderParameters | undefined

  #sharedUnprotectedHeader!: types.JWEHeaderParameters | undefined

  #unprotectedHeader!: types.JWEHeaderParameters | undefined

  #aad!: Uint8Array | undefined

  #cek!: Uint8Array | undefined

  #iv!: Uint8Array | undefined

  #keyManagementParameters?: types.JWEKeyManagementHeaderParameters

  /**
   * {@link FlattenedEncrypt} constructor
   *
   * @param plaintext Binary representation of the plaintext to encrypt.
   */
  constructor(plaintext: Uint8Array) {
    if (!(plaintext instanceof Uint8Array)) {
      throw new TypeError('plaintext must be an instance of Uint8Array')
    }
    this.#plaintext = plaintext
  }

  /**
   * Sets the JWE Key Management parameters to be used when encrypting.
   *
   * (ECDH-ES) Use of this method is needed for ECDH based algorithms to set the "apu" (Agreement
   * PartyUInfo) or "apv" (Agreement PartyVInfo) parameters.
   *
   * @param parameters JWE Key Management parameters.
   */
  setKeyManagementParameters(parameters: types.JWEKeyManagementHeaderParameters): this {
    assertNotSet(this.#keyManagementParameters, 'setKeyManagementParameters')
    this.#keyManagementParameters = parameters
    return this
  }

  /**
   * Sets the JWE Protected Header on the FlattenedEncrypt object.
   *
   * @param protectedHeader JWE Protected Header.
   */
  setProtectedHeader(protectedHeader: types.JWEHeaderParameters): this {
    assertNotSet(this.#protectedHeader, 'setProtectedHeader')
    this.#protectedHeader = protectedHeader
    return this
  }

  /**
   * Sets the JWE Shared Unprotected Header on the FlattenedEncrypt object.
   *
   * @param sharedUnprotectedHeader JWE Shared Unprotected Header.
   */
  setSharedUnprotectedHeader(sharedUnprotectedHeader: types.JWEHeaderParameters): this {
    assertNotSet(this.#sharedUnprotectedHeader, 'setSharedUnprotectedHeader')
    this.#sharedUnprotectedHeader = sharedUnprotectedHeader
    return this
  }

  /**
   * Sets the JWE Per-Recipient Unprotected Header on the FlattenedEncrypt object.
   *
   * @param unprotectedHeader JWE Per-Recipient Unprotected Header.
   */
  setUnprotectedHeader(unprotectedHeader: types.JWEHeaderParameters): this {
    assertNotSet(this.#unprotectedHeader, 'setUnprotectedHeader')
    this.#unprotectedHeader = unprotectedHeader
    return this
  }

  /**
   * Sets the Additional Authenticated Data on the FlattenedEncrypt object.
   *
   * @param aad Additional Authenticated Data.
   */
  setAdditionalAuthenticatedData(aad: Uint8Array): this {
    this.#aad = aad
    return this
  }

  /**
   * Sets a content encryption key to use, by default a random suitable one is generated for the JWE
   * enc" (Encryption Algorithm) Header Parameter.
   *
   * @deprecated You should not use this method. It is only really intended for test and vector
   *   validation purposes.
   *
   * @param cek JWE Content Encryption Key.
   */
  setContentEncryptionKey(cek: Uint8Array): this {
    assertNotSet(this.#cek, 'setContentEncryptionKey')
    this.#cek = cek
    return this
  }

  /**
   * Sets the JWE Initialization Vector to use for content encryption, by default a random suitable
   * one is generated for the JWE enc" (Encryption Algorithm) Header Parameter.
   *
   * @deprecated You should not use this method. It is only really intended for test and vector
   *   validation purposes.
   *
   * @param iv JWE Initialization Vector.
   */
  setInitializationVector(iv: Uint8Array): this {
    assertNotSet(this.#iv, 'setInitializationVector')
    this.#iv = iv
    return this
  }

  /**
   * Encrypts and resolves the value of the Flattened JWE object.
   *
   * @param key Public Key or Secret to encrypt the JWE with. See
   *   {@link https://github.com/panva/jose/issues/210#jwe-alg Algorithm Key Requirements}.
   * @param options JWE Encryption options.
   */
  async encrypt(
    key: types.CryptoKey | types.KeyObject | types.JWK | Uint8Array,
    options?: types.EncryptOptions,
  ): Promise<types.FlattenedJWE> {
    if (!this.#protectedHeader && !this.#unprotectedHeader && !this.#sharedUnprotectedHeader) {
      throw new JWEInvalid(
        'either setProtectedHeader, setUnprotectedHeader, or sharedUnprotectedHeader must be called before #encrypt()',
      )
    }

    if (
      !isDisjoint(this.#protectedHeader, this.#unprotectedHeader, this.#sharedUnprotectedHeader)
    ) {
      throw new JWEInvalid(
        'JWE Protected, JWE Shared Unprotected and JWE Per-Recipient Header Parameter names must be disjoint',
      )
    }

    const joseHeader: types.JWEHeaderParameters = {
      ...this.#protectedHeader,
      ...this.#unprotectedHeader,
      ...this.#sharedUnprotectedHeader,
    }

    validateCrit(JWEInvalid, new Map(), options?.crit, this.#protectedHeader, joseHeader)

    if (joseHeader.zip !== undefined && joseHeader.zip !== 'DEF') {
      throw new JOSENotSupported(
        'Unsupported JWE "zip" (Compression Algorithm) Header Parameter value.',
      )
    }

    if (joseHeader.zip !== undefined && !this.#protectedHeader?.zip) {
      throw new JWEInvalid(
        'JWE "zip" (Compression Algorithm) Header Parameter MUST be in a protected header.',
      )
    }

    const { alg, enc } = joseHeader

    if (typeof alg !== 'string' || !alg) {
      throw new JWEInvalid('JWE "alg" (Algorithm) Header Parameter missing or invalid')
    }

    if (typeof enc !== 'string' || !enc) {
      throw new JWEInvalid('JWE "enc" (Encryption Algorithm) Header Parameter missing or invalid')
    }

    let encryptedKey: Uint8Array | undefined

    if (this.#cek && (alg === 'dir' || alg === 'ECDH-ES')) {
      throw new TypeError(
        `setContentEncryptionKey cannot be called with JWE "alg" (Algorithm) Header ${alg}`,
      )
    }

    checkKeyType(alg === 'dir' ? enc : alg, key, 'encrypt')

    let cek: types.CryptoKey | Uint8Array
    {
      let parameters: { [propName: string]: unknown } | undefined
      const k = await normalizeKey(key, alg)
      ;({ cek, encryptedKey, parameters } = await encryptKeyManagement(
        alg,
        enc,
        k,
        this.#cek,
        this.#keyManagementParameters,
      ))

      if (parameters) {
        if (options && unprotected in options) {
          if (!this.#unprotectedHeader) {
            this.setUnprotectedHeader(parameters)
          } else {
            this.#unprotectedHeader = { ...this.#unprotectedHeader, ...parameters }
          }
        } else if (!this.#protectedHeader) {
          this.setProtectedHeader(parameters)
        } else {
          this.#protectedHeader = { ...this.#protectedHeader, ...parameters }
        }
      }
    }

    let additionalData: Uint8Array
    let protectedHeaderS: string
    let protectedHeaderB: Uint8Array
    let aadMember: string | undefined
    if (this.#protectedHeader) {
      protectedHeaderS = b64u(JSON.stringify(this.#protectedHeader))
      protectedHeaderB = encode(protectedHeaderS)
    } else {
      protectedHeaderS = ''
      protectedHeaderB = new Uint8Array()
    }

    if (this.#aad) {
      aadMember = b64u(this.#aad)
      const aadMemberBytes = encode(aadMember)
      additionalData = concat(protectedHeaderB, encode('.'), aadMemberBytes)
    } else {
      additionalData = protectedHeaderB
    }

    let plaintext = this.#plaintext
    if (joseHeader.zip === 'DEF') {
      plaintext = await compress(plaintext).catch((cause) => {
        throw new JWEInvalid('Failed to compress plaintext', { cause })
      })
    }

    const { ciphertext, tag, iv } = await encrypt(enc, plaintext, cek, this.#iv, additionalData)

    const jwe: types.FlattenedJWE = {
      ciphertext: b64u(ciphertext),
    }

    if (iv) {
      jwe.iv = b64u(iv)
    }

    if (tag) {
      jwe.tag = b64u(tag)
    }

    if (encryptedKey) {
      jwe.encrypted_key = b64u(encryptedKey)
    }

    if (aadMember) {
      jwe.aad = aadMember
    }

    if (this.#protectedHeader) {
      jwe.protected = protectedHeaderS
    }

    if (this.#sharedUnprotectedHeader) {
      jwe.unprotected = this.#sharedUnprotectedHeader
    }

    if (this.#unprotectedHeader) {
      jwe.header = this.#unprotectedHeader
    }

    return jwe
  }
}
