--- AWS AssumeRole implementation for Kong
-- Uses IRSA credentials to assume another role with optional external ID
-- Implements AWS STS AssumeRole API

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

local Assume = {}
Assume.__index = Assume

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

-- Cache multiple assumed role credentials indexed by role_arn
-- Use LRU cache with max 500 entries and 1 hour TTL
-- Each worker process has its own cache (no cross-worker synchronization needed)
local assume_credentials_cache, cache_err = lrucache.new(500)
if not assume_credentials_cache then
  error("Failed to create lrucache: " .. (cache_err or "unknown error"))
end

--- Parse XML response from STS AssumeRole
-- @param xml_body XML response body
-- @return table Credentials table or nil, error_message
local function parse_sts_assume_role_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 assume role credentials are still valid
-- @param role_arn The role ARN to check cache for
-- @return boolean True if valid, false otherwise
local function is_assume_cache_valid(role_arn)
  local cached = assume_credentials_cache:get(role_arn)
  
  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 AssumeRole API
-- @param opts Options table containing source_credentials, role_arn, external_id, session_name, region, timeout
-- @return table Credentials or nil, error_message
local function assume_role(opts)
  local source_credentials = opts.source_credentials
  local role_arn = opts.role_arn
  local external_id = opts.external_id
  local session_name = opts.session_name or "kong-plugin-session-assume"
  local region = opts.region or "us-east-1"
  
  if not source_credentials or not source_credentials.access_key_id then
    return nil, "Source credentials are required"
  end
  
  if not role_arn or role_arn == "" then
    return nil, "Role ARN is required"
  end
  
  local sts_endpoint = DEFAULT_STS_ENDPOINT
  local sts_host = "sts.amazonaws.com"
  if region and region ~= "us-east-1" then
    sts_endpoint = string.format("https://sts.%s.amazonaws.com/", region)
    sts_host = string.format("sts.%s.amazonaws.com", region)
  end
  
  -- Build request parameters
  local params = {
    Action = "AssumeRole",
    Version = "2011-06-15",
    RoleArn = role_arn,
    RoleSessionName = session_name,
  }
  
  if external_id and external_id ~= "" then
    params.ExternalId = external_id
  end
  
  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 body = table.concat(query_parts, "&")
  
  -- Sign the request
  local signed_headers, sign_err = aws_signer.sign_request({
    method = "POST",
    host = sts_host,
    path = "/",
    body = body,
    credentials = source_credentials,
    region = region,
    service = "sts",
    headers = {
      ["Content-Type"] = "application/x-www-form-urlencoded",
    }
  })
  
  if not signed_headers then
    return nil, "Failed to sign request: " .. sign_err
  end
  
  signed_headers["Content-Length"] = tostring(#body)
  
  -- Make the request with retry
  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 = body,
    headers = signed_headers,
  })
  
  if err or not res then
    return nil, "Failed to call STS AssumeRole API: " .. (err or "unknown error")
  end
  
  if res.status ~= 200 then
    kong.log.err("STS AssumeRole API error response: ", res.body, ", trace_id: ",
        kong.ctx.shared.trace_id or "-")
    return nil, string.format("STS AssumeRole API returned error status %d: %s", 
      res.status, res.body or "")
  end
  
  local credentials, parse_err = parse_sts_assume_role_response(res.body)
  if parse_err or not credentials then
    return nil, "Failed to parse STS AssumeRole response: " .. (parse_err or "unknown error")
  end
  
  return credentials
end

--- Create a new Assume role credentials fetcher
-- @param opts Options table
-- @param opts.irsa_opts Options to pass to IRSA.new() for getting source credentials
-- @param opts.role_arn IAM role ARN to assume (required)
-- @param opts.external_id External ID for AssumeRole (optional)
-- @param opts.session_name Session name (default: "kong-plugin-session-assume")
-- @param opts.region AWS region (default: "us-east-1")
-- @param opts.timeout Timeout in milliseconds for STS API calls (default: 10000)
-- @return table Assume instance
function Assume.new(opts)
  opts = opts or {}
  
  local self = setmetatable({}, Assume)
  
  self.role_arn = opts.role_arn
  self.external_id = opts.external_id
  self.session_name = opts.session_name or "kong-plugin-session-assume"
  self.region = opts.region or "us-east-1"
  self.timeout = opts.timeout or 10000
  
  -- Validate required parameters
  if not self.role_arn or self.role_arn == "" then
    error("role_arn must be specified")
  end
  
  -- Initialize IRSA for getting source credentials
  local irsa_opts = opts.irsa_opts or {}
  -- Pass timeout to IRSA if not already specified in irsa_opts
  if not irsa_opts.timeout then
    irsa_opts.timeout = self.timeout
  end
  self.irsa = irsa.new(irsa_opts)
  
  kong.log.notice("Assume initialization - role_arn: ", self.role_arn, ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  kong.log.notice("Assume initialization - external_id: ", self.external_id or "not set", ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  kong.log.notice("Assume initialization - session_name: ", self.session_name, ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  kong.log.notice("Assume initialization - region: ", self.region, ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  
  return self
end

--- Refresh assumed role credentials
-- @return boolean Success or nil, error_message
function Assume:refresh()
  -- First, get the source credentials from IRSA
  local source_credentials, irsa_err = self.irsa:get_credentials()
  if irsa_err or not source_credentials then
    return nil, "Failed to get IRSA credentials: " .. irsa_err
  end
  
  kong.log.notice("Got IRSA credentials, performing AssumeRole to: ", self.role_arn, ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  
  -- Use source credentials to assume the target role
  local credentials, assume_err = assume_role({
    source_credentials = source_credentials,
    role_arn = self.role_arn,
    external_id = self.external_id,
    session_name = self.session_name,
    region = self.region,
    timeout = self.timeout,
  })
  
  if assume_err or not credentials then
    return nil, assume_err
  end
  
  -- Cache the assumed role credentials indexed by role_arn
  -- TTL is set to 1 hour (3600 seconds)
  assume_credentials_cache:set(self.role_arn, {
    access_key_id = credentials.access_key_id,
    secret_access_key = credentials.secret_access_key,
    session_token = credentials.session_token,
    expiration = credentials.expiration,
  }, 3600)
  
  kong.log.notice("AssumeRole successful, credentials cached for role: ", self.role_arn, ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
  
  return true
end

--- Get assumed role AWS credentials (with automatic refresh if needed)
-- @return table Credentials containing access_key_id, secret_access_key, session_token or nil, error_message
function Assume:get_credentials()
  -- First check: avoid lock if cache is valid (fast path)
  if is_assume_cache_valid(self.role_arn) then
    local cached = assume_credentials_cache:get(self.role_arn)
    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 = "assume_role:" .. self.role_arn
  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: " .. lock_err
  end
  
  local elapsed, err = lock:lock(lock_key)
  if not elapsed then
    return nil, "Failed to acquire lock: " .. err
  end
  
  -- Second check: another request might have refreshed while we were waiting
  if not is_assume_cache_valid(self.role_arn) then
    kong.log.notice("Assume role 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 assume role credentials: " .. refresh_err
    end
    kong.log.notice("Assume role credentials refreshed successfully, trace_id: ",
        kong.ctx.shared.trace_id or "-")
  else
    lock:unlock()
  end
  
  -- Return cached credentials
  local cached = assume_credentials_cache:get(self.role_arn)
  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)
-- @param role_arn Optional role ARN to clear specific cache entry. If nil, clears all cache.
function Assume:clear_cache(role_arn)
  if role_arn then
    assume_credentials_cache:delete(role_arn)
  else
    assume_credentials_cache:flush_all()
  end
end

-- Expose internal functions for testing
Assume._parse_sts_assume_role_response = parse_sts_assume_role_response
Assume._is_assume_cache_valid = is_assume_cache_valid

return Assume