local http_client = require "kong.plugins.o11y-proxy.http_client"
local cjson = require "cjson.safe"
local lrucache = require "resty.lrucache"

local JWKS = {}
JWKS.__index = JWKS

-- Configurable cache size for JWKS keys. Defaults to 100 entries if not set.
local KEY_CACHE_SIZE = tonumber(os.getenv("JWKS_KEY_CACHE_SIZE")) or 100

-- Configurable TTL for JWKS keys in seconds. Defaults to 1 hour.
-- Auth0 keys rotate infrequently (manual or ~10 year expiration), but 1 hour provides
-- a good balance between performance and picking up key rotations promptly.
-- Auth0 recommends ~1 hour.
local KEY_CACHE_TTL = tonumber(os.getenv("JWKS_KEY_CACHE_TTL")) or 3600 -- 1 hour
-- Cache for Public Keys (PEM format)
-- Key: "kid", Value: PEM string
local key_cache = lrucache.new(KEY_CACHE_SIZE)

-- Helper: decode base64url-encoded string to binary
local function base64url_to_binary(input)
  if type(input) ~= "string" then
    return nil, "input must be a string"
  end

  -- Convert from base64url to standard base64
  input = input:gsub("-", "+"):gsub("_", "/")
  local padding = #input % 4
  if padding == 2 then
    input = input .. "=="
  elseif padding == 3 then
    input = input .. "="
  elseif padding ~= 0 then
    return nil, "invalid base64url padding"
  end

  local decoded = ngx.decode_base64(input)
  if not decoded then
    return nil, "failed to base64url-decode input"
  end

  return decoded
end

-- Helper: encode length in ASN.1 DER format
local function asn1_length(len)
  if len < 128 then
    return string.char(len)
  elseif len < 256 then
    return string.char(0x81, len)
  else
    return string.char(0x82, math.floor(len / 256), len % 256)
  end
end

-- Helper: encode an integer in ASN.1 DER format
local function asn1_integer(bytes)
  -- Add leading zero if high bit is set (to indicate positive number)
  if bytes:byte(1) >= 128 then
    bytes = string.char(0) .. bytes
  end
  return string.char(0x02) .. asn1_length(#bytes) .. bytes
end

-- Helper: encode a sequence in ASN.1 DER format
local function asn1_sequence(content)
  return string.char(0x30) .. asn1_length(#content) .. content
end

-- Helper: encode a bit string in ASN.1 DER format
local function asn1_bitstring(content)
  -- Prepend with 0x00 to indicate no unused bits
  local bs = string.char(0) .. content
  return string.char(0x03) .. asn1_length(#bs) .. bs
end

-- Convert JWK to PEM
local function jwk_to_pem(jwk)
  if jwk.kty ~= "RSA" then
    return nil, "Unsupported key type: " .. (jwk.kty or "nil")
  end

  -- Prefer x5c if available to preserve existing behavior
  if jwk.x5c and type(jwk.x5c) == "table" and jwk.x5c[1] then
    -- adding certificate header and footer
    return "-----BEGIN CERTIFICATE-----\n" ..
        jwk.x5c[1] ..
        "\n-----END CERTIFICATE-----"
  end

  -- Construct RSA public key from modulus (n) and exponent (e)
  if not (jwk.n and jwk.e) then
    return nil, "JWK missing required n or e fields"
  end

  local n_bin, err = base64url_to_binary(jwk.n)
  if not n_bin then
    return nil, "Failed to decode JWK modulus (n): " .. (err or "unknown error")
  end

  local e_bin, err2 = base64url_to_binary(jwk.e)
  if not e_bin then
    return nil, "Failed to decode JWK exponent (e): " .. (err2 or "unknown error")
  end

  -- Build RSAPublicKey ::= SEQUENCE { modulus INTEGER, publicExponent INTEGER }
  local rsa_public_key = asn1_sequence(asn1_integer(n_bin) .. asn1_integer(e_bin))

  -- RSA algorithm OID: 1.2.840.113549.1.1.1 with NULL parameters
  local rsa_oid = string.char(
    0x06, 0x09, -- OID tag and length
    0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x01, -- rsaEncryption OID
    0x05, 0x00 -- NULL parameters
  )
  local algorithm_identifier = asn1_sequence(rsa_oid)

  -- Build SubjectPublicKeyInfo ::= SEQUENCE { algorithm, subjectPublicKey BIT STRING }
  local subject_public_key_info = asn1_sequence(algorithm_identifier .. asn1_bitstring(rsa_public_key))

  -- Convert to PEM format
  local der_b64 = ngx.encode_base64(subject_public_key_info)

  -- Format PEM with 64-character lines
  local pem_lines = {"-----BEGIN PUBLIC KEY-----"}
  for i = 1, #der_b64, 64 do
    table.insert(pem_lines, der_b64:sub(i, i + 63))
  end
  table.insert(pem_lines, "-----END PUBLIC KEY-----")

  return table.concat(pem_lines, "\n")
end

function JWKS.new()
  return setmetatable({}, JWKS)
end

-- for unit test
JWKS.jwk_to_pem = jwk_to_pem

function JWKS:get_public_key(kid, jwks_url)
  -- 1. Check cache
  local cached = key_cache:get(kid)
  if cached then
    kong.log.debug("JWKS: Public key found in cache for kid: ", kid)
    return cached
  end

  kong.log.debug("JWKS: Fetching keys from: ", jwks_url)
  -- 2. Fetch JWKS
  local client = http_client.new()
  local res, err = client:request_uri(jwks_url, {
      method = "GET",
      ssl_verify = true -- ensure JWKS is fetched with TLS certificate verification enabled
  })

  if not res or res.status ~= 200 then
     return nil, "Failed to fetch JWKS: " .. (err or (res and res.status or "unknown"))
  end

  local body, parse_err = cjson.decode(res.body)
  if not body then
    return nil, "Failed to parse JWKS response: " .. (parse_err or "invalid JSON")
  end
  if not body.keys then
    return nil, "Invalid JWKS response: missing 'keys' field"
  end

  -- 3. Parse and Cache ALL keys
  for _, key in ipairs(body.keys) do
     if key.kid then
        local pem, err = jwk_to_pem(key)
        if pem then
          key_cache:set(key.kid, pem, KEY_CACHE_TTL)
        else
          kong.log.warn("Failed to convert JWK to PEM for kid " .. key.kid .. ": " .. (err or "unknown"))
        end
     end
  end

  kong.log.debug("JWKS: Successfully fetched and parsed JWKS")
  -- 4. Return requested key
  local key = key_cache:get(kid)
  if not key then
    kong.log.warn("JWKS: kid ", kid, " not found in fetched JWKS")
    return nil, "Key with kid " .. kid .. " not found in JWKS"
  end
  return key
end

return JWKS