local jwt = require "resty.jwt"
local JWKS = require "kong.plugins.auth0-validator.jwks"
local lrucache = require "resty.lrucache"
local resty_md5 = require "resty.md5"
local str = require "resty.string"

local Auth0ValidatorHandler = {
  VERSION = "1.0.0",
  -- higher than o11y-proxy
  PRIORITY = 1010,
}

-- JWKS instance is created at module load time, which in Kong means per worker.
-- Each worker maintains its own JWKS instance; it is not shared across workers or nodes.
local jwks_instance = JWKS.new()

-- Token Cache (per-worker)
-- Cache validation results (bool) keyed by MD5(token).
-- This lrucache lives in the worker's Lua VM and is not shared across workers or the Kong cluster,
-- so different workers may redundantly fetch JWKS keys on first use.
local token_cache_size = tonumber(os.getenv("AUTH0_TOKEN_CACHE_SIZE")) or 2000
local token_cache = lrucache.new(token_cache_size)

local function get_cache_key(token)
  local md5 = resty_md5:new()
  md5:update(token)
  local digest = md5:final()
  return str.to_hex(digest)
end

function Auth0ValidatorHandler:access(conf)
  local token_header = conf.token_header or "Authorization"
  kong.log.debug("Auth0ValidatorHandler: access phase started")
  local auth_header = kong.request.get_header(token_header)

  if auth_header then
    auth_header = auth_header:match("^%s*(.-)%s*$")
  end
  if not auth_header then
    kong.log.warn("Missing authorization header")
    return kong.response.exit(401, { message = "Missing authorization header" })
  end

  if string.sub(auth_header, 1, 7):lower() == "bearer " then
    token = string.sub(auth_header, 8)
  else
    kong.log.err("Invalid authorization header format. Header: ", auth_header)
    return kong.response.exit(401, { message = "Invalid authorization header format" })
  end

  if not token or token == "" then
    kong.log.err("Invalid token format. Header: ", auth_header)
    return kong.response.exit(401, { message = "Invalid token format" })
  end

  -- Check cache
  local cache_key = get_cache_key(token)
  local cached_valid = token_cache:get(cache_key)
  if cached_valid then
    kong.log.debug("Auth0ValidatorHandler: Token found in cache")
    return
  end

  -- Decode without verifying to get key id
  local jwt_obj = jwt:load_jwt(token)
  if not jwt_obj.valid then
    kong.log.warn("Auth0ValidatorHandler: Invalid JWT structure")
    return kong.response.exit(401, { message = "Invalid JWT structure" })
  end

  local kid = jwt_obj.header.kid
  if not kid then
    kong.log.warn("Auth0ValidatorHandler: No kid in JWT header")
    return kong.response.exit(401, { message = "No kid in JWT header" })
  end

  -- Validate that the JWKS URL is well-formed and matches the expected issuer host
  local function validate_jwks_url(jwks_url, issuer)
    if type(jwks_url) ~= "string" or type(issuer) ~= "string" then
      return nil, "JWKS URL or issuer not a string"
    end

    -- Enforce HTTPS for JWKS fetching
    if string.sub(jwks_url, 1, 8) ~= "https://" then
      return nil, "JWKS URL must use https scheme"
    end

    -- Extract hosts
    local issuer_host = string.match(issuer, "^https?://([^/]+)")
    local jwks_host = string.match(jwks_url, "^https?://([^/]+)")

    if not issuer_host or not jwks_host then
      return nil, "Unable to parse host from issuer or JWKS URL"
    end

    if issuer_host ~= jwks_host then
      return nil, "JWKS URL host does not match issuer host"
    end

    -- Ensure we are calling the well-known JWKS endpoint
    if not string.find(jwks_url, "/.well%-known/jwks%.json", 1, false) then
      return nil, "JWKS URL path is not the expected well-known JWKS path"
    end

    return true
  end

  local issuer = conf.issuer
  if string.sub(issuer, -1) ~= "/" then
    issuer = issuer .. "/"
  end

  local jwks_url = conf.jwks_custom_url
  if not jwks_url then
    jwks_url = issuer .. ".well-known/jwks.json"
  end

  local ok, jwks_err = validate_jwks_url(jwks_url, issuer)
  if not ok then
    kong.log.crit("Auth0ValidatorHandler: Invalid JWKS URL configuration: ", jwks_err, " (URL: ", jwks_url, ", issuer: ", issuer, ")")
    return kong.response.exit(500, { message = "Internal Server Error" })
  end
  kong.log.debug("Auth0ValidatorHandler: key ID [", kid, "] fetching from [", jwks_url, "]")
  local pub_key, err = jwks_instance:get_public_key(kid, jwks_url)
  if not pub_key then
    kong.log.err("Failed to get public key: ", err)
    return kong.response.exit(401, { message = "Unable to verify token signature" })
  end

  local verified_jwt = jwt:verify(pub_key, token)
  if not verified_jwt.verified then
    kong.log.warn("Auth0ValidatorHandler: Invalid token signature: ", verified_jwt.reason)
    return kong.response.exit(401, { message = "Invalid token: " .. (verified_jwt.reason or "signature mismatch") })
  end

  -- Check Claims
  local claims = verified_jwt.payload

  -- 1. Expiration
  if not claims.exp then
    kong.log.warn("Auth0ValidatorHandler: Token missing expiration claim")
    return kong.response.exit(401, { message = "Token missing expiration claim" })
  end

  local current_time = ngx.time()
  if claims.exp <= current_time then
    kong.log.warn("Auth0ValidatorHandler: Token has expired. Exp: ", claims.exp, ", Current: ", current_time)
    return kong.response.exit(401, { message = "Token has expired" })
  end

  -- 2. Issuer
  if claims.iss ~= issuer then
    kong.log.warn("Auth0ValidatorHandler: Invalid issuer. Expected: ", issuer, ", Got: ", claims.iss)
     return kong.response.exit(401, { message = "Invalid issuer" })
  end

  -- 3. Audience
  local aud = claims.aud
  local aud_valid = false
  if type(aud) == "string" then
    if aud == conf.audience then aud_valid = true end
  elseif type(aud) == "table" then
    for _, a in ipairs(aud) do
      if a == conf.audience then
        aud_valid = true
        break
      end
    end
  end

  if not aud_valid then
    kong.log.warn("Auth0ValidatorHandler: Invalid audience. Expected: ", conf.audience, ", Got: ", claims.aud)
    return kong.response.exit(401, { message = "Invalid audience" })
  end

  -- Token is valid, update cache
  -- claims.exp is guaranteed to exist at this point
  local ttl = claims.exp - ngx.time()
  if ttl > 0 then
    token_cache:set(cache_key, true, ttl)
    kong.log.debug("Auth0ValidatorHandler: Token verified and cached")
  end
end

return Auth0ValidatorHandler