ywt_core

//// <style>
////   .goto-top {
////     display: block;
////     font-size: 0.85em;
////     text-align: right;
////   }
//// </style>
//// <script>
//// (callback => document.readyState !== 'loading' ? callback() : document.addEventListener('DOMContentLoaded', callback, { once: true }))(() => {
////   const goToTop = document.createElement('a')
////   goToTop.classList.add('goto-top')
////   goToTop.setAttribute('href', '#')
////   goToTop.textContent = 'Back to top ↑'
////   for (const member of document.querySelectorAll('.member')) {
////     member.insertBefore(goToTop.cloneNode(true), null)
////   }
//// })
//// </script>

import bigi.{type BigInt}
import gleam/bit_array
import gleam/dynamic/decode.{type Decoder}
import gleam/json
import gleam/option.{type Option, None, Some}
import ywt/algorithm.{type Algorithm}
import ywt/internal/core
import ywt/internal/jose_json

/// A key that can create JWT signatures.
///
/// Treat signing keys as secrets. Anyone with a signing key can create tokens
/// that your verifiers may trust.
pub opaque type SignKey {
  SignEcdsa(
    id: Option(String),
    curve: core.NamedCurve,
    digest_type: core.DigestType,
    public_key: BitArray,
    private_key: BitArray,
  )
  SignRsaSimple(
    id: Option(String),
    digest_type: core.DigestType,
    public_exponent: BigInt,
    modulus: BigInt,
    private_exponent: BigInt,
    padding: core.Padding,
  )
  SignRsaFull(
    id: Option(String),
    digest_type: core.DigestType,
    public_exponent: BigInt,
    modulus: BigInt,
    private_exponent: BigInt,
    first_prime_factor: BigInt,
    second_prime_factor: BigInt,
    first_factor_crt_exponent: BigInt,
    second_factor_crt_exponent: BigInt,
    first_crt_coefficient: BigInt,
    other_primes_info: List(#(BigInt, BigInt, BigInt)),
    padding: core.Padding,
  )
  SignHmac(id: Option(String), digest_type: core.DigestType, secret: BitArray)
}

/// Creates an HMAC-SHA256 signing key.
///
/// The secret must be at least 256 bits. HMAC keys are symmetric, so every
/// verifier that receives this key can also sign tokens. The length check does
/// not prove the bytes are random; use `ywt.generate_key` or a CSPRNG.
///
/// ```gleam
/// let assert Ok(key) = sign_key.hs256(secret_32_bytes)
/// ```
pub fn hs256(secret: BitArray) -> Result(SignKey, Nil) {
  // A key of the same size as the hash output (for instance, 256 bits for
  // "HS256") or larger MUST be used with this algorithm. (JWA 3.2)
  hmac_key(None, core.Sha256, secret)
}

/// Creates an HMAC-SHA384 signing key.
///
/// The secret must be at least 384 bits. HMAC keys are symmetric, so every
/// verifier that receives this key can also sign tokens. The length check does
/// not prove the bytes are random; use `ywt.generate_key` or a CSPRNG.
///
/// ```gleam
/// let assert Ok(key) = sign_key.hs384(secret_48_bytes)
/// ```
pub fn hs384(secret: BitArray) -> Result(SignKey, Nil) {
  hmac_key(None, core.Sha384, secret)
}

/// Creates an HMAC-SHA512 signing key.
///
/// The secret must be at least 512 bits. HMAC keys are symmetric, so every
/// verifier that receives this key can also sign tokens. The length check does
/// not prove the bytes are random; use `ywt.generate_key` or a CSPRNG.
///
/// ```gleam
/// let assert Ok(key) = sign_key.hs512(secret_64_bytes)
/// ```
pub fn hs512(secret: BitArray) -> Result(SignKey, Nil) {
  hmac_key(None, core.Sha512, secret)
}

fn hmac_key(
  id: Option(String),
  digest_type: core.DigestType,
  secret: BitArray,
) -> Result(SignKey, Nil) {
  let minimum_bits = case digest_type {
    core.Sha256 -> 256
    core.Sha384 -> 384
    core.Sha512 -> 512
  }

  case bit_array.bit_size(secret) >= minimum_bits {
    True -> Ok(SignHmac(id, digest_type, secret))
    False -> Error(Nil)
  }
}

/// Assigns a random key id to a signing key.
///
/// The id is written to JWT headers as `kid`, letting verifiers choose the right
/// key during rotation.
///
/// ```gleam
/// let key = sign_key.with_random_id(key)
/// ```
pub fn with_random_id(key: SignKey) -> SignKey {
  let id = Some(get_random_id())
  case key {
    SignEcdsa(..) -> SignEcdsa(..key, id:)
    SignHmac(..) -> SignHmac(..key, id:)
    SignRsaFull(..) -> SignRsaFull(..key, id:)
    SignRsaSimple(..) -> SignRsaSimple(..key, id:)
  }
}

@external(erlang, "ywt_core_ffi", "random_id")
@external(javascript, "../ywt_core_ffi.mjs", "random_id")
fn get_random_id() -> String

/// Decodes a private JWK into a signing key.
///
/// Only decode private keys from trusted storage. Unknown JWK fields are ignored;
/// `key_ops` and `use` are not enforced. HMAC and RSA signing JWKs must include
/// `alg`; EC signing JWKs may omit `alg` when the curve determines it.
/// For raw JSON strings, prefer `parse_jwk` so duplicate member handling is
/// consistent on all targets.
///
/// ```gleam
/// json.parse(private_jwk, sign_key.decoder())
/// ```
pub fn decoder() -> Decoder(SignKey) {
  // TODO: do we want a PEM parser as well?

  use id <- decode.optional_field("kid", None, decode.map(decode.string, Some))
  use kty <- decode.field("kty", decode.string)
  case kty {
    // TODO: what about octet keys without an explicit algorithm set?
    "oct" -> oct_decoder(id)
    "EC" -> ec_decoder(id)
    "RSA" -> rsa_decoder(id)
    _ -> decode.failure(SignHmac(id, core.Sha256, <<>>), "kty")
  }
}

/// Parses a private JWK into a signing key.
///
/// Duplicate JSON object members are handled consistently on all targets: the
/// lexically last member value is used, matching JavaScript's `JSON.parse`.
///
/// ```gleam
/// sign_key.parse_jwk(private_jwk)
/// ```
pub fn parse_jwk(private_jwk: String) -> Result(SignKey, json.DecodeError) {
  jose_json.parse(private_jwk, decoder())
}

fn oct_decoder(id: Option(String)) -> Decoder(SignKey) {
  use alg <- decode.field("alg", algorithm.decoder())
  use k <- decode.field("k", core.bits_decoder())

  let digest_type = algorithm.digest_type(alg)
  case algorithm.is_hmac(alg) {
    True ->
      case hmac_key(id, digest_type, k) {
        Ok(key) -> decode.success(key)
        Error(_) -> decode.failure(SignHmac(id, digest_type, k), "k")
      }
    False -> decode.failure(SignHmac(id, digest_type, k), "alg")
  }
}

fn ec_decoder(id: Option(String)) -> Decoder(SignKey) {
  // An Elliptic Curve public key is represented by a pair of coordinates
  // drawn from a finite field, which together define a point on an
  // Elliptic Curve.  The following members MUST be present for all
  // Elliptic Curve public keys:
  //   o  "crv"
  //   o  "x"
  // The following member MUST also be present for Elliptic Curve public
  // keys for the three curves defined in the following section:
  //   o  "y"
  use crv <- decode.field("crv", core.curve_decoder())

  let coord_size = core.named_curve_size(crv)
  use x <- decode.field("x", core.bits_of_length_decoder(coord_size))
  use y <- decode.field("y", core.bits_of_length_decoder(coord_size))

  // The "d" (ECC private key) parameter contains the Elliptic Curve
  // private key value.  It is represented as the base64url encoding of
  // the octet string representation of the private key value, as defined
  // in Section 2.3.7 of SEC1 [SEC1].  The length of this octet string
  // MUST be ceiling(log-base-2(n)/8) octets (where n is the order of the
  // curve).
  use private_key <- decode.field("d", core.bits_of_length_decoder(coord_size))

  // The first octet of the OCTET STRING indicates whether the key is
  // compressed or uncompressed.  The uncompressed form is indicated
  // by 0x04 and the compressed form is indicated by either 0x02 or
  // 0x03 (see 2.3.3 in [SEC1]).  The public key MUST be rejected if
  // any other value is included in the first octet.
  let public_key = <<0x4, x:bits, y:bits>>

  let digest_type = case crv {
    core.Secp256r1 -> core.Sha256
    core.Secp384r1 -> core.Sha384
    core.Secp521r1 -> core.Sha512
  }

  let sign_key = SignEcdsa(id, crv, digest_type, public_key, private_key)

  // The curve uniquely determines the algorithm (RFC 7518 §3.1). Use the
  // curve-derived algorithm as the fallback when "alg" is absent (RFC 7517
  // §4.4: "alg" is optional). When present, it must agree with the curve.
  let expected_alg = algorithm.ec_algorithm(crv)
  use alg <- decode.optional_field("alg", expected_alg, algorithm.decoder())

  case alg == expected_alg {
    True -> decode.success(sign_key)
    False -> decode.failure(sign_key, "alg")
  }
}

fn rsa_decoder(id: Option(String)) -> Decoder(SignKey) {
  use alg <- decode.field("alg", algorithm.decoder())

  let field_decoder = core.int_decoder()
  // The "e" (exponent) parameter contains the exponent value for the RSA
  // public key.  It is represented as a Base64urlUInt-encoded value.
  // For instance, when representing the value 65537, the octet sequence
  // to be base64url-encoded MUST consist of the three octets [1, 0, 1];
  // the resulting representation for this value is "AQAB".
  use public_exponent <- decode.field("e", field_decoder)
  // The "n" (modulus) parameter contains the modulus value for the RSA
  // public key.  It is represented as a Base64urlUInt-encoded value.
  use modulus <- decode.field("n", field_decoder)

  // The "d" (private exponent) parameter contains the private exponent
  // value for the RSA private key.  It is represented as a Base64urlUInt-
  // encoded value.
  use private_exponent <- decode.field("d", field_decoder)

  // The parameter "d" is REQUIRED for RSA private keys.  The others enable
  // optimizations and SHOULD be included by producers of JWKs
  // representing RSA private keys.  If the producer includes any of the
  // other private key parameters, then all of the others MUST be present,
  // with the exception of "oth", which MUST only be present when more
  // than two prime factors were used.

  let digest_type = algorithm.digest_type(alg)
  use padding <- decode.then(case algorithm.padding(alg) {
    Ok(padding) -> decode.success(padding)
    Error(_) -> decode.failure(core.RsaPkcs1Padding, "alg")
  })

  use _ <- decode.then(case core.rsa_modulus_is_large_enough(modulus) {
    True -> decode.success(Nil)
    False -> decode.failure(Nil, "n")
  })

  rsa_private_parameters_decoder(
    id:,
    digest_type:,
    padding:,
    public_exponent:,
    modulus:,
    private_exponent:,
  )
}

fn rsa_private_parameters_decoder(
  id id: Option(String),
  digest_type digest_type: core.DigestType,
  padding padding: core.Padding,
  public_exponent public_exponent: BigInt,
  modulus modulus: BigInt,
  private_exponent private_exponent: BigInt,
) -> Decoder(SignKey) {
  let field_decoder = core.int_decoder()
  let optional_int = decode.map(field_decoder, Some)
  let simple_key =
    SignRsaSimple(
      id:,
      digest_type:,
      public_exponent:,
      modulus:,
      private_exponent:,
      padding:,
    )

  // If the producer includes any of "p", "q", "dp", "dq", or "qi", then all
  // of them MUST be present.

  // The "p" (first prime factor) parameter contains the first prime
  // factor.  It is represented as a Base64urlUInt-encoded value.
  use first_prime_factor <- decode.optional_field("p", None, optional_int)
  // The "q" (second prime factor) parameter contains the second prime
  // factor.  It is represented as a Base64urlUInt-encoded value.
  use second_prime_factor <- decode.optional_field("q", None, optional_int)
  // The "dp" (first factor CRT exponent) parameter contains the Chinese
  // Remainder Theorem (CRT) exponent of the first factor.  It is
  // represented as a Base64urlUInt-encoded value.
  use first_factor_crt_exponent <- decode.optional_field(
    "dp",
    None,
    optional_int,
  )
  // The "dq" (second factor CRT exponent) parameter contains the CRT
  // exponent of the second factor.  It is represented as a Base64urlUInt-
  // encoded value.
  use second_factor_crt_exponent <- decode.optional_field(
    "dq",
    None,
    optional_int,
  )
  // The "qi" (first CRT coefficient) parameter contains the CRT
  // coefficient of the second factor.  It is represented as a
  // Base64urlUInt-encoded value.
  use first_crt_coefficient <- decode.optional_field("qi", None, optional_int)

  // The "oth" (other primes info) parameter contains an array of
  // information about any third and subsequent primes, should they exist.
  // When only two primes have been used (the normal case), this parameter
  // MUST be omitted.  When three or more primes have been used, the
  // number of array elements MUST be the number of primes used minus two.
  // For more information on this case, see the description of the
  // OtherPrimeInfo parameters in Appendix A.1.2 of RFC 3447 [RFC3447],
  // upon which the following parameters are modeled.  If the consumer of
  // a JWK does not support private keys with more than two primes and it
  // encounters a private key that includes the "oth" parameter, then it
  // MUST NOT use the key.  Each array element MUST be an object with the
  // following members:
  //   o  "r" (prime factor)
  //   o  "d" (factor CRT exponent)
  //   o  "t" (factor CRT coefficient)
  use other_primes_info <- decode.optional_field(
    "oth",
    [],
    decode.list(other_primes_info_decoder()),
  )

  case
    first_prime_factor,
    second_prime_factor,
    first_factor_crt_exponent,
    second_factor_crt_exponent,
    first_crt_coefficient,
    other_primes_info
  {
    None, None, None, None, None, [] -> decode.success(simple_key)

    Some(first_prime_factor),
      Some(second_prime_factor),
      Some(first_factor_crt_exponent),
      Some(second_factor_crt_exponent),
      Some(first_crt_coefficient),
      other_primes_info
    ->
      decode.success(SignRsaFull(
        id:,
        digest_type:,
        public_exponent:,
        modulus:,
        private_exponent:,
        first_prime_factor:,
        second_prime_factor:,
        first_factor_crt_exponent:,
        second_factor_crt_exponent:,
        first_crt_coefficient:,
        other_primes_info:,
        padding:,
      ))

    _, _, _, _, _, _ -> decode.failure(simple_key, "RSA CRT parameters")
  }
}

fn other_primes_info_decoder() -> Decoder(#(BigInt, BigInt, BigInt)) {
  let field_decoder = core.int_decoder()
  // The "r" (prime factor) parameter within an "oth" array member
  // represents the value of a subsequent prime factor.  It is represented
  // as a Base64urlUInt-encoded value.
  use factor <- decode.field("r", field_decoder)
  // The "d" (factor CRT exponent) parameter within an "oth" array member
  // represents the CRT exponent of the corresponding prime factor.  It is
  // represented as a Base64urlUInt-encoded value.
  use exponent <- decode.field("d", field_decoder)
  // The "t" (factor CRT coefficient) parameter within an "oth" array
  // member represents the CRT coefficient of the corresponding prime
  // factor.  It is represented as a Base64urlUInt-encoded value.
  use coefficient <- decode.field("t", field_decoder)

  decode.success(#(factor, exponent, coefficient))
}

/// Returns the key id, if one is set.
///
/// ```gleam
/// sign_key.id(key)
/// ```
pub fn id(key: SignKey) -> Result(String, Nil) {
  option.to_result(key.id, Nil)
}

// -- INTERNALS -------------------------------------------------------------------------

@internal
pub fn match(sign_key: SignKey, on_ecdsa, on_rsa_simple, on_rsa_full, on_hmac) {
  case sign_key {
    SignEcdsa(id:, curve:, digest_type:, public_key:, private_key:) ->
      on_ecdsa(id, curve, digest_type, public_key, private_key)
    SignHmac(id:, digest_type:, secret:) -> on_hmac(id, digest_type, secret)
    SignRsaFull(
      id:,
      digest_type:,
      public_exponent:,
      modulus:,
      private_exponent:,
      first_prime_factor:,
      second_prime_factor:,
      first_factor_crt_exponent:,
      second_factor_crt_exponent:,
      first_crt_coefficient:,
      other_primes_info:,
      padding:,
    ) ->
      on_rsa_full(
        id,
        digest_type,
        public_exponent,
        modulus,
        private_exponent,
        first_prime_factor,
        second_prime_factor,
        first_factor_crt_exponent,
        second_factor_crt_exponent,
        first_crt_coefficient,
        other_primes_info,
        padding,
      )
    SignRsaSimple(
      id:,
      digest_type:,
      public_exponent:,
      modulus:,
      private_exponent:,
      padding:,
    ) ->
      on_rsa_simple(
        id,
        digest_type,
        public_exponent,
        modulus,
        private_exponent,
        padding,
      )
  }
}

@internal
pub fn algorithm(key: SignKey) -> Algorithm {
  case key {
    SignEcdsa(digest_type: core.Sha256, ..) -> algorithm.es256
    SignEcdsa(digest_type: core.Sha384, ..) -> algorithm.es384
    SignEcdsa(digest_type: core.Sha512, ..) -> algorithm.es512
    SignHmac(digest_type: core.Sha256, ..) -> algorithm.hs256
    SignHmac(digest_type: core.Sha384, ..) -> algorithm.hs384
    SignHmac(digest_type: core.Sha512, ..) -> algorithm.hs512

    SignRsaSimple(digest_type: core.Sha256, padding: core.RsaPkcs1Padding, ..)
    | SignRsaFull(digest_type: core.Sha256, padding: core.RsaPkcs1Padding, ..) ->
      algorithm.rs256

    SignRsaSimple(digest_type: core.Sha384, padding: core.RsaPkcs1Padding, ..)
    | SignRsaFull(digest_type: core.Sha384, padding: core.RsaPkcs1Padding, ..) ->
      algorithm.rs384

    SignRsaSimple(digest_type: core.Sha512, padding: core.RsaPkcs1Padding, ..)
    | SignRsaFull(digest_type: core.Sha512, padding: core.RsaPkcs1Padding, ..) ->
      algorithm.rs512

    SignRsaSimple(
      digest_type: core.Sha256,
      padding: core.RsaPkcs1PssPadding,
      ..,
    ) -> algorithm.ps256

    SignRsaFull(digest_type: core.Sha256, padding: core.RsaPkcs1PssPadding, ..) ->
      algorithm.ps256

    SignRsaSimple(
      digest_type: core.Sha384,
      padding: core.RsaPkcs1PssPadding,
      ..,
    ) -> algorithm.ps384

    SignRsaFull(digest_type: core.Sha384, padding: core.RsaPkcs1PssPadding, ..) ->
      algorithm.ps384

    SignRsaSimple(
      digest_type: core.Sha512,
      padding: core.RsaPkcs1PssPadding,
      ..,
    ) -> algorithm.ps512

    SignRsaFull(digest_type: core.Sha512, padding: core.RsaPkcs1PssPadding, ..) ->
      algorithm.ps512
  }
}

@internal
pub fn to_jwk(key: SignKey) -> json.Json {
  let fields = case key {
    SignHmac(..) -> [
      #("kty", json.string("oct")),
      #("k", core.bits_to_json(key.secret)),
    ]
    SignEcdsa(..) -> {
      let coord_size = core.named_curve_size(key.curve)
      let assert <<0x4, x:bytes-size(coord_size), y:bytes-size(coord_size)>> =
        key.public_key
      [
        #("kty", json.string("EC")),
        #("crv", json.string(core.named_curve(key.curve))),
        #("x", core.bits_to_json(x)),
        #("y", core.bits_to_json(y)),
        #("d", core.bits_to_json(key.private_key)),
      ]
    }
    SignRsaSimple(..) -> [
      #("kty", json.string("RSA")),
      #("e", core.int_to_json(key.public_exponent)),
      #("n", core.int_to_json(key.modulus)),
      #("d", core.int_to_json(key.private_exponent)),
    ]
    SignRsaFull(..) -> {
      let fields = [
        #("kty", json.string("RSA")),
        #("e", core.int_to_json(key.public_exponent)),
        #("n", core.int_to_json(key.modulus)),
        #("d", core.int_to_json(key.private_exponent)),
        #("p", core.int_to_json(key.first_prime_factor)),
        #("q", core.int_to_json(key.second_prime_factor)),
        #("dp", core.int_to_json(key.first_factor_crt_exponent)),
        #("dq", core.int_to_json(key.second_factor_crt_exponent)),
        #("qi", core.int_to_json(key.first_crt_coefficient)),
      ]
      case key.other_primes_info {
        [] -> fields
        [_, ..] -> [
          #(
            "oth",
            json.array(key.other_primes_info, fn(info) {
              json.object([
                #("r", core.int_to_json(info.0)),
                #("d", core.int_to_json(info.1)),
                #("t", core.int_to_json(info.2)),
              ])
            }),
          ),
          ..fields
        ]
      }
    }
  }

  let fields = [#("alg", algorithm.to_json(algorithm(key))), ..fields]

  let fields = case key.id {
    Some(id) -> [#("kid", json.string(id)), ..fields]
    None -> fields
  }

  json.object(fields)
}