ywt_core

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

/// A cryptographic key that can create JWT signatures.
///
/// SignKeys are the most sensitive components in a JWT system - they can create
/// valid tokens that will be trusted by your services. Protect them accordingly.
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 a symmetric HMAC-SHA256 signing key for JWT authentication.
///
/// This is the most commonly used algorithm for simple applications where you have a shared
/// secret between your application and the service that will verify JWTs. The same key is
/// used for both signing and verification.
///
/// ## Security Considerations
/// - The secret must be cryptographically random and at least 256 bits (32 bytes) long
/// - Keep the secret absolutely confidential - anyone with access can forge tokens
/// - Rotate secrets regularly and have a key rotation strategy
/// - Never store secrets in code or version control
///
/// ## Example
/// ```gleam
/// import gleam/crypto
///
/// // Generate a secure random secret
/// let secret = crypto.strong_random_bytes(32)
/// let assert Ok(key) = sign_key.hs256(secret)
///
/// // Use this key to sign JWTs
/// ywt.encode(payload: [#("sub", json.string("user123"))], claims: [], key: key)
/// ```
///
/// ⚠️ **Warning:** HMAC keys are symmetric - the same key signs and verifies tokens.
/// This means every service that needs to verify tokens must have the secret,
/// increasing the attack surface.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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)
  case bit_array.bit_size(secret) >= 256 {
    True -> Ok(SignHmac(None, core.Sha256, secret))
    False -> Error(Nil)
  }
}

/// Creates a symmetric HMAC-SHA384 signing key for for JWT authentication.
///
/// Choose this over HS256 when you need additional security margins or your security
/// policy requires SHA-384.
///
/// ## Security Considerations
/// - The secret must be cryptographically random and at least 384 bits (48 bytes) long
/// - Keep the secret absolutely confidential - anyone with access can forge tokens
/// - Rotate secrets regularly and have a key rotation strategy
/// - Never store secrets in code or version control
///
/// ## Example
/// ```gleam
/// import gleam/crypto
///
/// // Generate a secure random secret
/// let secret = crypto.strong_random_bytes(48)
/// let assert Ok(key) = sign_key.hs384(secret)
///
/// // Use this key to sign JWTs
/// ywt.encode(payload: [#("sub", json.string("user123"))], claims: [], key: key)
/// ```
///
/// ⚠️ **Warning:** HMAC keys are symmetric - the same key signs and verifies tokens.
/// This means every service that needs to verify tokens must have the secret,
/// increasing the attack surface.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn hs384(secret: BitArray) -> Result(SignKey, Nil) {
  case bit_array.bit_size(secret) >= 384 {
    True -> Ok(SignHmac(None, core.Sha384, secret))
    False -> Error(Nil)
  }
}

/// Creates a symmetric HMAC-SHA512 signing key for JWT authentication.
///
/// This is the most commonly used algorithm for simple applications where you have a shared
/// secret between your application and the service that will verify JWTs. The same key is
/// used for both signing and verification.
///
/// ## Security Considerations
/// - The secret must be cryptographically random and at least 512 bits (64 bytes) long
/// - Keep the secret absolutely confidential - anyone with access can forge tokens
/// - Rotate secrets regularly and have a key rotation strategy
/// - Never store secrets in code or version control
///
/// ## Example
/// ```gleam
/// import gleam/crypto
///
/// // Generate a secure random secret
/// let secret = crypto.strong_random_bytes(64)
/// let assert Ok(key) = sign_key.hs512(secret)
///
/// // Use this key to sign JWTs
/// ywt.encode(payload: [#("sub", json.string("user123"))], claims: [], key: key)
/// ```
///
/// ⚠️ **Warning:** HMAC keys are symmetric - the same key signs and verifies tokens.
/// This means every service that needs to verify tokens must have the secret,
/// increasing the attack surface.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn hs512(secret: BitArray) -> Result(SignKey, Nil) {
  case bit_array.bit_size(secret) >= 512 {
    True -> Ok(SignHmac(None, core.Sha512, secret))
    False -> Error(Nil)
  }
}

/// Assigns a random key identifier to a signing key for key management and rotation.
///
/// This function generates a cryptographically random identifier and assigns it to
/// the key's `id` field. Key identifiers enable proper key selection during JWT
/// verification and are essential for implementing key rotation strategies.
///
/// ## Usage
/// Use this when you need to uniquely identify keys for distribution, storage,
/// or rotation purposes. The identifier will be included in JWT headers as the
/// `kid` (Key ID) field, allowing verifiers to select the correct key.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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 JSON Web Key (JWK) containing a private key into a signing key.
///
/// Use this for loading private keys from secure configuration or key management systems
/// when implementing JWT signing services.
///
/// ## Security Considerations
/// - Only accept keys from trusted sources
/// - Store private keys in secure key management systems, never in code
/// - Only keep private keys in your auth service, they should never leave secure boundaries
/// - Rotate private keys regularly and have a key rotation strategy
///
/// ## Example
/// ```gleam
/// // Load private key from secure configuration (never hardcode!)
/// let private_jwk = load_from_secure_vault("jwt-signing-key")
///
/// case json.parse(private_jwk, sign_key.decoder()) {
///   Ok(sign_key) -> {
///     // Use to create signed JWTs
///     encode(payload: [#("sub", json.string("user123"))],
///            claims: [],
///            key: sign_key)
///   }
///   Error(_) -> // Handle key loading error
/// }
/// ```
///
/// 🔐 **Critical:** Sign keys are the most important pieces of your security.
/// Treat them accordingly.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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")
  }
}

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)

  decode.success(SignHmac(id, digest_type, k))
}

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())

  // TODO: check for the length of the coordinates early, right now we accept
  // them but return false for all verify calls.
  use x <- decode.field("x", core.bits_decoder())
  use y <- decode.field("y", core.bits_decoder())

  // 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).

  // TODO: Verify length here too
  use private_key <- decode.field("d", core.bits_decoder())

  // 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
  }

  decode.success(SignEcdsa(id, crv, digest_type, public_key, private_key))
}

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")
  })

  let full =
    rsa_full_decoder(
      id:,
      digest_type:,
      padding:,
      public_exponent:,
      modulus:,
      private_exponent:,
    )

  decode.one_of(full, or: [
    decode.success(SignRsaSimple(
      id:,
      digest_type:,
      public_exponent:,
      modulus:,
      private_exponent:,
      padding:,
    )),
  ])
}

fn rsa_full_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()
  // The "p" (first prime factor) parameter contains the first prime
  // factor.  It is represented as a Base64urlUInt-encoded value.
  use first_prime_factor <- decode.field("p", field_decoder)
  // The "q" (second prime factor) parameter contains the second prime
  // factor.  It is represented as a Base64urlUInt-encoded value.
  use second_prime_factor <- decode.field("q", field_decoder)
  // 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.field("dp", field_decoder)
  // 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.field("dq", field_decoder)
  // 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.field("qi", field_decoder)
  // 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.
  use other_primes_info <- decode.optional_field(
    "oth",
    [],
    decode.list(other_primes_info_decoder()),
  )

  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:,
  ))
}

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))
}

/// Extracts the key identifier from a signing key, if present.
///
/// This function retrieves the optional key identifier that was assigned to a key.
/// Key identifiers are used for key selection during JWT verification.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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)
}