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/dynamic/decode.{type Decoder}
import gleam/json.{type Json}
import gleam/option.{type Option, None, Some}
import ywt/algorithm.{type Algorithm}
import ywt/internal/core.{
  RsaPkcs1Padding, RsaPkcs1PssPadding, Sha256, Sha384, Sha512,
}
import ywt/internal/jose_json
import ywt/sign_key

/// A key that can verify JWT signatures.
///
/// Asymmetric verification keys contain public key material. HMAC verification
/// keys contain the shared secret and can also be used to sign tokens.
pub opaque type VerifyKey {
  VerifyEcdsa(
    id: Option(String),
    curve: core.NamedCurve,
    digest_type: core.DigestType,
    public_key: BitArray,
  )
  VerifyRsa(
    id: Option(String),
    digest_type: Option(core.DigestType),
    exponent: BigInt,
    modulus: BigInt,
    padding: Option(core.Padding),
  )
  VerifyHmac(id: Option(String), digest_type: core.DigestType, secret: BitArray)
}

/// Derives a verification key from a signing key.
///
/// For RSA and ECDSA this extracts public key material. For HMAC this returns
/// the same shared secret, so treat the result as sensitive.
///
/// ```gleam
/// let verify_key = verify_key.derived(signing_key)
/// ```
pub fn derived(sign_key: sign_key.SignKey) -> VerifyKey {
  sign_key.match(
    sign_key,
    fn(id, curve, digest_type, public_key, _) {
      VerifyEcdsa(id, curve, digest_type, public_key)
    },
    fn(id, digest, e, n, _, padding) {
      VerifyRsa(id, Some(digest), e, n, Some(padding))
    },
    fn(id, digest, e, n, _, _, _, _, _, _, _, padding) {
      VerifyRsa(id, Some(digest), e, n, Some(padding))
    },
    fn(id, digest_type, secret) { VerifyHmac(id, digest_type, secret) },
  )
}

/// Decodes a public JWK into a verification key.
///
/// Only decode keys from trusted sources. Unknown JWK fields are ignored;
/// `key_ops`, `use`, certificates, and key pinning are not enforced.
/// This decoder accepts EC and RSA public keys, not HMAC shared-secret JWKs.
///
/// RSA JWKs may omit `alg`. In that case, JWT decoding uses the signed `alg`
/// header to specialize the key. Raw signature verification with that key fails
/// closed because there is no JWT header.
/// For raw JSON strings, prefer `parse_jwk` so duplicate member handling is
/// consistent on all targets.
///
/// ```gleam
/// json.parse(jwk_json, verify_key.decoder())
/// ```
pub fn decoder() -> Decoder(VerifyKey) {
  // 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 {
    "EC" -> ec_decoder(id)
    "RSA" -> rsa_decoder(id)
    _ -> decode.failure(VerifyHmac(id, Sha256, <<>>), "kty")
  }
}

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

/// Decodes a JWKS into a list of verification keys.
///
/// This has the same trust and validation limits as `decoder`. It expects a JSON
/// object with a `keys` list.
/// For raw JSON strings, prefer `parse_jwks` so duplicate member handling is
/// consistent on all targets.
///
/// ```gleam
/// json.parse(jwks_json, verify_key.set_decoder())
/// ```
pub fn set_decoder() -> Decoder(List(VerifyKey)) {
  decode.at(["keys"], decode.list(decoder()))
}

/// Parses a JWKS into a list of verification keys.
///
/// Duplicate JSON object members are handled consistently on all targets: the
/// lexically last member value is used, matching JavaScript's `JSON.parse`.
///
/// ```gleam
/// verify_key.parse_jwks(jwks_json)
/// ```
pub fn parse_jwks(jwks: String) -> Result(List(VerifyKey), json.DecodeError) {
  jose_json.parse(jwks, set_decoder())
}

fn ec_decoder(id: Option(String)) -> Decoder(VerifyKey) {
  // 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 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 -> Sha256
    core.Secp384r1 -> Sha384
    core.Secp521r1 -> Sha512
  }

  let verify_key = VerifyEcdsa(id, crv, digest_type, public_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(verify_key)
    False -> decode.failure(verify_key, "alg")
  }
}

fn rsa_decoder(id: Option(String)) -> Decoder(VerifyKey) {
  // 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 exponent <- decode.field("e", core.int_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", core.int_decoder())

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

  // RFC 7517 §4.4: "alg" is optional. When absent, digest_type and padding
  // are stored as None; the JWT "alg" header (part of the signed content)
  // is used at verify time via for_algorithm to specialize the key.
  use alg <- decode.optional_field(
    "alg",
    None,
    decode.map(algorithm.decoder(), Some),
  )

  case alg {
    None -> {
      let key =
        VerifyRsa(id:, digest_type: None, exponent:, modulus:, padding: None)
      decode.success(key)
    }

    Some(alg) ->
      case algorithm.padding(alg) {
        Ok(padding) -> {
          let digest_type = Some(algorithm.digest_type(alg))
          let padding = Some(padding)
          let key = VerifyRsa(id:, digest_type:, exponent:, modulus:, padding:)
          decode.success(key)
        }
        Error(_) ->
          decode.failure(VerifyRsa(id, None, exponent, modulus, None), "alg")
      }
  }
}

/// Encodes a verification key as a JWK.
///
/// Asymmetric keys encode only public key material. HMAC keys encode the shared
/// secret, so do not publish HMAC JWKs.
///
/// ```gleam
/// verify_key.derived(signing_key) |> verify_key.to_jwk
/// ```
pub fn to_jwk(key: VerifyKey) -> Json {
  let fields = case key {
    VerifyHmac(..) -> [
      #("kty", json.string("oct")),
      #("k", core.bits_to_json(key.secret)),
    ]

    VerifyEcdsa(..) -> {
      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)),
      ]
    }

    VerifyRsa(..) -> [
      #("kty", json.string("RSA")),
      #("e", core.int_to_json(key.exponent)),
      #("n", core.int_to_json(key.modulus)),
    ]
  }

  let fields = case algorithm(key) {
    Ok(alg) -> [#("alg", algorithm.to_json(alg)), ..fields]
    Error(Nil) -> fields
  }

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

  json.object(fields)
}

/// Encodes verification keys as a JWKS.
///
/// This is suitable for publishing asymmetric keys during rotation. Do not
/// publish HMAC keys, because their JWK form contains the shared secret.
///
/// ```gleam
/// verify_key.to_jwks([current_key, previous_key])
/// ```
pub fn to_jwks(keys: List(VerifyKey)) -> Json {
  json.object([#("keys", json.array(keys, to_jwk))])
}

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

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

@internal
pub fn match(verify_key: VerifyKey, on_ecdsa, on_rsa, on_hmac) {
  case verify_key {
    VerifyEcdsa(id:, curve:, digest_type:, public_key:) ->
      on_ecdsa(id, curve, digest_type, public_key)
    VerifyHmac(id:, digest_type:, secret:) -> on_hmac(id, digest_type, secret)
    VerifyRsa(id:, digest_type:, exponent:, modulus:, padding:) ->
      on_rsa(id, digest_type, exponent, modulus, padding)
  }
}

@internal
pub fn algorithm(key: VerifyKey) -> Result(Algorithm, Nil) {
  case key {
    VerifyEcdsa(digest_type: Sha256, ..) -> Ok(algorithm.es256)
    VerifyEcdsa(digest_type: Sha384, ..) -> Ok(algorithm.es384)
    VerifyEcdsa(digest_type: Sha512, ..) -> Ok(algorithm.es512)
    VerifyHmac(digest_type: Sha256, ..) -> Ok(algorithm.hs256)
    VerifyHmac(digest_type: Sha384, ..) -> Ok(algorithm.hs384)
    VerifyHmac(digest_type: Sha512, ..) -> Ok(algorithm.hs512)

    VerifyRsa(digest_type: Some(Sha256), padding: Some(RsaPkcs1Padding), ..) ->
      Ok(algorithm.rs256)

    VerifyRsa(digest_type: Some(Sha384), padding: Some(RsaPkcs1Padding), ..) ->
      Ok(algorithm.rs384)

    VerifyRsa(digest_type: Some(Sha512), padding: Some(RsaPkcs1Padding), ..) ->
      Ok(algorithm.rs512)

    VerifyRsa(digest_type: Some(Sha256), padding: Some(RsaPkcs1PssPadding), ..) ->
      Ok(algorithm.ps256)

    VerifyRsa(digest_type: Some(Sha384), padding: Some(RsaPkcs1PssPadding), ..) ->
      Ok(algorithm.ps384)

    VerifyRsa(digest_type: Some(Sha512), padding: Some(RsaPkcs1PssPadding), ..) ->
      Ok(algorithm.ps512)

    VerifyRsa(..) -> Error(Nil)
  }
}

/// Returns `key` locked to `alg`, or `Error(Nil)` if it is incompatible.
///
/// RSA JWKs without `alg` are specialized from the signed JWT `alg` header during
/// verification. Keys that already have an algorithm must match exactly.
@internal
pub fn for_algorithm(key: VerifyKey, alg: Algorithm) -> Result(VerifyKey, Nil) {
  case key {
    VerifyRsa(digest_type: None, id:, exponent:, modulus:, ..) ->
      case algorithm.padding(alg) {
        Ok(padding) -> {
          let digest_type = Some(algorithm.digest_type(alg))
          let padding = Some(padding)
          Ok(VerifyRsa(id:, digest_type:, exponent:, modulus:, padding:))
        }
        Error(Nil) -> Error(Nil)
      }
    _ ->
      case algorithm(key) == Ok(alg) {
        True -> Ok(key)
        False -> Error(Nil)
      }
  }
}