import std/[json, options, unittest]
import jwt_full
import buffoonery/auth

const testIssuer = "https://issuer.example"

proc keyNode(kid = "original"): JsonNode =
  %*{"kty": "oct", "alg": "HS256", "kid": kid,
     "k": "c2VjcmV0", "metadata": {"labels": ["original"]}}

proc newContext(): ApiAuthContext =
  initApiAuthContext("example.com", "test", @[testIssuer], testIssuer,
    @[testIssuer], "original", @[initJwk(keyNode())])

proc tokenFor(kid = "original", issuer = testIssuer): JWT =
  # Lookup only needs the header and issuer; signing is tested separately.
  createSignedJwt(initJoseHeader(%*{"alg": "HS256", "kid": kid}),
    initJwtClaims(%*{"iss": issuer}), initJwk(keyNode(kid)))

when compileOption("threads"):
  type WorkerInput = tuple[ctx: ApiAuthContext, ok: ptr bool]

  proc useContext(input: WorkerInput) {.thread.} =
    var ctx = input.ctx
    try:
      ctx.validAudiences[0] = "worker"
      let key = ctx.findSigningKey(tokenFor(), false)
      key["metadata"].get["labels"][0].str = "worker"
      ctx.addSigningKeys(testIssuer, @[initJwk(keyNode("worker"))])
      input.ok[] = ctx.findSigningKey(tokenFor("worker"), false).kid.get == "worker"
    except CatchableError:
      input.ok[] = false

suite "auth context value semantics":
  test "copies own their configuration and signing-key cache":
    var original = newContext()
    var copied = original
    copied.appDomain[0] = 'E'
    copied.validAudiences[0][0] = 'H'
    copied.trustedIssuers.add "https://other.example"
    copied.addSigningKeys(testIssuer, @[initJwk(keyNode("replacement"))])

    check original.appDomain == "example.com"
    check original.validAudiences == @[testIssuer]
    check original.trustedIssuers == @[testIssuer]
    check original.findSigningKey(tokenFor(), false).kid.get == "original"
    check copied.findSigningKey(tokenFor("replacement"), false).kid.get == "replacement"
    expect AuthError:
      discard copied.findSigningKey(tokenFor(), false)
    expect AuthError:
      discard original.findSigningKey(tokenFor("replacement"), false)

  test "constructor and added keys do not retain caller-owned JSON":
    let node = keyNode()
    var ctx = initApiAuthContext("example.com", "test", @[testIssuer],
      testIssuer, @[testIssuer], "original", @[initJwk(node)])
    node["metadata"]["labels"][0].str = "changed"
    check ctx.findSigningKey(tokenFor(), false)["metadata"].get["labels"][0].getStr == "original"

    let added = keyNode("added")
    ctx.addSigningKeys(testIssuer, @[initJwk(added)])
    added["metadata"]["labels"][0].str = "changed"
    check ctx.findSigningKey(tokenFor("added"), false)["metadata"].get["labels"][0].getStr == "original"

  test "returned key metadata is independent of both context copies":
    var original = newContext()
    var copied = original
    let key = copied.findSigningKey(tokenFor(), false)
    key["metadata"].get["labels"][0].str = "changed"
    check original.findSigningKey(tokenFor(), false)["metadata"].get["labels"][0].getStr == "original"
    check copied.findSigningKey(tokenFor(), false)["metadata"].get["labels"][0].getStr == "original"

  test "default contexts accept keys and reject missing keys without fetching":
    var ctx: ApiAuthContext
    expect AuthError:
      discard ctx.findSigningKey(tokenFor(), false)
    ctx.addSigningKeys(testIssuer, @[initJwk(keyNode())])
    check ctx.findSigningKey(tokenFor(), false).kid.get == "original"

  test "invalid key updates leave the existing cache intact":
    var ctx = newContext()
    let invalid = keyNode("invalid")
    invalid.delete("alg")
    expect AuthError:
      ctx.addSigningKeys(testIssuer, @[initJwk(invalid)])
    check ctx.findSigningKey(tokenFor(), false).kid.get == "original"

  test "immutable contexts can sign tokens that mutable copies validate":
    let original = newContext()
    let token = original.newApiAccessToken("user")
    var copied = original
    copied.validateJWT(token)
    check token.claims.sub.get == "user"

  test "cached keys preserve RSA and EC variants":
    var ctx: ApiAuthContext
    let nodes = @[
      %*{"kty": "RSA", "alg": "RS256", "kid": "rsa-public",
         "n": "AQAB", "e": "AQAB"},
      %*{"kty": "RSA", "alg": "RS256", "kid": "rsa-private",
         "n": "AQAB", "e": "AQAB", "d": "Ag", "p": "Aw", "q": "BQ",
         "oth": [{"r": "Bw", "d": "Ag", "t": "Aw"}]},
      %*{"kty": "EC", "alg": "ES256", "kid": "ec-public",
         "crv": "P-256", "x": "AQ", "y": "Ag"},
      %*{"kty": "EC", "alg": "ES256", "kid": "ec-private",
         "crv": "P-256", "x": "AQ", "y": "Ag", "d": "Aw"}]
    for node in nodes:
      let key = initJwk(node)
      ctx.addSigningKeys(testIssuer, @[key])
      let restored = ctx.findSigningKey(tokenFor(key.kid.get), false)
      check restored.keyKind == key.keyKind
      for name, value in node.pairs:
        check restored[name].get == value
      case key.keyKind
      of RsaPublic: check restored.rsaPub == key.rsaPub
      of RsaPrivate: check restored.rsaPrv == key.rsaPrv
      of EcPublic: check restored.ecPub == key.ecPub
      of EcPrivate: check restored.ecPrv == key.ecPrv
      of Octet: discard

  when compileOption("threads"):
    test "workers mutate independent context copies":
      var original = newContext()
      var workers: array[2, Thread[WorkerInput]]
      var succeeded: array[2, bool]
      for i in 0 ..< workers.len:
        createThread(workers[i], useContext, (original, addr succeeded[i]))
      joinThreads(workers)
      check succeeded == [true, true]
      check original.validAudiences == @[testIssuer]
      check original.findSigningKey(tokenFor(), false)["metadata"].get["labels"][0].getStr == "original"