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 gleam/dict
import gleam/dynamic
import gleam/dynamic/decode.{type Decoder, type Dynamic}
import gleam/float
import gleam/int
import gleam/json.{type Json}
import gleam/list
import gleam/option.{None, Some}
import gleam/order
import gleam/time/duration.{type Duration}
import gleam/time/timestamp.{type Timestamp}
import ywt/internal/core.{type Error}

/// A rule that validates or writes a JWT header or payload field.
///
/// Required claims must be present when decoding. Optional claims are checked
/// only when the token includes the corresponding field.
pub opaque type Claim {
  HeaderClaim(
    name: String,
    is_required: Bool,
    verify: fn(Dynamic) -> Result(Nil, Error),
    value: fn() -> Json,
  )
  PayloadClaim(
    name: String,
    is_required: Bool,
    verify: fn(Dynamic) -> Result(Nil, Error),
    value: fn() -> Json,
  )
}

/// The `typ` header describes what kind of token this is.
///
/// The field is optional in the JWT specification, but some systems require
/// `"JWT"` or another explicit type. It is not a security boundary; still
/// validate issuer, audience, expiration, and signature.
///
/// ```gleam
/// let claims = [
///   claim.typ("JWT"),
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn typ(expected: String) -> Claim {
  let name = "typ"
  let value = fn() { json.string(expected) }

  let verify =
    validate(
      name,
      decode.string,
      fn(found) { core.InvalidType(expected, found) },
      fn(v) { expected == v },
    )

  HeaderClaim(name:, is_required: True, value:, verify:)
}

/// The `iat` claim records when the token was issued.
///
/// This writes the current time when creating a token, but it does not reject
/// future-dated tokens when verifying. Use `not_before` for that.
///
/// ```gleam
/// let claims = [
///   claim.issued_at(),
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn issued_at() -> Claim {
  let name = "iat"
  let value = fn() { encode_numeric_date(timestamp.system_time()) }
  let verify = fn(data) {
    case decode.run(data, numeric_date_decoder()) {
      Ok(_) -> Ok(Nil)
      Error(error) -> Error(core.ClaimDecodingError(name, error))
    }
  }

  PayloadClaim(name:, is_required: False, value:, verify:)
}

/// The `iss` claim identifies who issued the token.
///
/// Use this to reject tokens from the wrong issuer. Issuer strings often point
/// at the service that also publishes verification keys, but ywt does not fetch
/// or trust those keys automatically.
///
/// ```gleam
/// let claims = [
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn issuer(issuer: String, others: List(String)) -> Claim {
  string_claim("iss", issuer, others, core.InvalidIssuer)
}

/// The `aud` claim identifies who the token is meant for.
///
/// Tokens with an `aud` field are rejected by default unless you add this claim
/// and the value matches one of the accepted audiences. ywt validates both
/// string and array-valued `aud` fields.
///
/// ```gleam
/// let claims = [
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn audience(primary: String, others: List(String)) -> Claim {
  let name = "aud"
  let expected = [primary, ..others]
  let value = fn() { json.string(primary) }

  let verify =
    validate(
      name,
      audience_decoder(),
      invalid_audience(expected, _),
      any_audience_matches(_, expected),
    )

  PayloadClaim(name:, is_required: True, value:, verify:)
}

/// The `nbf` claim says the token must not be accepted before a time.
///
/// Tokens with `nbf` are checked by default with zero leeway. Add this claim
/// when you want to set the value while signing or allow clock skew while
/// verifying. Keep leeway small.
///
/// ```gleam
/// let claims = [
///   claim.not_before(timestamp.system_time(), leeway: duration.minutes(1)),
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn not_before(time time: Timestamp, leeway leeway: Duration) -> Claim {
  let name = "nbf"
  let value = fn() { encode_numeric_date(time) }

  let verify =
    validate(
      name,
      numeric_date_decoder(),
      core.TokenNotYetValid,
      fn(not_before) {
        let now = timestamp.system_time()
        case timestamp.compare(not_before, timestamp.add(now, leeway)) {
          order.Gt -> False
          _ -> True
        }
      },
    )

  PayloadClaim(name:, is_required: True, value:, verify:)
}

/// The `exp` claim says when the token expires.
///
/// Prefer short lifetimes. Tokens with `exp` are checked by default with zero
/// leeway, but ywt will not add an expiration unless you include this claim.
/// `max_age` is used when writing the token; verification checks the token's
/// `exp` value plus `leeway`.
///
/// ```gleam
/// let claims = [
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn expires_at(max_age max_age: Duration, leeway leeway: Duration) -> Claim {
  let name = "exp"

  let value = fn() {
    timestamp.system_time()
    |> timestamp.add(max_age)
    |> encode_numeric_date
  }

  let verify =
    validate(name, numeric_date_decoder(), core.TokenExpired, fn(expires_at) {
      let now = timestamp.system_time()
      case timestamp.compare(timestamp.add(expires_at, leeway), now) {
        order.Gt -> True
        _ -> False
      }
    })

  PayloadClaim(name:, is_required: True, value:, verify:)
}

/// The `jti` claim gives the token a unique id.
///
/// ywt does not store token ids. Use this with your own denylist or session
/// store if you need revocation.
///
/// ```gleam
/// let claims = [
///   claim.id("session-123", []),
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn id(id: String, others: List(String)) -> Claim {
  string_claim("jti", id, others, core.InvalidId)
}

/// The `sub` claim identifies the user or entity the token is about.
///
/// Prefer stable internal ids. Avoid putting personal data in the subject,
/// because JWT payloads are readable by whoever has the token.
///
/// ```gleam
/// let claims = [
///   claim.subject("user_123", []),
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn subject(subject: String, others: List(String)) -> Claim {
  string_claim("sub", subject, others, core.InvalidSubject)
}

/// A custom claim stores an application-specific value.
///
/// The decoded token value must exactly match `value`. Use this for simple
/// fixed requirements such as role, tenant, or token class checks.
///
/// ```gleam
/// let claims = [
///   claim.custom("role", "admin", json.string, decode.string),
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn custom(
  name name: String,
  value value: a,
  encode encode: fn(a) -> Json,
  decoder decoder: Decoder(a),
) -> Claim {
  let on_error = fn(_) { core.InvalidCustomClaim(name) }
  let check = fn(decoded) { decoded == value }
  let value = fn() { encode(value) }
  let verify = validate(name, decoder, on_error, check)

  PayloadClaim(name:, is_required: True, value:, verify:)
}

fn validate(
  name: String,
  decoder: Decoder(a),
  on_error: fn(a) -> Error,
  validate: fn(a) -> Bool,
) -> fn(Dynamic) -> Result(Nil, Error) {
  fn(data) {
    case decode.run(data, decoder) {
      Ok(value) ->
        case validate(value) {
          True -> Ok(Nil)
          False -> Error(on_error(value))
        }
      Error(error) -> Error(core.ClaimDecodingError(name, error))
    }
  }
}

fn string_claim(
  name: String,
  primary: String,
  others: List(String),
  error: fn(List(String), String) -> Error,
) {
  let expected = [primary, ..others]
  let value = fn() { json.string(primary) }
  let verify =
    validate(name, decode.string, error(expected, _), list.contains(expected, _))

  PayloadClaim(name:, is_required: True, value:, verify:)
}

fn audience_decoder() -> Decoder(List(String)) {
  decode.one_of(decode.map(decode.string, list.wrap), or: [
    decode.list(decode.string),
  ])
}

fn any_audience_matches(found: List(String), expected: List(String)) -> Bool {
  use audience <- list.any(found)
  list.contains(expected, audience)
}

fn invalid_audience(expected: List(String), found: List(String)) -> Error {
  core.InvalidAudience(expected, found)
}

fn timestamp_from_numeric_date_float(seconds: Float) -> Timestamp {
  let whole_seconds = seconds |> float.floor |> float.truncate
  let fractional_seconds = seconds -. int.to_float(whole_seconds)
  let nanoseconds = float.round(fractional_seconds *. 1_000_000_000.0)

  timestamp.from_unix_seconds_and_nanoseconds(
    seconds: whole_seconds,
    nanoseconds: nanoseconds,
  )
}

/// Decodes a JWT numeric date into a timestamp.
///
/// JWT numeric dates are seconds since the Unix epoch.
///
/// ```gleam
/// decode.field("exp", claim.numeric_date_decoder())
/// ```
pub fn numeric_date_decoder() -> Decoder(Timestamp) {
  decode.one_of(decode.map(decode.int, timestamp.from_unix_seconds), or: [
    decode.map(decode.float, timestamp_from_numeric_date_float),
  ])
}

/// Encodes a timestamp as a JWT numeric date.
///
/// Fractional seconds are truncated to keep encoded JWTs compact.
///
/// ```gleam
/// #("checked_at", claim.encode_numeric_date(timestamp.system_time()))
/// ```
pub fn encode_numeric_date(timestamp: Timestamp) {
  json.int(float.truncate(timestamp.to_unix_seconds(timestamp)))
}

/// Allows a claim to be absent during verification.
///
/// When the field is present it is still validated. If the claim is already
/// optional, it is returned unchanged.
///
/// ```gleam
/// let claims = [
///   claim.id("session-123", []) |> claim.optional,
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
///   claim.issuer("https://auth.example.com", []),
///   claim.audience("https://api.example.com", []),
/// ]
/// ```
pub fn optional(claim: Claim) -> Claim {
  case claim {
    HeaderClaim(is_required: True, ..) ->
      HeaderClaim(..claim, is_required: False)

    PayloadClaim(is_required: True, ..) ->
      PayloadClaim(..claim, is_required: False)

    HeaderClaim(is_required: False, ..)
    | PayloadClaim(is_required: False, ..) -> claim
  }
}

// -- CLAIM ENCODE/DECODE ------------------------------------------------------

/// Validates claims against a JWT payload.
///
/// This is a low-level function used internally by ywt.
///
/// This does not verify the token signature. Use the platform `ywt.decode`
/// function for normal JWT verification.
///
/// Tokens with `exp` and `nbf` are checked by default with zero leeway. Tokens
/// with `aud` are rejected by default unless you pass an `audience` claim.
/// Audience validation accepts string or array values.
pub fn verify(payload: Dynamic, claims: List(Claim)) -> Result(Nil, Error) {
  verify_with_header(dynamic.nil(), payload, claims)
}

/// Validates claims against both the JWT header and payload.
///
/// This is a low-level helper used by ywt. It adds default checks for `exp`,
/// `nbf`, and `aud` when you have not configured those claims yourself.
/// Audience validation accepts string or array values.
@internal
pub fn verify_with_header(
  header: Dynamic,
  payload: Dynamic,
  claims: List(Claim),
) -> Result(Nil, Error) {
  claims
  |> add_default_claim(expires_at(duration.seconds(0), duration.seconds(0)))
  |> add_default_claim(not_before(timestamp.system_time(), duration.seconds(0)))
  |> add_default_claim(reject_present_audience())
  |> list.try_each(verify_single_claim(header, payload, _))
}

fn reject_present_audience() -> Claim {
  let name = "aud"
  let value = fn() { json.string("") }

  let verify =
    validate(name, audience_decoder(), invalid_audience([], _), fn(_) { False })

  PayloadClaim(name:, is_required: False, value:, verify:)
}

fn add_default_claim(claims, claim: Claim) {
  // TODO: @Speed
  case has_payload_claim(claims, claim.name) {
    True -> claims
    False -> [optional(claim), ..claims]
  }
}

fn has_payload_claim(claims, name) {
  use claim <- list.any(claims)
  case claim {
    PayloadClaim(name: claim_name, ..) -> name == claim_name
    HeaderClaim(..) -> False
  }
}

/// Verify a single claim, handling required vs optional claims.
fn verify_single_claim(
  header: Dynamic,
  payload: Dynamic,
  claim: Claim,
) -> Result(Nil, Error) {
  let decoder = {
    use value <- decode.optional_field(
      claim.name,
      None,
      decode.map(decode.dynamic, Some),
    )
    decode.success(value)
  }

  let result = case claim {
    HeaderClaim(..) -> decode.run(header, decoder)
    PayloadClaim(..) -> decode.run(payload, decoder)
  }

  case result {
    Ok(Some(claim_value)) -> claim.verify(claim_value)
    Ok(None) if !claim.is_required -> Ok(Nil)
    Ok(None) -> Error(core.MissingClaim(claim.name))
    Error(error) -> Error(core.ClaimDecodingError(claim.name, error))
  }
}

/// Encodes claims and application data into a JSON object.
///
/// This is a low-level function used internally by ywt.
///
/// Claims take precedence when a field appears in both lists. Put security
/// fields such as `exp`, `iss`, `aud`, and `sub` in claims instead of payload
/// data.
///
/// ```gleam
/// let claims = [
///   claim.subject("user_123", []),
///   claim.expires_at(max_age: duration.minutes(15), leeway: duration.minutes(1)),
/// ]
///
/// claim.encode(claims, [
///   #("sub", json.string("ignored")),
///   #("role", json.string("admin")),
/// ])
/// ```
pub fn encode(claims: List(Claim), data: List(#(String, Json))) -> Json {
  // we roundtrip using a dict here to ensure consistent behaviour when a key
  // appears in both lists or mulitple times.
  // We want claims to take precedence, so we add them second.
  dict.new()
  |> list.fold(from: _, over: data, with: fn(dict, pair) {
    dict.insert(dict, pair.0, pair.1)
  })
  |> list.fold(from: _, over: claims, with: fn(dict, claim) {
    case claim {
      PayloadClaim(..) -> dict.insert(dict, claim.name, claim.value())
      _ -> dict
    }
  })
  |> dict.to_list
  |> json.object
}

@internal
pub fn encode_header(claims: List(Claim), data: List(#(String, Json))) -> Json {
  dict.new()
  |> list.fold(from: _, over: data, with: fn(dict, pair) {
    dict.insert(dict, pair.0, pair.1)
  })
  |> list.fold(from: _, over: claims, with: fn(dict, claim) {
    case claim {
      HeaderClaim(..) -> dict.insert(dict, claim.name, claim.value())
      _ -> dict
    }
  })
  |> dict.to_list
  |> json.object
}