local Assume = require "kong.plugins.o11y-proxy.assume"
local AwsSigner = require "kong.plugins.o11y-proxy.aws_signer"
local Router = require "kong.plugins.o11y-proxy.router"
local Authn = require "kong.plugins.o11y-proxy.authn"
local Authz = require "kong.plugins.o11y-proxy.authz"
local ResourceMapper = require "kong.plugins.o11y-proxy.resource_mapper"
local cjson = require "cjson.safe"
local lrucache = require "resty.lrucache"
local kong_meta = require "kong.meta"

local O11YProxyHandler = {
  VERSION = kong_meta.version,
  PRIORITY = 1050,
}

-- Cache of Assume role instances indexed by role_arn
-- Use LRU cache to handle potentially hundreds of different role ARNs
-- Max 500 entries, no TTL (instances should persist for worker lifetime)
local assume_instances, cache_err = lrucache.new(500)
if not assume_instances or cache_err then
  error("Failed to create assume_instances cache: " .. (cache_err or "unknown error"))
end

-- Get the request path after removing Kong's route prefix (if strip_path is enabled)
-- Returns: the path component after the stripped prefix
local function get_request_path_without_prefix()
  local request_path = kong.request.get_path()
  local stripped_prefix = kong.request.get_forwarded_prefix()

  if stripped_prefix and stripped_prefix ~= "" then
    -- Remove the stripped prefix from the full path to get the remaining path
    if request_path:sub(1, #stripped_prefix) == stripped_prefix then
      local path_after_prefix = request_path:sub(#stripped_prefix + 1)
      if path_after_prefix == "" then
        return "/"
      end
      -- Ensure the path starts with /
      if path_after_prefix:sub(1, 1) ~= "/" then
        return "/" .. path_after_prefix
      end
      return path_after_prefix
    else
      -- Fallback: prefix doesn't match, use full path
      return request_path
    end
  else
    -- No prefix was stripped, use the full path
    return request_path
  end
end

-- Parse URL into components (scheme, host, port, path, query)
-- Returns: scheme, host, port, path, query or nil, error_message on failure
local function parse_url(url)
  if not url or url == "" then
    return nil, "URL cannot be empty"
  end

  -- Match URL pattern: scheme://host:port/path?query
  -- Pattern: ([^:]+)://([^/]+)(/[^?]+)?%?(.*)
  -- Group 1: scheme (http or https)
  -- Group 2: host:port (everything before first /)
  -- Group 3: path (optional, starts with /)
  -- Group 4: query (optional, after ?)
  local scheme, host_port, path, query = url:match("([^:]+)://([^/]+)(/[^?]+)%?(.*)")

  if not scheme then
    scheme, host_port, path = url:match("([^:]+)://([^/]+)(/.*)")
  end
  if not scheme then
    scheme, host_port = url:match("([^:]+)://([^/]+)")
  end
  if not scheme or not host_port then
    return nil, "Invalid URL format. Expected: scheme://host[:port][/path][?query]"
  end

  local host, port
  local colon_pos = host_port:find(":")
  if colon_pos then
    host = host_port:sub(1, colon_pos - 1)
    local port_str = host_port:sub(colon_pos + 1)
    port = tonumber(port_str)
    if not port or port < 1 or port > 65535 then
      return nil, "Port must be a valid port number (1-65535)"
    end
  else
    host = host_port
    port = (scheme == "https") and 443 or 80
  end

  if not host or host == "" then
    return nil, "Invalid URL format (no host found)"
  end

  path = path or "/"
  if query == "" then
    query = nil
  end

  return scheme, host, port, path, query
end

-- Validate authentication and authorization with external services
-- Returns: user_info on success, exits with 401/403/500 on auth/authz failure
local function validate_auth_and_authz(conf)
  if conf and conf.require_auth == false then
    kong.log.notice("Authentication disabled by plugin config, skipping authn/authz, trace_id: ",
        kong.ctx.shared.trace_id or "-")
    return nil
  end

  -- Step 1: Authn
  local user_info = nil
  local account_user_meta = kong.request.get_header("account_user_meta")

  if account_user_meta and account_user_meta ~= "" then
    -- Parse user info from account_user_meta header
    -- Expected format: {"id": 12345678,"email": "xxx","first_name": "xxx","last_name": "xxx"}
    local parsed, err = cjson.decode(account_user_meta)
    if type(parsed) == "table" and parsed.id then
      user_info = { id = parsed.id }
      kong.log.notice("User info parsed from account_user_meta header, user_id: ", user_info.id,
          ", trace_id: ", kong.ctx.shared.trace_id or "-")
    else
      kong.log.notice(
        "fallback to authn using bearer token at Authorization header due to failed to parse account_user_meta header: ",
        err, "account_user_meta: ", account_user_meta, ", trace_id: ",
        kong.ctx.shared.trace_id or "-")
    end
  end

  if not user_info then
    -- Authn fallback using bearer token at Authorization header
    if conf and conf.authn_service_url and conf.authn_service_url ~= "" then
      local token_header = "Authorization"
      local auth0_token = kong.request.get_header(token_header)

      if not auth0_token or auth0_token == "" then
        kong.log.err("Missing Auth0 token in ", token_header, " header, trace_id: ",
            kong.ctx.shared.trace_id or "-")
        return kong.response.exit(401, { message = "Authentication token is required" })
      end

      -- Remove "Bearer " prefix if present (case-insensitive)
      auth0_token = auth0_token:gsub("^[Bb][Ee][Aa][Rr][Ee][Rr]%s+", "")

      local authn_result, authn_err = Authn.authenticate(conf.authn_service_url, auth0_token, conf.authn_timeout, conf.authn_cache_ttl)

      if authn_err or not authn_result then
        kong.log.err("Failed to authenticate: ", authn_err, ", trace_id: ",
            kong.ctx.shared.trace_id or "-")
        return kong.response.exit(500, { message = "Failed to validate authentication: " .. (authn_err or "unknown error") })
      end

      if not authn_result.check_result then
        kong.log.warn("Authentication failed for token, trace_id: ",
            kong.ctx.shared.trace_id or "-")
        return kong.response.exit(401, { message = "Authentication failed" })
      end

      user_info = { id = authn_result.user_id }
      kong.log.notice("Authentication successful for user_id: ", user_info.id, ", trace_id: ",
          kong.ctx.shared.trace_id or "-")
    else
      kong.log.err("Authentication failed: requires user_info from authentication, trace_id: ",
          kong.ctx.shared.trace_id or "-")
      return kong.response.exit(401, { message = "Authentication failed" })
    end
  end

  -- Step 2: Authorization - check user permissions
  if conf.authz_service_url and conf.authz_service_url ~= "" then

    local request_path = get_request_path_without_prefix()
    local org_id = kong.request.get_header("X-Org-Id")
    local biz_type = kong.request.get_header("X-Biz-Type")
    local current_project_id = kong.request.get_header("X-Current-Project-Id")
    local project_id = current_project_id

    if not project_id or project_id == "" then
      project_id = kong.request.get_header("X-Project-Id")
    end
    local cluster_id = kong.request.get_header("X-Application-Id")
    if not biz_type or biz_type == "" then
      return kong.response.exit(400, { message = "Invalid business type" })
    end
    local resource_uri, req_uri = ResourceMapper.get_authz_uris(request_path, biz_type, org_id, project_id, cluster_id)

    local authz_result, authz_err = Authz.check_permission(
      conf.authz_service_url,
      user_info.id,
      resource_uri,
      req_uri,
      conf.authz_timeout,
      conf.authz_cache_ttl
    )

    if authz_err or not authz_result then
      kong.log.err("Failed to check authorization: ", authz_err,
        ", resource_uri: ", resource_uri, ", req_uri: ", req_uri,
        ", user_id: ", user_info.id, ", trace_id: ", kong.ctx.shared.trace_id or "-")
      return kong.response.exit(403, { message = "Failed to validate authorization: " .. (authz_err or "unknown error") })
    end

    if not authz_result.has_permission then
      kong.log.warn("Authorization denied for user_id ", user_info.id, ", resource: ", resource_uri,
          ", trace_id: ", kong.ctx.shared.trace_id or "-")
      return kong.response.exit(403, { message = "Access denied: insufficient permissions" })
    end

    kong.log.notice("Authorization successful for user_id: ", user_info.id, ", resource: ", resource_uri,
        ", trace_id: ", kong.ctx.shared.trace_id or "-")
  else
    kong.log.notice("No authz_service_url configured, skipping authorization, trace_id: ",
        kong.ctx.shared.trace_id or "-")
  end

  return user_info
end

-- Get routing information from router service
-- Returns: route_data on success, nil and error_message on failure, exits with 500 on router failure
local function get_routing_info(conf)
  local region_id = kong.request.get_header("X-Region-Id")
  local org_id = kong.request.get_header("X-Org-Id")
  local biz_type = kong.request.get_header("X-Biz-Type")
  local project_id = kong.request.get_header("X-Project-Id")
  local application_id = kong.request.get_header("X-Application-Id")

  local route_data, router_err = Router.query_target(conf.router_service_url, region_id, org_id, biz_type, project_id, application_id, conf.router_timeout, conf.router_cache_ttl)
  if router_err or not route_data then
    kong.log.err("Failed to get routing data from router service: ", router_err, ", trace_id: ",
        kong.ctx.shared.trace_id or "-")
    return kong.response.exit(500, { message = "Failed to determine routing target: " .. (router_err or "unknown error") })
  end

  return route_data
end

-- Get AWS credentials via AssumeRole for the given access_principal
-- Returns: credentials on success, nil if no access_principal, exits with 500 on credential failure
local function get_aws_credentials(conf, route_data)
  if not route_data.access_principal or route_data.access_principal == "" then
    -- No role to assume, return nil
    kong.log.notice("No access_principal provided, skipping AWS AssumeRole, trace_id: ",
        kong.ctx.shared.trace_id or "-")
    return nil
  end

  local role_arn = route_data.access_principal

  -- Get or create Assume instance for this role
  -- Using LRU cache to handle potentially hundreds of different role ARNs
  -- Note: Concurrent creation of multiple instances for the same role_arn is acceptable because:
  -- 1. Assume instances are stateless (credentials stored in global lrucache)
  -- 2. Only the last created instance will be kept in cache
  -- 3. The actual credential refresh is already protected by locks in assume.lua
  -- 4. At most one redundant instance will be created during initialization, which is negligible
  local instance = assume_instances:get(role_arn)
  if not instance then
    local success, result = pcall(function()
      return Assume.new({
        region = conf.region,
        role_arn = role_arn,
        external_id = route_data.access_external_id,
        timeout = conf.aws_sts_timeout,
      })
    end)

    if success then
      -- Set with nil TTL (persist for worker lifetime)
      assume_instances:set(role_arn, result, nil)
      instance = result
      kong.log.notice("Assume role instance initialized for role: ", role_arn, ", trace_id: ",
          kong.ctx.shared.trace_id or "-")
    else
      kong.log.err("Failed to initialize Assume role: ", result, ", trace_id: ",
          kong.ctx.shared.trace_id or "-")
      return kong.response.exit(500, { message = "Failed to initialize AWS role assumption" })
    end
  end

  -- Get credentials from the Assume instance
  -- The get_credentials() call is protected by locks in assume.lua to prevent concurrent refresh
  local creds, err = instance:get_credentials()
  if not creds then
    kong.log.err("Failed to get AssumeRole credentials: ", err, ", trace_id: ",
        kong.ctx.shared.trace_id or "-")
    return kong.response.exit(500, { message = "Failed to obtain AWS credentials" })
  end

  return creds
end

-- Sign request with AWS SigV4 and set signed headers
-- Returns: true on success, exits with 500 on signing failure
local function resolve_sigv4_region(conf, host)
  local region = nil

  if host and host ~= "" then
    -- Expected format: {api-id}.execute-api.{region}.amazonaws.com
    region = host:match("%.execute%-api%.([^.]+)%.amazonaws%.com$")
  end

  if (not region or region == "") and conf and conf.region and conf.region ~= "" then
    region = conf.region
  end

  if not region or region == "" then
    region = "us-west-2"
  end

  return region
end

-- Sign request with AWS SigV4 and set signed headers
-- Returns: true on success, exits with 500 on signing failure
local function sign_and_set_headers(conf, credentials, host, port, scheme, final_path, original_query)
  local method = kong.request.get_method()
  local body = kong.request.get_raw_body() or ""
  local original_headers = kong.request.get_headers()

  -- Prioritize configured region from plugin config (schema provides default "us-west-2")
  local region = resolve_sigv4_region(conf, host)

  local sign_opts = {
    method = method,
    host = host .. (((scheme == "https" and port ~= 443) or (scheme == "http" and port ~= 80)) and (":" .. port) or ""),
    path = final_path,
    headers = original_headers,
    body = body,
    credentials = {
      access_key_id = credentials.access_key_id,
      secret_access_key = credentials.secret_access_key,
      session_token = credentials.session_token,
    },
    region = region,
    service = "execute-api",
    query_string = original_query or "",
  }

  local signed_headers, err = AwsSigner.sign_request(sign_opts)
  if not signed_headers then
    kong.log.err("Failed to sign request: ", err, ", trace_id: ",
        kong.ctx.shared.trace_id or "-")
    return kong.response.exit(500, { message = "Failed to sign AWS request" })
  end

  -- Set the signed headers for the upstream request
  kong.service.request.set_header("Authorization", signed_headers["Authorization"])
  kong.service.request.set_header("X-Amz-Date", signed_headers["X-Amz-Date"])
  kong.service.request.set_header("Host", signed_headers["Host"])

  if signed_headers["Content-Type"] then
    kong.service.request.set_header("Content-Type", signed_headers["Content-Type"])
  end

  if signed_headers["X-Amz-Security-Token"] then
    kong.service.request.set_header("X-Amz-Security-Token", signed_headers["X-Amz-Security-Token"])
  end

  return true
end

-- Forward request to upstream target
-- Returns: nothing, sets Kong service target and logs forwarding info
local function forward_request(scheme, host, port, final_path, original_query)
  kong.service.set_target(host, port)
  kong.service.request.set_scheme(scheme)
  kong.service.request.set_path(final_path)
  if original_query and original_query ~= "" then
    kong.service.request.set_raw_query(original_query)
  end

  kong.log.notice("Forwarding request to ", scheme, "://", host, ":", port, final_path,
      (original_query and ("?" .. original_query) or ""), ", trace_id: ",
      kong.ctx.shared.trace_id or "-")
end

function O11YProxyHandler:access(conf)
  -- Bypass OPTIONS requests (CORS preflight)
  if kong.request.get_method() == "OPTIONS" then
    return kong.response.exit(204)
  end
  -- Step 1: Validate authentication and authorization if configured
  validate_auth_and_authz(conf)

  -- Step 2: Get routing information from router service
  local route_data = get_routing_info(conf)

  -- Step 3: Parse target URL
  local target_url = route_data.dns_address
  local scheme, host, port, base_path, query = parse_url(target_url)
  if not scheme then
    kong.log.err("Failed to parse target URL: ", target_url, ", trace_id: ",
        kong.ctx.shared.trace_id or "-")
    return kong.response.exit(500, { message = "Failed to find dataplane routing info" })
  end

  -- Step 4: Calculate final path and query
  local path_after_prefix = get_request_path_without_prefix()
  -- Avoid double slashes when base_path is "/" and path_after_prefix starts with "/"
  local final_path
  if base_path == "/" and path_after_prefix:sub(1, 1) == "/" then
    final_path = path_after_prefix
  else
    final_path = base_path .. path_after_prefix
  end
  local original_query = kong.request.get_raw_query()

  -- Step 5: If access_principal exists, get credentials and sign request; otherwise forward directly
  if route_data.access_principal and route_data.access_principal ~= "" then
    local credentials = get_aws_credentials(conf, route_data)
    if credentials then
      sign_and_set_headers(conf, credentials, host, port, scheme, final_path, original_query)
    end
  end

  -- Step 6: Forward request to upstream
  forward_request(scheme, host, port, final_path, original_query)
end

-- Expose internal functions for testing
O11YProxyHandler._parse_url = parse_url

return O11YProxyHandler