ywt_core

import gleam/dict
import gleam/dynamic/decode.{type Decoder, type Dynamic}
import gleam/float
import gleam/json.{type Json}
import gleam/list
import gleam/order
import gleam/time/duration.{type Duration}
import gleam/time/timestamp.{type Timestamp}
import ywt/internal/core.{type ParseError}

/// A Claim is a validation rule on the JWT.
/// You can for example validate that the issuer has a value you expect,
/// or that the JWT is not expired yet.
pub opaque type Claim {
  Claim(
    name: String,
    is_required: Bool,
    verify: fn(Dynamic) -> Result(Nil, ParseError),
    value: fn() -> Json,
  )
}

/// Adds the current timestamp as the "issued at" (iat) claim.
///
/// Include this in JWTs to track when tokens were issued, useful for debugging
/// and audit trails. To reject tokens that were created in the future, use `not_before`.
///
/// ## Example
/// ```gleam
/// let claims = [
///   claim.expires_at(max_age: duration.hours(1), leeway: duration.minutes(5)),
///   ...
/// ]
/// ```
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn issued_at() -> Claim {
  let value = fn() { encode_numeric_date(timestamp.system_time()) }
  let verify = fn(_) { Ok(Nil) }

  Claim(name: "iat", is_required: False, value:, verify:)
}

/// Validates that the JWT comes from a trusted issuer, and includes that issuer
/// in generated tokens.
///
/// Use this to ensure JWTs are only accepted from authorized authentication services,
/// preventing token misuse across different services. It is common practice to
/// set the issuer to a URL with a `.well-known/jwks.json` endpoint, from which
/// public verification keys can be fetched.
///
/// ## Example
/// ```gleam
/// // Only accept tokens from your auth service
/// let claims = [
///   claim.issuer("https://auth.example.com", [])
/// ]
/// ```
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn issuer(issuer: String, others: List(String)) -> Claim {
  string_claim("iss", issuer, others, core.InvalidIssuer)
}

/// Validates that the JWT is intended for your service, and includes the primary
/// audience in generated tokens.
///
/// Include this to ensure JWTs are only used by their intended recipients,
/// preventing tokens meant for other services from being accepted.
///
/// ## Example
/// ```gleam
/// // For an API service
/// let claims = [
///   claim.audience("api.myapp.com", ["admin-api.myapp.com"]),
/// ]
/// ```
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn audience(primary: String, others: List(String)) -> Claim {
  string_claim("aud", primary, others, core.InvalidAudience)
}

/// Adds a "not before" (nbf) claim that prevents the JWT from being valid until the current time.
///
/// Use this for JWTs that should only become valid at a specific time, such as
/// scheduled operations or delayed activations.
///
/// ## Example
/// ```gleam
/// // Token becomes valid immediately, allowing 1 minute clock skew
/// let claims = [
///   claim.not_before(timestamp.system_time(), leeway: duration.minutes(1)),
///   claim.expires_at(max_age: duration.hours(2), leeway: duration.minutes(1))
/// ]
/// ```
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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
        }
      },
    )

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

/// Sets an expiration time for the JWT, automatically rejecting expired tokens.
///
/// **ALWAYS include this** - it's your primary defense against token replay attacks
/// and limits the damage from compromised tokens.
///
/// ## Security Considerations
/// - Make expiration times as short as reasonably possible
/// - Include reasonable leeway for clock differences
///
/// ## Example
/// ```gleam
/// // Short-lived API access token
/// let api_claims = [
///   claim.expires_at(max_age: duration.minutes(30), leeway: duration.minutes(5))
/// ]
///
/// // Longer-lived refresh token
/// let session_claims = [
///   claim.expires_at(max_age: duration.hours(24 * 7), leeway: duration.minutes(5))
/// ]
/// ```
///
/// 🚨 **Critical:** Never create or parse JWTs without expiration times -
/// they would be valid forever if compromised.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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
      }
    })

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

/// Adds a unique identifier (jti) to the JWT for tracking and revocation.
///
/// Use this when you need to track specific tokens or implement token revocation
/// mechanisms.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn id(id: String, others: List(String)) -> Claim {
  string_claim("jti", id, others, core.InvalidId)
}

/// Identifies the subject (typically a user) that the JWT represents.
///
/// Use this to bind JWTs to specific users or entities, enabling proper
/// authorization decisions. This could be an internal user id.
///
/// ## Security Considerations
/// - Use stable, unique identifiers (user IDs, not usernames)
/// - Don't include personally identifiable information
///
/// ## Example
/// ```gleam
/// // Bind token to a specific user
/// let claims = [
///   claim.subject("user_12345", []),
///   claim.expires_at(max_age: duration.hours(1), leeway: duration.minutes(5)),
///   claim.audience("api.myapp.com", [])
/// ]
/// ```
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn subject(subject: String, others: List(String)) -> Claim {
  string_claim("sub", subject, others, core.InvalidSubject)
}

/// Creates a custom claim with a specific name and value that must match exactly during validation.
///
/// Use this for application-specific claims like roles, permissions, or custom
/// business logic requirements.
///
/// ## Example
/// ```gleam
/// // Add role-based authorization
/// let claims = [
///   claim.custom(
///     name: "role",
///     value: "admin",
///     encode: json.string,
///     decoder: decode.string
///   ),
///   claim.custom(
///     name: "permissions",
///     value: ["read", "write", "delete"],
///     encode: json.array(_, json.string),
///     decoder: decode.list(decode.string)
///   ),
///   claim.expires_at(max_age: duration.hours(1), leeway: duration.minutes(5))
/// ]
/// ```
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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)

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

fn validate(
  name: String,
  decoder: Decoder(a),
  on_error: fn(a) -> ParseError,
  validate: fn(a) -> Bool,
) -> fn(Dynamic) -> Result(Nil, ParseError) {
  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) -> ParseError,
) {
  let expected = [primary, ..others]
  let value = fn() { json.string(primary) }
  let verify =
    validate(name, decode.string, error(expected, _), list.contains(expected, _))

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

/// Decodes JWT numeric date values (seconds since Unix epoch) into timestamp objects.
///
/// Use this when manually processing JWT claims that contain timestamp values.
///
/// ## Example
/// ```gleam
/// // Manually decode an expiration claim
/// let exp_decoder = decode.field("exp", numeric_date_decoder())
/// ```
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn numeric_date_decoder() -> Decoder(Timestamp) {
  decode.map(decode.int, timestamp.from_unix_seconds)
}

/// Encodes a timestamp into the JWT numeric date format (seconds since Unix epoch).
///
/// Use this when manually creating JWT payloads with timestamp values.
///
/// ## Example
/// ```gleam
/// // Add custom timestamp to payload
/// let payload = [
///   #("custom_date", encode_numeric_date(my_timestamp))
/// ]
/// ```
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn encode_numeric_date(timestamp: Timestamp) {
  json.int(float.truncate(timestamp.to_unix_seconds(timestamp)))
}

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

/// Validates all claims against a JWT payload.
///
/// This function checks each claim in the provided list against the JWT payload,
/// returning the first validation error encountered or success if all claims pass.
///
/// 🚨 This is a low-level function used internally by ywt.
///
/// ## Example
/// ```gleam
/// let payload = // ... decoded JWT payload
///
/// let security_claims = [
///   expires_at(max_age: duration.hours(1), leeway: duration.minutes(5)),
///   issuer("https://auth.myapp.com", []),
///   audience("https://api.myapp.com", []),
///   custom("role", "admin", json.string, decode.string)
/// ]
///
/// case verify(payload, security_claims) {
///   Ok(Nil) -> {
///     // All claims validated successfully
///     process_authenticated_request(payload)
///   }
///   Error(TokenExpired(expired_at)) -> {
///     log.info("Token expired at: " <> timestamp.to_string(expired_at))
///     return_authentication_error()
///   }
///   Error(InvalidIssuer(expected, actual)) -> {
///     log.warning("Invalid issuer - expected: " <> expected <> ", got: " <> actual)
///     return_forbidden_error()
///   }
///   Error(error) -> {
///     log.error("Claim validation failed: " <> string.inspect(error))
///     return_bad_request_error()
///   }
/// }
/// ```
///
/// ⚠️ **Important:** This function only validates claims - it does not verify
/// cryptographic signatures. Use the full `decode()` function for complete JWT validation.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
pub fn verify(payload: Dynamic, claims: List(Claim)) -> Result(Nil, ParseError) {
  list.try_each(claims, verify_single_claim(payload, _))
}

/// Verify a single claim, handling required vs optional claims.
fn verify_single_claim(
  payload: Dynamic,
  claim: Claim,
) -> Result(Nil, ParseError) {
  let decoder = decode.at([claim.name], decode.dynamic)

  // TODO: MissingClaim and CalimDecodeError are 2 different cases
  case decode.run(payload, decoder) {
    Ok(claim_value) -> claim.verify(claim_value)
    // Optional claim missing is OK
    Error(_) if !claim.is_required -> Ok(Nil)
    Error(_) -> Error(core.MissingClaim(claim.name))
  }
}

/// Encode claims and application data into JSON.
///
/// This function merges application data with JWT claims into a single JSON object
/// for token encoding. When the same field appears in both lists, claims take
/// precedence to ensure security-critical values cannot be overridden.
///
/// 🚨 This is a low-level function used internally by ywt.
///
/// ## Example
/// ```gleam
/// let user_data = [
///   #("name", json.string("Alice")),
///   #("role", json.string("user")),
///   #("exp", json.int(1234567890))  // This will be overridden
/// ]
///
/// let security_claims = [
///   expires_at(max_age: duration.hours(1), leeway: duration.minutes(5)),
///   custom("role", "admin", json.string, decode.string)  // This takes precedence
/// ]
///
/// let payload = encode(security_claims, user_data)
/// // Result: {"name": "Alice", "role": "admin", "exp": <actual_expiry>}
/// ```
///
/// 💡 **Best Practice:** Place security-sensitive fields in claims rather than
/// data to validate them and ensure they cannot be accidentally overridden.
///
/// <div style="text-align: right;">
///      <a href="#">
///          <small>Back to top ↑</small>
///      </a>
/// </div>
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) {
    dict.insert(dict, claim.name, claim.value())
  })
  |> dict.to_list
  |> json.object
}