--- IRSA (IAM Roles for Service Accounts) credentials fetcher for Kong
-- Fetches AWS credentials using Web Identity Token from Kubernetes ServiceAccount
-- Based on AWS_WEB_IDENTITY_TOKEN_FILE environment variable

local http_client = require "kong.plugins.o11y-proxy.http_client"
local resty_lock = require "resty.lock"
local lrucache = require "resty.lrucache"
local utils = require "kong.plugins.o11y-proxy.utils"

local IRSA = {}
IRSA.__index = IRSA

local DEFAULT_STS_ENDPOINT = "https://sts.amazonaws.com/"

-- Read environment variables at module load time (in master process)
-- This is the recommended approach for Kong plugins
-- Reference: https://github.com/Kong/kong/issues/7083#issuecomment-826197079
-- In Kong/OpenResty, os.getenv() only works in the master process during module loading
-- After that, worker processes can't access environment variables unless declared in nginx config
local ENV_AWS_WEB_IDENTITY_TOKEN_FILE = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
local ENV_AWS_ROLE_ARN = os.getenv("AWS_ROLE_ARN")
local ENV_AWS_ROLE_SESSION_NAME = os.getenv("AWS_ROLE_SESSION_NAME")
local ENV_AWS_REGION = os.getenv("AWS_REGION")
local ENV_AWS_DEFAULT_REGION = os.getenv("AWS_DEFAULT_REGION")

-- Use LRU cache for IRSA credentials
-- Single entry cache with 1 hour TTL, each worker has its own cache
local CACHE_KEY = "irsa_credentials"
local credentials_cache, cache_err = lrucache.new(10)  -- Small cache, only need 1 entry
if not credentials_cache then
  error("Failed to create lrucache: " .. (cache_err or "unknown error"))
end

--- Read token from file
-- @param token_file Path to the token file
-- @return string Token content or nil, error_message
local function read_token_file(token_file)
  local file, err = io.open(token_file, "r")
  if not file then
    return nil, "Failed to open token file: " .. err
  end
  
  local token = file:read("*all")
  file:close()
  
  if not token or token == "" then
    return nil, "Token file is empty"
  end
  
  return token
end

--- Parse XML response from STS AssumeRoleWithWebIdentity
-- Simple XML parser for STS response (avoiding external XML dependencies)
-- @param xml_body XML response body
-- @return table Credentials table or nil, error_message
local function parse_sts_response(xml_body)
  if not xml_body then
    return nil, "Empty response body"
  end
  
  -- Extract credentials from XML using pattern matching
  local access_key = xml_body:match("<AccessKeyId>([^<]+)</AccessKeyId>")
  local secret_key = xml_body:match("<SecretAccessKey>([^<]+)</SecretAccessKey>")
  local session_token = xml_body:match("<SessionToken>([^<]+)</SessionToken>")
  local expiration = xml_body:match("<Expiration>([^<]+)</Expiration>")
  
  if not access_key or not secret_key or not session_token then
    return nil, "Failed to parse credentials from STS response"
  end
  
  return {
    access_key_id = access_key,
    secret_access_key = secret_key,
    session_token = session_token,
    expiration = expiration,
  }
end

--- Check if cached credentials are still valid
-- @return boolean True if valid, false otherwise
local function is_cache_valid()
  local cached = credentials_cache:get(CACHE_KEY)
  
  if not cached or
     not cached.access_key_id or
     not cached.secret_access_key or
     not cached.session_token or
     not cached.expiration then
    return false
  end
  
  local expiration_time = utils.parse_iso8601(cached.expiration)
  if not expiration_time then
    return false
  end
  
  -- Refresh 5 minutes before expiration
  local now = os.time()
  local buffer_time = 300  -- 5 minutes
  
  return (expiration_time - now) > buffer_time
end

--- Call AWS STS AssumeRoleWithWebIdentity API
-- @param opts Options table containing token, role_arn, session_name, region, timeout
-- @return table Credentials or nil, error_message
local function assume_role_with_web_identity(opts)
  local token = opts.token
  local role_arn = opts.role_arn
  local session_name = opts.session_name or "kong-plugin-session"
  local region = opts.region or "us-east-1"
  
  local sts_endpoint = DEFAULT_STS_ENDPOINT
  if region and region ~= "us-east-1" then
    sts_endpoint = string.format("https://sts.%s.amazonaws.com/", region)
  end
  
  local params = {
    Action = "AssumeRoleWithWebIdentity",
    Version = "2011-06-15",
    RoleArn = role_arn,
    RoleSessionName = session_name,
    WebIdentityToken = token,
  }
  
  local query_parts = {}
  for key, value in pairs(params) do
    table.insert(query_parts, string.format("%s=%s", 
      key, 
      ngx.escape_uri(value)))
  end
  local query_string = table.concat(query_parts, "&")
  
  local httpc = http_client.new({
    timeout = opts.timeout or 10000,  -- Default 10 seconds for AWS STS
    max_retries = 3,
    retry_delay = 100,
  })
  
  local res, err = httpc:request_uri(sts_endpoint, {
    method = "POST",
    body = query_string,
    headers = {
      ["Content-Type"] = "application/x-www-form-urlencoded",
      ["Content-Length"] = tostring(#query_string),
    },
  })
  
  if err or not res then
    return nil, "Failed to call STS API: " .. (err or "unknown error")
  end
  
  if res.status ~= 200 then
    return nil, string.format("STS API returned error status %d: %s", 
      res.status, res.body or "")
  end
  
  local credentials, parse_err = parse_sts_response(res.body)
  if parse_err or not credentials then
    return nil, "Failed to parse STS response: " .. (parse_err or "unknown error")
  end
  
  return credentials
end

--- Create a new IRSA credentials fetcher
-- @param opts Options table
-- @param opts.token_file Path to Web Identity Token file (default: AWS_WEB_IDENTITY_TOKEN_FILE env var)
-- @param opts.role_arn IAM role ARN to assume (default: AWS_ROLE_ARN env var)
-- @param opts.session_name Session name (default: AWS_ROLE_SESSION_NAME env var or "kong-plugin-session")
-- @param opts.region AWS region (default: AWS_REGION or AWS_DEFAULT_REGION env var)
-- @param opts.timeout Timeout in milliseconds for STS API calls (default: 10000)
-- @return table IRSA instance
function IRSA.new(opts)
  opts = opts or {}
  
  local self = setmetatable({}, IRSA)
  
  -- Get configuration from options or pre-loaded environment variables
  -- Environment variables are read at module load time (master process)
  -- This avoids the need for nginx env directives
  self.token_file = opts.token_file or ENV_AWS_WEB_IDENTITY_TOKEN_FILE
  self.role_arn = opts.role_arn or ENV_AWS_ROLE_ARN
  self.session_name = opts.session_name or ENV_AWS_ROLE_SESSION_NAME or "kong-plugin-session"
  self.region = opts.region or ENV_AWS_REGION or ENV_AWS_DEFAULT_REGION or "us-east-1"
  self.timeout = opts.timeout or 10000
  
  -- Log configuration for debugging
  kong.log.notice("IRSA initialization - token_file: ", self.token_file or "nil", ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  kong.log.notice("IRSA initialization - role_arn: ", self.role_arn or "nil", ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  kong.log.notice("IRSA initialization - session_name: ", self.session_name, ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  kong.log.notice("IRSA initialization - region: ", self.region, ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  
  -- Validate required parameters
  if not self.token_file or self.token_file == "" then
    error("token_file must be specified or AWS_WEB_IDENTITY_TOKEN_FILE environment variable must be set")
  end
  
  if not self.role_arn or self.role_arn == "" then
    error("role_arn must be specified or AWS_ROLE_ARN environment variable must be set")
  end
  
  return self
end

--- Refresh credentials by calling STS AssumeRoleWithWebIdentity
-- @return boolean Success or nil, error_message
function IRSA:refresh()
  local token, err = read_token_file(self.token_file)
  if err or not token then
    return nil, err
  end
  
  local credentials, sts_err = assume_role_with_web_identity({
    token = token,
    role_arn = self.role_arn,
    session_name = self.session_name,
    region = self.region,
    timeout = self.timeout,
  })
  
  if sts_err or not credentials then
    return nil, sts_err
  end
  
  -- Cache credentials with 1 hour TTL
  credentials_cache:set(CACHE_KEY, {
    access_key_id = credentials.access_key_id,
    secret_access_key = credentials.secret_access_key,
    session_token = credentials.session_token,
    expiration = credentials.expiration,
  }, 3600)
  
  return true
end

--- Get AWS credentials (with automatic refresh if needed)
-- @return table Credentials containing access_key_id, secret_access_key, session_token or nil, error_message
function IRSA:get_credentials()
  -- First check: avoid lock if cache is valid (fast path)
  if is_cache_valid() then
    local cached = credentials_cache:get(CACHE_KEY)
    return {
      access_key_id = cached.access_key_id,
      secret_access_key = cached.secret_access_key,
      session_token = cached.session_token,
      expiration = cached.expiration,
    }
  end
  
  -- Cache invalid, acquire lock to refresh
  local lock_key = "irsa_refresh"
  local lock, lock_err = resty_lock:new("o11y_proxy_locks", {
    timeout = 30,    -- Wait up to 30 seconds to acquire lock
    exptime = 60,    -- Lock expires after 60 seconds (safety measure)
  })
  
  if not lock then
    return nil, "Failed to create lock: " .. tostring(lock_err)
  end
  
  local elapsed, err = lock:lock(lock_key)
  if not elapsed then
    return nil, "Failed to acquire lock: " .. tostring(err)
  end
  
  -- Second check: another request might have refreshed while we were waiting
  if not is_cache_valid() then
    kong.log.notice("IRSA credentials cache invalid or expired, refreshing..., trace_id: ",
        kong.ctx.shared.trace_id or "-")
    local success, refresh_err = self:refresh()
    lock:unlock()
    
    if not success then
      return nil, "Failed to refresh IRSA credentials: " .. refresh_err
    end
    kong.log.notice("IRSA credentials refreshed successfully, trace_id: ",
        kong.ctx.shared.trace_id or "-")
  else
    lock:unlock()
  end
  
  local cached = credentials_cache:get(CACHE_KEY)
  if not cached then
    -- This should not happen if cache is valid, but handle gracefully
    return nil, "Credentials cache is unexpectedly empty"
  end
  
  return {
    access_key_id = cached.access_key_id,
    secret_access_key = cached.secret_access_key,
    session_token = cached.session_token,
    expiration = cached.expiration,
  }
end

--- Clear the credentials cache (useful for testing)
function IRSA:clear_cache()
  credentials_cache:delete(CACHE_KEY)
end

-- Expose internal functions for testing
IRSA._read_token_file = read_token_file
IRSA._parse_sts_response = parse_sts_response
IRSA._is_cache_valid = is_cache_valid
IRSA._assume_role_with_web_identity = assume_role_with_web_identity

return IRSA