import std/json, std/logging, std/options, std/sequtils, std/strutils

import ../private/crypto
import ../private/encoding
import ../private/jsonutils

import ./claims
import ./joseheader
import ./jwa
import ./jwk
import ./jwssig

const VALID_SIGNATURE_ALGORITHMS = [
  HS256, HS384, HS512, RS256, RS384, RS512,
  ES256, ES384, ES512, PS256, PS384, PS512
]

type
  InvalidSignature* = object of CatchableError

  JWS* = object
    payloadB64: string
    signatures: seq[JwsSignature]
    compactSerialization: Option[string]

func payloadB64*(jws: JWS): string = jws.payloadB64
func signatures*(jws: JWS): seq[JwsSignature] = jws.signatures

func `[]`*(jws: JWS, idx: int): JwsSignature = jws.signatures[idx]
func len*(jws: JWS): int = jws.signatures.len

func compactSerialization*(jws: JWS): Option[string] =
  jws.compactSerialization

func jsonSerialization*(jws: JWS, flatten = false): string =
  if flatten and jws.len != 1:
    raise newException(ValueError, "A JWS must have exactly one signature " &
      "in order to be serialized in the Flattened JSON Serialization " &
      "Syntax. This JWS has " & $(jws.len) & " signatures.")

  elif flatten:
    return $(%*{
      "payload": jws.payloadB64,
      "protected": jws[0].protected.rawB64,
      "header": jws[0].header.rawB64,
      "signature": jws[0].signatureB64
    })

  else:
    return $(%*{
      "payload": jws.payloadB64,
      "signatures": jws.signatures
    })

func `$`*(jws: JWS): string =
  if jws.len == 1: return jws.compactSerialization.get
  else: return jws.jsonSerialization

proc validateHeader(h: JoseHeader) =
  if not h.alg.isSome:
    raise newException(ValueError, "missing required alg header parameter")

  if not VALID_SIGNATURE_ALGORITHMS.contains(h.alg.get):
    raise newException(ValueError,
      $(h.alg) & " is not a valid signature algorithm.")

proc initJWS*(b64: string): JWS =
  let parts = b64.split('.')

  if parts.len != 3:
    raise newException(ValueError,
      "invalid JWS, expected three parts but only found " & $parts.len)

  debug "Base64url-encoded payload is:\n\t" & parts[1]

  result = JWS(
    payloadB64: parts[1],
    signatures: @[ initJwsSignature(
      protected = initJoseHeader(parts[0]),
      header = initJoseHeader(alg = none[JwtAlgorithm]()),
      signatureB64 = parts[2]) ])
  result.compactSerialization = some(b64)

proc initJWS*(n: JsonNode): JWS =
  if not n.hasKey("payload"):
    raise newException(ValueError, "invalid JWS: missing 'payload'")

  if n.hasKey("signatures"):
    result = JWS(
      payloadB64: n.reqStrVal("payload"),
      signatures: n["signatures"].getElems.mapIt(initJwsSignature(it)))
    if result.len == 0:
      raise newException(ValueError, "Invalid JWS: no signatures.")
  else:
    result = JWS(
      payloadB64: n.reqStrVal("payload"),
      signatures: @[initJwsSignature(n)])

proc sign*(payload: string, header: JoseHeader, key: JWK, payloadIsB64Encoded = false): JWS =
  ## Create a JWS signing the given payload, which can be any arbitrary data.
  ## The returned JWS can be serialized with the compact serialization. All
  ## header data will be protected.

  validateHeader(header)

  let payloadB64 = if payloadIsB64Encoded: payload
                   else: b64UrlEncode(payload)

  let valueToSign = $header & "." & payloadB64
  let signatureB64 =
    b64UrlEncode(computeSignature(valueToSign, header.alg.get, key))

  debug "Signed a value:\n\ttoSign   : " & valueToSign &
    "\n\tsignature: " & signatureB64

  result = JWS(
    payloadB64: payloadB64,
    signatures: @[
      initJwsSignature(
        protected = header,
        signatureB64 = signatureB64)])

  result.compactSerialization =
    some($header & "." & result.payloadB64 & "." & signatureB64)

proc sign*(claims: JwtClaims, header: JoseHeader, key: JWK): JWS =
  ## Create a JWS signing a set of JWT claims. The returned JWS will be a valid
  ## JWT.
  sign($claims, header, key, payloadIsB64Encoded = true)

proc sign*(
    payload: string,
    unprotected: JoseHeader,
    protected: JoseHeader,
    key: JWK,
    payloadIsB64Encoded = false): JWS =
  ## Create a JWS signing the given payload, which can be arbitrary data. The
  ## returned JWS can be serialized with the compact serialization.

  let combinedHeader = combine(unprotected, protected)
  validateHeader(combinedHeader)

  let payloadB64 = if payloadIsB64Encoded: payload
                   else: b64UrlEncode(payload)

  let valueToSign = $protected & "." & payloadB64
  let signatureB64 = b64UrlEncode(
    computeSignature(valueToSign, combinedHeader.alg.get, key))

  debug "Signed a value:\n\ttoSign   : " & valueToSign &
    "\n\tsignature: " & signatureB64

  result = JWS(
    payloadB64: payloadB64,
    signatures: @[
      initJwsSignature(
        protected = protected,
        header = unprotected,
        signatureB64 = signatureB64)])

  result.compactSerialization =
    some($combinedHeader & "." & result.payloadB64 & "." & signatureB64)

proc sign*(
    jws: JWS,
    unprotected: JoseHeader,
    protected: JoseHeader,
    key: JWK): JWS =
  ## Create a new JWS signing the payload again with the newly supplied header
  ## and key. The resulting JWS token will have multiple signatures and cannot
  ## be serialized with the compact serialization.

  let combinedHeader = combine(unprotected, protected)
  validateHeader(combinedHeader)

  let valueToSign = $protected & "." & jws.payloadB64
  let signatureB64 = b64UrlEncode(
    computeSignature(valueToSign, combinedHeader.alg.get, key))

  debug "Signed a value:\n\ttoSign   : " & valueToSign &
    "\n\tsignature: " & signatureB64

  result = JWS(
    payloadB64: jws.payloadB64,
    signatures: jws.signatures &
      @[initJwsSignature(
        protected = protected,
        header = unprotected,
        signatureB64 = signatureB64)])

  result.compactSerialization = none[string]()

proc validate*(jws: JWS, alg: JwtAlgorithm, key: JWK, sigIdx = 0) =

  if jws.len == 0:
    raise newException(InvalidSignature, "JWS has no signature.")

  if jws.len < sigIdx+1:
    raise newException(InvalidSignature,
      "No signature at index " & $sigIdx & ". There are only " & $jws.len &
      " signatures on this JWS.")

  let combinedHeader = combine(jws[sigIdx].header, jws[sigIdx].protected)

  if combinedHeader.alg.isNone:
    raise newException(InvalidSignature, "JWS header has no value for 'alg'.")

  if alg != combinedHeader.alg.get:
    raise newException(InvalidSignature,
      "JWS alg (" & $(combinedHeader.alg.get) & ") does not match the " &
      "requested alg (" & $alg & ")")

  if not VALID_SIGNATURE_ALGORITHMS.contains(alg):
    raise newException(InvalidSignature,
      "'" & $alg & "' is not a valid signing algorithm.")

  let payload = $(jws[sigIdx].protected) & "." & jws.payloadB64

  debug "Verifying a JWS:\n\tpayload  : " & payload &
    "\n\tsignature: " & jws[sigIdx].signatureB64


  if not verifySignature(
    payload = payload,
    signature = b64UrlDecode(jws[sigIdx].signatureB64),
    alg = alg,
    key = key):

    raise newException(InvalidSignature, "failed to verify signature value.")


proc tryValidate*(jws: JWS, alg: JwtAlgorithm, key: JWK): bool =
  try:
    jws.validate(alg, key)
    return true
  except: return false