import gleam/bit_array
import gleam/bool
import gleam/dynamic/decode.{type Decoder, type Dynamic}
import gleam/json.{type Json}
import gleam/list
import gleam/option.{type Option, None, Some}
import gleam/result
import gleam/string
import ywt/algorithm.{type Algorithm}
import ywt/claim.{type Claim}
import ywt/internal/base64url
import ywt/internal/core.{type Error}
import ywt/internal/jose_json
import ywt/sign_key.{type SignKey}
import ywt/verify_key.{type VerifyKey}

// -- HEADER -------------------------------------------------------------------

/// Header represents the values we are interested in in a JOSE (JWT) header.
type Header {
  Header(alg: Algorithm, kid: Option(String), b64: Bool)
}

fn header_decoder() -> Decoder(Header) {
  use alg <- decode.field("alg", algorithm.decoder())
  use kid <- decode.optional_field("kid", None, decode.map(decode.string, Some))
  use b64 <- decode.optional_field("b64", True, decode.bool)

  decode.success(Header(alg:, kid:, b64:))
}

fn encode_header(
  alg: Algorithm,
  kid: Option(String),
  claims: List(Claim),
) -> Json {
  let fields = [
    #("alg", algorithm.to_json(alg)),
  ]
  let fields = case kid {
    Some(kid) -> [#("kid", json.string(kid)), ..fields]
    None -> fields
  }

  claim.encode_header(claims, fields)
}

// -- PARSE---------------------------------------------------------------------

/// Extracts payload from a JWT without verifying its signature or claims.
pub fn decode_unsafely_without_validation(
  jwt: String,
  payload_decoder: Decoder(payload),
) -> Result(payload, Error) {
  use #(_, raw_payload, _) <- result.try(split(jwt))

  use dynamic_payload <- result.try(part(
    raw_payload,
    core.InvalidPayloadEncoding,
    core.InvalidPayloadJson,
  ))

  decode.run(dynamic_payload, payload_decoder)
  |> result.map_error(core.PayloadDecodingError)
}

/// Verifies a JWT signature and validates all claims, returning the decoded payload if successful.
pub fn decode(
  jwt jwt: String,
  using decoder: Decoder(payload),
  claims claims: List(Claim),
  keys keys: List(VerifyKey),
  verify verify: fn(BitArray, BitArray, VerifyKey, fn(Bool) -> result) -> result,
  resolve resolve: fn(Result(payload, Error)) -> result,
) -> result {
  let continuation_or_error = {
    // 1. Verify that the JWT contains at least one period ('.') character.
    // 6. Determine whether the JWT is a JWS or a JWE using any of the
    //    methods described in Section 9 of [JWE].
    use #(raw_header, raw_payload, raw_signature) <- result.try(split(jwt))

    // 2. Let the Encoded JOSE Header be the portion of the JWT before the first
    //    period ('.') character.
    use dynamic_header <- result.try(part(
      raw_header,
      core.InvalidHeaderEncoding,
      core.InvalidHeaderJson,
    ))

    use parsed_header <- result.try(
      decode.run(dynamic_header, header_decoder())
      |> result.map_error(core.HeaderDecodingError),
    )
    // RFC 7797 Section 3 defines `b64`; `false` means the JWS Payload is
    // represented without base64url encoding. Section 7 updates JWT to say
    // JWTs MUST NOT use `b64` with a `false` value.
    use Nil <- result.try(reject_unencoded_payload_header(parsed_header))
    use Nil <- result.try(reject_critical_header(dynamic_header))

    // 6. Base64url-decode the encoded representation of the JWS Payload,
    //    following the restriction that no line breaks, whitespace, or
    //    other additional characters have been used.
    use dynamic_payload <- result.try(part(
      raw_payload,
      core.InvalidPayloadEncoding,
      core.InvalidPayloadJson,
    ))

    // 10. Verify that the resulting octet sequence is a UTF-8-encoded
    //     representation of a completely valid JSON object conforming to
    //     RFC 7159 [RFC7159]; let the JWT Claims Set be this JSON object.
    use Nil <- result.try(claim.verify_with_header(
      dynamic_header,
      dynamic_payload,
      claims,
    ))

    // Decode the payload using the provided decoder
    use parsed_payload <- result.try(
      decode.run(dynamic_payload, decoder)
      |> result.map_error(core.PayloadDecodingError),
    )

    // 7. Base64url-decode the encoded representation of the JWS Signature,
    //    following the restriction that no line breaks, whitespace, or
    //    other additional characters have been used.
    use signature <- result.try(
      base64url.decode(raw_signature)
      |> result.replace_error(core.InvalidSignatureEncoding),
    )

    let matching_keys = list.filter_map(keys, key_for_header(_, parsed_header))
    use <- bool.guard(
      when: list.is_empty(matching_keys),
      return: Error(core.NoMatchingKey),
    )

    // 8. Validate the JWS Signature against the JWS Signing Input
    //    ASCII(BASE64URL(UTF8(JWS Protected Header)) || '.' ||
    //    BASE64URL(JWS Payload)) in the manner defined for the algorithm
    //    being used, which MUST be accurately represented by the value of
    //    the "alg" (algorithm) Header Parameter, which MUST be present.
    //    See Section 10.6 for security considerations on algorithm
    //    validation.  Record whether the validation succeeded or not.
    let verify_message = <<raw_header:utf8, ".", raw_payload:utf8>>

    let verify = fn(key: VerifyKey, next: fn(Bool) -> result) -> result {
      verify(verify_message, signature, key, next)
    }

    let on_result = fn(result: Bool) -> result {
      case result {
        True -> resolve(Ok(parsed_payload))
        False -> resolve(Error(core.InvalidSignature))
      }
    }

    Ok(fn() { verify_loop(matching_keys, False, verify, on_result) })
  }

  case continuation_or_error {
    Ok(continuation) -> continuation()
    Error(error) -> resolve(Error(error))
  }
}

fn reject_critical_header(header: Dynamic) -> Result(Nil, Error) {
  // RFC 7515 Section 4.1.11: "crit" marks extensions that recipients MUST
  // understand and process. ywt supports no critical JOSE extensions, so any
  // "crit" parameter is rejected instead of ignored.
  let decoder = decode.optionally_at(["crit"], Nil, decode.failure(Nil, "crit"))

  decode.run(header, decoder)
  |> result.replace_error(core.UnsupportedCriticalHeader)
}

fn reject_unencoded_payload_header(header: Header) -> Result(Nil, Error) {
  case header.b64 {
    True -> Ok(Nil)
    False -> Error(core.UnsupportedUnencodedPayload)
  }
}

fn verify_loop(
  keys: List(VerifyKey),
  result: Bool,
  verify verify: fn(VerifyKey, fn(Bool) -> result) -> result,
  then on_result: fn(Bool) -> result,
) -> result {
  case keys {
    [] -> on_result(result)
    [key, ..keys] -> {
      use is_valid <- verify(key)
      let result = result || is_valid
      verify_loop(keys, result, verify, on_result)
    }
  }
}

/// Split a JWT/JWS into its 3 parts - header, payload, and signature.
fn split(jwt: String) -> Result(#(String, String, String), Error) {
  case string.split(jwt, on: ".") {
    [header, payload, signature] -> Ok(#(header, payload, signature))
    _ -> Error(core.MalformedToken)
  }
}

/// Decode a raw base64-encoded part into a Gleam value with specific error types.
fn part(
  part: String,
  encoding_error: Error,
  json_error: fn(json.DecodeError) -> Error,
) -> Result(Dynamic, Error) {
  // 3. Base64url decode the Encoded JOSE Header following the
  //    restriction that no line breaks, whitespace, or other additional
  //    characters have been used.
  use decoded <- result.try(
    base64url.decode(part)
    |> result.replace_error(encoding_error),
  )

  // 4. Verify that the resulting octet sequence is a UTF-8-encoded
  //    representation of a completely valid JSON object conforming to
  //    RFC 7159 [RFC7159]; let the JOSE Header be this JSON object.
  jose_json.parse_bits(decoded, decode.dynamic)
  |> result.map_error(json_error)
}

/// Pairs up a key candidate with the JWT header, specializing it to the header's
/// algorithm. Returns `Ok(specialized_key)` if the key is compatible with the
/// header's kid and alg, or `Error(Nil)` if it should be skipped.
fn key_for_header(key: VerifyKey, header: Header) -> Result(VerifyKey, Nil) {
  // 5. Verify that the resulting JOSE Header includes only parameters
  //    and values whose syntax and semantics are both understood and
  //    supported or that are specified as being ignored when not
  //    understood.
  //
  // The producer SHOULD include sufficient information in the Header
  // Parameters to identify the key used, unless the application uses
  // another means or convention to determine the key used.  Validation of
  // the signature or MAC fails when the algorithm used requires a key
  // (which is true of all algorithms except for "none") and the key used
  // cannot be determined.
  //
  // Kid matching: if the header has a kid, the key must have the same kid.
  case header.kid, verify_key.id(key) {
    Some(header_kid), Ok(key_id) if header_kid != key_id -> Error(Nil)
    Some(_), Error(_) -> Error(Nil)
    // Algorithm matching and specialization (handles no-alg RSA keys).
    _, _ -> verify_key.for_algorithm(key, header.alg)
  }
}

// -- SIGN ---------------------------------------------------------------------

/// Creates a signed JWT containing the specified payload and claims.
pub fn encode(
  payload payload: List(#(String, Json)),
  claims claims: List(Claim),
  key key: SignKey,
  sign sign: fn(BitArray, SignKey, fn(BitArray) -> String) -> result,
) -> result {
  let alg = sign_key.algorithm(key)
  let header =
    encode_json(encode_header(alg, option.from_result(sign_key.id(key)), claims))
  let payload = encode_json(claim.encode(claims, payload))

  let verify_message = header <> "." <> payload

  use signed <- sign(<<verify_message:utf8>>, key)

  let signature = bit_array.base64_url_encode(signed, False)

  verify_message <> "." <> signature
}

fn encode_json(part: Json) -> String {
  bit_array.base64_url_encode(<<json.to_string(part):utf8>>, False)
}