ywt_core

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
import ywt/sign_key

/// A cryptographic key that can verify JWT signatures but cannot create them.
pub opaque type VerifyKey {
  VerifyEcdsa(
    id: Option(String),
    curve: core.NamedCurve,
    digest_type: core.DigestType,
    public_key: BitArray,
  )
  VerifyRsa(
    id: Option(String),
    digest_type: core.DigestType,
    exponent: BigInt,
    modulus: BigInt,
    padding: core.Padding,
  )
  VerifyHmac(id: Option(String), digest_type: core.DigestType, secret: BitArray)
}

/// Construct a verification key for token validation from a signing key.
///
/// ## Security Considerations
/// - For asymmetric keys (RSA/ECDSA), this extracts only the public key
/// - For symmetric keys (HMAC), this returns the same secret (use with caution)
/// - Distribute verify keys to services that only need to check token validity
/// - Never give signing keys to services that only need verification
///
/// ## Example
/// ```gleam
/// import gleam/crypto
///
/// let secret = crypto.strong_random_bytes(32)
/// // Create a signing key
/// let sign_key = sign_key.hs256(secret)
///
/// // Extract verification key for a service
/// let verify_key = verify_key.derived(sign_key)
///
/// // This token can be used to verify but not create tokens
/// ywt.decode(jwt_token, using: payload_decoder, claims: [], keys: [verify_key])
/// ```
///
/// 💡 **Best Practice:** Only your auth service should have signing keys.
/// All other services should only have verification keys.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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, digest, e, n, padding) },
    fn(id, digest, e, n, _, _, _, _, _, _, _, padding) {
      VerifyRsa(id, digest, e, n, padding)
    },
    fn(id, digest_type, secret) { VerifyHmac(id, digest_type, secret) },
  )
}

/// Decodes a JSON Web Key (JWK) containing a public key into a verification key.
///
/// Use this when receiving public keys from external sources, key management services,
/// or when implementing key rotation with public key distribution.
///
/// This function only supports RSA or Elliptic Curve keys.
///
/// ## Security Considerations
/// - Only accept keys from trusted sources
/// - Only accepts asymmetric public keys (EC, RSA) - never processes private key material
/// - Does not implement key pinning or certificate validation
///
/// ## Example
/// ```gleam
/// import gleam/json
///
/// // Receive a public key from your identity provider
/// let jwk_json = "{
///   \"kty\": \"EC\",
///   \"crv\": \"P-256\",
///   \"alg\": \"ES256\",
///   \"x\": \"...\",
///   \"y\": \"...\"
/// }"
///
/// case json.parse(jwk_json, verify_key.decoder()) {
///   Ok(verify_key) -> {
///     // Use this key to verify tokens from the identity provider
///     ywt.parse(token, using: decoder, claims: [], keys: [verify_key])
///   }
///   Error(_) -> // Handle invalid key format
/// }
/// ```
///
/// 📝 **Note:** Always validate that keys come from trusted sources and consider
/// implementing key rotation mechanisms.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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, core.Sha256, <<>>), "kty")
  }
}

/// Decodes a JSON Web Key Set (JWKS) into a list of verification keys.
///
/// Use this when receiving public keys from external sources, key management services,
/// or when implementing key rotation with public key distribution.
///
/// ## Security Considerations
/// - Only accept keys from trusted sources
/// - Only accepts asymmetric public keys (EC, RSA) - never processes private key material
/// - Does not implement key pinning or certificate validation
///
/// ## Example
/// ```gleam
/// import gleam/json
///
/// // Receive a public key from your identity provider
/// let jwk_json = "{\"keys\": [{\"kty\": \"EC\", ...}]}"
///
/// case json.parse(jwk_json, verify_key.set_decoder()) {
///   Ok(verify_keys) -> {
///     ywt.parse(token, using: decoder, claims: [], keys: verify_keys)
///   }
///   Error(_) -> // Handle invalid key format
/// }
/// ```
///
/// 📝 **Note:** Always validate that keys come from trusted sources and consider
/// implementing key rotation mechanisms.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn set_decoder() -> Decoder(List(VerifyKey)) {
  decode.at(["keys"], decode.list(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())

  // 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 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(VerifyEcdsa(id, crv, digest_type, public_key))
}

fn rsa_decoder(id: Option(String)) -> Decoder(VerifyKey) {
  use alg <- decode.field("alg", algorithm.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 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())

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

  decode.success(VerifyRsa(id:, digest_type:, exponent:, modulus:, padding:))
}

/// Encodes a verification key into a JSON Web Key (JWK).
///
/// This function serializes a `VerifyKey` into the standard JWK format defined in RFC 7517,
/// making it suitable for key distribution, storage, or exchange with other JWT libraries
/// and services.
///
/// ## Security Considerations
/// - Only converts verification keys (public keys and HMAC secrets for verification)
/// - HMAC keys will include the secret in the JWK - distribute carefully
/// - For asymmetric keys, only public information is included (safe to distribute)
/// - Include key IDs (`kid`) to enable proper key rotation and selection
/// - Consider key expiration and rotation policies when distributing JWKs
///
/// ## Example
/// ```gleam
/// // Convert an ECDSA public key to JWK
/// let verify_key = verify_key.derived(ecdsa_sign_key)
/// let jwk = verify_key.to_jwk(verify_key)
///
/// // Serve this JWK for other services to verify your tokens
/// json.to_string(jwk)
/// // Returns: {"kty":"EC","crv":"P-256","x":"...","y":"...","alg":"ES256","kid":"key1"}
/// ```
///
/// ⚠️ **Warning:** For HMAC keys, the resulting JWK contains the shared secret.
/// Only distribute HMAC JWKs to trusted services that need to verify tokens.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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 = [#("alg", algorithm.to_json(algorithm(key))), ..fields]

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

  json.object(fields)
}

/// Converts a list of verification keys to a JSON Web Key Set (JWKS) format.
///
/// This creates a standard JWKS structure containing multiple keys, which is the
/// standard way to publish multiple verification keys for key rotation, different
/// algorithms, or multi-tenant scenarios.
///
/// ## Usage
/// Use this to:
/// - Implement a JWKS endpoint (`/.well-known/jwks.json`)
/// - Support key rotation with multiple active keys
/// - Provide keys for different algorithms or use cases
/// - Enable external services to verify your JWTs
///
/// ## Security Considerations
/// - JWKS endpoints are typically public - ensure only verification keys are included
/// - Implement proper caching headers for JWKS endpoints (but allow for key rotation)
/// - For HMAC keys, be extremely careful about distribution
/// - Consider rate limiting your JWKS endpoint to prevent abuse
/// - Include appropriate CORS headers if serving to web applications
///
/// ## Example
/// ```gleam
/// // Create a JWKS with multiple keys for rotation
/// let current_key = verify_key.derived(current_signing_key)
/// let previous_key = verify_key.derived(previous_signing_key)
/// let backup_key = verify_key.derived(backup_signing_key)
///
/// let jwks = verify_key.to_jwks([current_key, previous_key, backup_key])
///
/// // Serve this at /.well-known/jwks.json
/// json.to_string(jwks)
/// // Returns: {"keys":[{"kty":"EC",...},{"kty":"RSA",...},{"kty":"EC",...}]}
/// ```
///
/// 💡 **Best Practice:** Include multiple keys during rotation periods to ensure
/// tokens signed with old keys remain valid during the transition period.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn to_jwks(keys: List(VerifyKey)) -> Json {
  json.object([#("keys", json.array(keys, to_jwk))])
}

/// Extracts the key identifier from a verification 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: 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) -> Algorithm {
  case key {
    VerifyEcdsa(digest_type: core.Sha256, ..) -> algorithm.es256
    VerifyEcdsa(digest_type: core.Sha384, ..) -> algorithm.es384
    VerifyEcdsa(digest_type: core.Sha512, ..) -> algorithm.es512
    VerifyHmac(digest_type: core.Sha256, ..) -> algorithm.hs256
    VerifyHmac(digest_type: core.Sha384, ..) -> algorithm.hs384
    VerifyHmac(digest_type: core.Sha512, ..) -> algorithm.hs512

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

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

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

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

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

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