///|
/// JWT signing algorithm.
pub(all) enum Algorithm {
  HS256
  HS384
  HS512
} derive(Debug, Eq)

///|
/// Convert algorithm to its JOSE name string.
pub fn Algorithm::to_str(self : Algorithm) -> String {
  match self {
    HS256 => "HS256"
    HS384 => "HS384"
    HS512 => "HS512"
  }
}

///|
/// JWT error type.
pub(all) enum JwtError {
  InvalidToken(String)
  InvalidSignature
  InvalidAlgorithm(String)
  InvalidHeader(String)
  InvalidClaims(String)
  TokenExpired
} derive(Debug, Eq)

///|
/// JWT header structure.
pub(all) struct JwtHeader {
  mut alg : Algorithm
  typ : String
  kid : String?
} derive(Debug, Eq)

///|
/// Default header for HS256.
pub fn JwtHeader::hs256() -> JwtHeader {
  { alg: HS256, typ: "JWT", kid: None }
}

///|
/// Encode header to JSON.
fn JwtHeader::to_json(self : JwtHeader) -> Json {
  let obj : Map[String, Json] = Map([])
  obj["alg"] = Json::string(self.alg.to_str())
  obj["typ"] = Json::string(self.typ)
  match self.kid {
    Some(k) => obj["kid"] = Json::string(k)
    None => ()
  }
  Json::object(obj)
}

///|
/// Sign a JWT token with HS256 (or HS384/HS512).
///
/// `header_json` and `claims_json` should be pre-built JSON strings.
pub fn sign(
  _alg : Algorithm,
  header_json : String,
  claims_json : String,
  secret : Bytes,
) -> String {
  let header_b64 = base64url_encode_str(header_json)
  let claims_b64 = base64url_encode_str(claims_json)
  let signing_input = header_b64 + "." + claims_b64
  let mac = hmac_sha256(secret, @utf8.encode(signing_input))
  let sig_b64 = base64url_encode(mac)
  signing_input + "." + sig_b64
}

///|
/// Convenience: sign HS256 with a string secret and pre-made header + claims JSON.
pub fn sign_hs256(
  header : JwtHeader,
  claims_json : String,
  secret : Bytes,
) -> String {
  let header_json = header.to_json().stringify()
  sign(HS256, header_json, claims_json, secret)
}

///|
/// Verify a JWT token's signature. Returns Ok(()) if valid, Err otherwise.
pub fn verify(token : String, secret : Bytes) -> Result[Unit, JwtError] {
  let parts = token.split(".").to_array()
  if parts.length() != 3 {
    return Err(InvalidToken("token must have 3 parts separated by '.'"))
  }
  let header_b64 = parts[0].to_owned()
  let claims_b64 = parts[1].to_owned()
  let sig_b64 = parts[2].to_owned()
  let signing_input = header_b64 + "." + claims_b64
  // Decode and check header
  let header_bytes = match base64url_decode(header_b64) {
    Some(b) => b
    None => return Err(InvalidHeader("cannot decode header base64url"))
  }
  ignore(header_bytes)
  // header_str not needed for signature verification
  // Compute expected signature
  let expected_mac = hmac_sha256(secret, @utf8.encode(signing_input))
  // Decode provided signature
  let provided_mac = match base64url_decode(sig_b64) {
    Some(b) => b
    None => return Err(InvalidSignature)
  }
  // Constant-time-ish comparison
  if expected_mac.length() != provided_mac.length() {
    return Err(InvalidSignature)
  }
  let mut diff = 0
  for i in 0.. Result[Json, JwtError] {
  let parts = token.split(".").to_array()
  if parts.length() != 3 {
    return Err(InvalidToken("token must have 3 parts"))
  }
  match base64url_decode(parts[1].to_owned()) {
    Some(claims_bytes) =>
      try {
        let j = @json.parse(@utf8.decode_lossy(claims_bytes))
        Ok(j)
      } catch {
        _ => Err(InvalidClaims("claims is not valid JSON"))
      }
    None => Err(InvalidClaims("cannot decode claims base64url"))
  }
}