local kong_meta = require "kong.meta"

-- Cookie Filter Plugin
-- Filters cookies before forwarding to upstream, only keeping whitelisted fields
-- This prevents non-ASCII or unwanted cookies from causing upstream errors
local CookieFilter = { PRIORITY = 1280, VERSION = kong_meta.version }

-- Cookie sid field names by environment (same as cookie-auth plugin)
local COOKIE_SID_FIELDS = {
    prod = "account_sid",
    dev = "account_sid_dev",
    staging = "account_sid_staging"
}

-- Extract value from cookie header by cookie name
local function extract_value_from_cookie(cookie_header, cookie_name)
    if not cookie_header then
        return nil
    end
    -- Match <cookie_name>=<value> format
    local pattern = cookie_name .. "=([^;]+)"
    local val = string.match(cookie_header, pattern)
    return val
end

-- Parse cookie header into key-value pairs
local function parse_cookies(cookie_header)
    if not cookie_header then
        return {}
    end

    local cookies = {}
    for pair in string.gmatch(cookie_header, "([^;]+)") do
        pair = pair:match("^%s*(.-)%s*$") -- trim whitespace
        local eq_pos = string.find(pair, "=")
        if eq_pos then
            local key = string.sub(pair, 1, eq_pos - 1)
            key = key:match("^%s*(.-)%s*$") -- trim key
            local value = string.sub(pair, eq_pos + 1)
            cookies[key] = value
        end
    end
    return cookies
end

-- Check if a cookie key matches any pattern in the whitelist
local function is_whitelisted(key, whitelist)
    for _, pattern in ipairs(whitelist) do
        if ngx.re.match(key, "^" .. pattern .. "$", "jo") then
            return true
        end
    end
    return false
end

-- Build cookie header from filtered cookies
local function build_cookie_header(cookies)
    local parts = {}
    for key, value in pairs(cookies) do
        table.insert(parts, key .. "=" .. value)
    end
    if #parts == 0 then
        return nil
    end
    return table.concat(parts, "; ")
end

-- Check if origin matches any pattern to be filtered
local function is_filter_origin(origin, filt_origins)
    if not origin or not filt_origins or #filt_origins == 0 then
        return false
    end
    for _, pattern in ipairs(filt_origins) do
        if ngx.re.match(origin, "^" .. pattern .. "$", "jo") then
            return true
        end
    end
    return false
end

local function is_filter(conf)
    local cookie_header = kong.request.get_header("cookie")
    if not cookie_header then
        return false
    end

    -- If whitelist is empty, don't filter anything
    if not conf.whitelist or #conf.whitelist == 0 then
        kong.log.debug("cookie-filter: whitelist is empty, trace_id: ",
            kong.ctx.shared.trace_id or "-")
        return false
    end

    -- Determine the cookie name based on environment
    local env = conf.env_type
    local sid_field = COOKIE_SID_FIELDS[env]
    if not sid_field then
        kong.log.notice("Unknown env_type: ", env, ", defaulting to prod cookie name, trace_id: ",
            kong.ctx.shared.trace_id or "-")
        sid_field = COOKIE_SID_FIELDS["prod"]
    end

    -- If sid cookie is present, filter cookies
    local sid = extract_value_from_cookie(cookie_header, sid_field)
    if sid then
        kong.log.debug("cookie-filter: sid found: ", sid_field, ", trace_id: ",
            kong.ctx.shared.trace_id or "-")
        return true
    end

    -- If request origin matches allow_origins, filter cookies
    local origin = kong.request.get_header("origin")
    if is_filter_origin(origin, conf.filt_origins) then
        kong.log.debug("cookie-filter: filter matched origin: ", origin, ", trace_id: ",
            kong.ctx.shared.trace_id or "-")
        return true
    end

    kong.log.debug("cookie-filter: skipped for no sid field: ", sid_field, " and not match filter origin: , trace_id: ",
        kong.ctx.shared.trace_id or "-")

    return false
end

function CookieFilter:access(conf)
    if not is_filter(conf) then
        return
    end

    local cookie_header = kong.request.get_header("cookie")
    if not cookie_header then
        return
    end

    kong.log.notice("cookie-filter: origin cookie: ", cookie_header, ", trace_id: ",
        kong.ctx.shared.trace_id or "-")

    local cookies = parse_cookies(cookie_header)
    local filtered_cookies = {}
    local filtered_count = 0

    for key, value in pairs(cookies) do
        if is_whitelisted(key, conf.whitelist) then
            filtered_cookies[key] = value
        else
            filtered_count = filtered_count + 1
            kong.log.debug("cookie-filter: filtered out cookie field: ", key, ", trace_id: ",
                kong.ctx.shared.trace_id or "-")
        end
    end

    if filtered_count > 0 then
        local new_cookie_header = build_cookie_header(filtered_cookies)

        kong.log.notice("cookie-filter: new cookie: ", new_cookie_header, ", trace_id: ",
            kong.ctx.shared.trace_id or "-")

        if new_cookie_header then
            kong.service.request.set_header("cookie", new_cookie_header)
        else
            -- No whitelisted cookies, remove the header entirely
            kong.service.request.clear_header("cookie")
            kong.log.notice("cookie-filter: removed all cookies (none in whitelist), trace_id: ",
                kong.ctx.shared.trace_id or "-")
        end
    end
end

return CookieFilter