--- AWS Signature Version 4 implementation for Kong
-- Implements AWS SigV4 signing for HTTP requests
-- Reference: https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_sigv.html

local resty_sha256 = require("resty.sha256")
local str = require("resty.string")
local ffi = require("ffi")

local AwsSigner = {}

-- FFI definitions for OpenSSL HMAC
ffi.cdef[[
typedef struct hmac_ctx_st HMAC_CTX;
typedef struct engine_st ENGINE;
typedef struct evp_md_st EVP_MD;

const EVP_MD *EVP_sha256(void);
unsigned char *HMAC(const EVP_MD *evp_md, const void *key, int key_len,
                    const unsigned char *d, size_t n, unsigned char *md,
                    unsigned int *md_len);
]]

local C = ffi.C
local EVP_sha256 = C.EVP_sha256

--- HMAC-SHA256 implementation using OpenSSL via FFI
-- @param key The secret key
-- @param data The data to sign
-- @return string Binary HMAC-SHA256 result
local function hmac_sha256(key, data)
  local md = ffi.new("unsigned char[32]")
  local md_len = ffi.new("unsigned int[1]")
  
  local result = C.HMAC(
    EVP_sha256(),
    key, #key,
    ffi.cast("const unsigned char*", data), #data,
    md, md_len
  )
  
  if result == nil then
    error("HMAC-SHA256 computation failed")
  end
  
  return ffi.string(md, md_len[0])
end

--- Canonicalize a path according to AWS SigV4 rules
-- @param path The URI path
-- @return string Canonical path
local function canonicalize_path(path)
  if not path or path == "" then
    return "/"
  end
  
  local segments = {}
  for segment in path:gmatch("/([^/]*)") do
    if segment == "" or segment == "." then
      -- Skip empty and current directory segments
    elseif segment == ".." then
      -- Go up one directory (intentionally discards components at top level)
      segments[#segments] = nil
    else
      -- URL decode and re-encode the segment
      local unescaped = ngx.unescape_uri(segment)
      unescaped = unescaped:gsub("[^%w%-%._~]", function(c)
        return string.format("%%%02X", string.byte(c))
      end)
      segments[#segments + 1] = unescaped
    end
  end
  
  if #segments == 0 then
    return "/"
  end
  
  -- Preserve trailing slash if present
  local result = "/" .. table.concat(segments, "/")
  if path:sub(-1) == "/" and result ~= "/" then
    result = result .. "/"
  end
  
  return result
end

--- Build canonical headers string and signed headers list
-- @param host Host name (may include port)
-- @param method HTTP method
-- @param headers Optional request headers table
-- @param amz_date ISO8601 formatted timestamp
-- @param session_token Optional AWS session token
-- @return string, string Canonical headers string and signed headers list (semicolon-separated)
local function build_canonical_headers(host, method, headers, amz_date, session_token)
  -- Remove standard ports from host (80 for HTTP, 443 for HTTPS)
  -- AWS SigV4 expects host without standard ports
  local canonical_host = host
  if host:match(":80$") or host:match(":443$") then
    canonical_host = host:match("^(.+):%d+$")
  end
  
  -- Prepare canonical headers
  local canonical_headers = {
    ["host"] = canonical_host,
    ["x-amz-date"] = amz_date,
  }
  
  -- Only include content-type for methods that typically have a body
  -- For GET, DELETE, HEAD requests, omit content-type unless explicitly provided
  local has_body_method = method ~= "GET" and method ~= "DELETE" and method ~= "HEAD"
  if has_body_method then
    canonical_headers["content-type"] = headers and headers["Content-Type"] or "application/x-www-form-urlencoded"
  elseif headers and headers["Content-Type"] then
    -- If explicitly provided in headers even for GET/DELETE, include it
    canonical_headers["content-type"] = headers["Content-Type"]
  end
  
  if session_token then
    canonical_headers["x-amz-security-token"] = session_token
  end
  
  -- Build signed headers list (must be sorted)
  local signed_headers_list = {"host", "x-amz-date"}
  if canonical_headers["content-type"] then
    table.insert(signed_headers_list, "content-type")
  end
  if session_token then
    table.insert(signed_headers_list, "x-amz-security-token")
  end
  table.sort(signed_headers_list)
  
  -- Build canonical headers string
  local canonical_headers_str = ""
  for _, h in ipairs(signed_headers_list) do
    -- Trim whitespace from header values (AWS requirement)
    local header_value = tostring(canonical_headers[h])
    header_value = header_value:match("^%s*(.-)%s*$") -- trim leading/trailing whitespace
    canonical_headers_str = canonical_headers_str .. h .. ":" .. header_value .. "\n"
  end
  
  local signed_headers = table.concat(signed_headers_list, ";")
  
  return canonical_headers_str, signed_headers
end

--- Canonicalize query string according to AWS SigV4 rules
-- @param query_string The query string
-- @return string Canonical query string
local function canonicalize_query_string(query_string)
  if not query_string or query_string == "" then
    return ""
  end
  
  local params = {}
  for key, val in query_string:gmatch("([^&=]+)=?([^&]*)") do
    -- URL decode
    key = ngx.unescape_uri(key)
    val = ngx.unescape_uri(val)
    -- URL encode again with proper encoding
    key = key:gsub("[^%w%-%._~]", function(c)
      return string.format("%%%02X", string.byte(c))
    end)
    val = val:gsub("[^%w%-%._~]", function(c)
      return string.format("%%%02X", string.byte(c))
    end)
    params[#params + 1] = key .. "=" .. val
  end
  
  table.sort(params)
  return table.concat(params, "&")
end

--- Generate AWS Signature Version 4 for HTTP requests
-- @param opts Options table
-- @param opts.method HTTP method (e.g., "POST", "GET")
-- @param opts.host Host name (e.g., "sts.amazonaws.com")
-- @param opts.path Request path (default: "/")
-- @param opts.headers Optional request headers table
-- @param opts.body Request body (default: "")
-- @param opts.credentials Credentials table containing access_key_id, secret_access_key, and optional session_token
-- @param opts.region AWS region (default: "us-east-1")
-- @param opts.service AWS service name (default: "sts")
-- @param opts.query_string Optional query string for GET requests (default: "")
-- @return table Signed headers ready to use in HTTP request or nil, error_message
function AwsSigner.sign_request(opts)
  if not opts then
    return nil, "Options table is required"
  end
  
  local method = opts.method
  local host = opts.host
  local path = opts.path or "/"
  local headers = opts.headers or {}
  local body = opts.body or ""
  local credentials = opts.credentials
  local region = opts.region or "us-east-1"
  local service = opts.service or "sts"
  local query_string = opts.query_string or ""
  
  -- Validate required parameters
  if not method or method == "" then
    return nil, "HTTP method is required"
  end
  
  if not host or host == "" then
    return nil, "Host is required"
  end
  
  if not credentials or not credentials.access_key_id or not credentials.secret_access_key then
    return nil, "Valid credentials with access_key_id and secret_access_key are required"
  end
  
  -- Get current timestamp
  local current_time = ngx.time()
  local date_stamp = os.date("!%Y%m%d", current_time)
  local amz_date = os.date("!%Y%m%dT%H%M%SZ", current_time)
  
  -- Canonicalize path and query string
  path = canonicalize_path(path)
  query_string = canonicalize_query_string(query_string)
  
  -- Build canonical headers and signed headers list
  local canonical_headers_str, signed_headers = build_canonical_headers(host, method, headers, amz_date, credentials.session_token)
  
  -- Calculate payload hash (SHA256 of body)
  local sha256 = resty_sha256:new()
  sha256:update(body)
  local payload_hash = sha256:final()
  local payload_hash_hex = str.to_hex(payload_hash)
  
  -- Build canonical request
  local canonical_request = table.concat({
    method,
    path,
    query_string,
    canonical_headers_str,
    signed_headers,
    payload_hash_hex
  }, "\n")
  
  -- Create credential scope
  local credential_scope = date_stamp .. "/" .. region .. "/" .. service .. "/aws4_request"
  
  -- Calculate canonical request hash
  local sha256_2 = resty_sha256:new()
  sha256_2:update(canonical_request)
  local canonical_request_hash = sha256_2:final()
  local canonical_request_hash_hex = str.to_hex(canonical_request_hash)
  
  -- Create string to sign
  local string_to_sign = table.concat({
    "AWS4-HMAC-SHA256",
    amz_date,
    credential_scope,
    canonical_request_hash_hex
  }, "\n")
  
  -- Calculate signature using HMAC-SHA256
  local k_date = hmac_sha256("AWS4" .. credentials.secret_access_key, date_stamp)
  local k_region = hmac_sha256(k_date, region)
  local k_service = hmac_sha256(k_region, service)
  local k_signing = hmac_sha256(k_service, "aws4_request")
  local signature_bin = hmac_sha256(k_signing, string_to_sign)
  
  -- Convert signature to hex
  local signature = str.to_hex(signature_bin)
  
  -- Build authorization header
  local authorization = string.format(
    "AWS4-HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s",
    credentials.access_key_id,
    credential_scope,
    signed_headers,
    signature
  )
  
  -- Remove standard ports from host (80 for HTTP, 443 for HTTPS)
  local canonical_host = host
  if host:match(":80$") or host:match(":443$") then
    canonical_host = host:match("^(.+):%d+$")
  end
  
  -- Return all headers needed for the request
  local result_headers = {
    ["Authorization"] = authorization,
    ["X-Amz-Date"] = amz_date,
    ["Host"] = canonical_host,
  }
  
  -- Only include content-type if it was part of the signature
  local has_body_method = method ~= "GET" and method ~= "DELETE" and method ~= "HEAD"
  if has_body_method or (headers and headers["Content-Type"]) then
    result_headers["Content-Type"] = headers and headers["Content-Type"] or "application/x-www-form-urlencoded"
  end
  
  if credentials.session_token then
    result_headers["X-Amz-Security-Token"] = credentials.session_token
  end
  
  return result_headers
end

-- Expose internal functions for testing
AwsSigner._canonicalize_path = canonicalize_path
AwsSigner._canonicalize_query_string = canonicalize_query_string
AwsSigner._build_canonical_headers = build_canonical_headers
AwsSigner._hmac_sha256 = hmac_sha256

return AwsSigner