local http = require "resty.http"

local _M = {}

local DEFAULT_MAX_REDIRECTS = 2
local LOOPBACK_FOLLOW_HOST = "127.0.0.1"
local FOLLOW_CACHE_CONFIG_HEADER = "X-AGC-Follow-Cache-Config"
local FOLLOW_CACHE_KEY_HEADER = "X-AGC-Follow-Cache-Key"

local function normalize_boolean(value, default)
    if value == nil then
        return default
    end

    if value == true or value == "true" then
        return true
    end

    if value == false or value == "false" then
        return false
    end

    return default
end

local function normalize_max_redirects(value)
    local max_redirects = tonumber(value) or DEFAULT_MAX_REDIRECTS
    if max_redirects < 1 then
        return 1
    end
    if max_redirects > 5 then
        return 5
    end

    return max_redirects
end

local function parse_location(location, orig_scheme, orig_host)
    if not location or location == "" then
        return nil, nil, nil, nil
    end

    ngx.log(ngx.INFO, "[RedirectFollow] Parsing location: " .. location)

    local scheme, host, path = location:match("^(https?)://([^/]+)(.*)$")
    if scheme and host then
        if path == "" then
            path = "/"
        end
        return scheme .. "://" .. host .. path, scheme, host, path
    end

    local host2, path2 = location:match("^//([^/]+)(.*)$")
    if host2 then
        if path2 == "" then
            path2 = "/"
        end
        return orig_scheme .. "://" .. host2 .. path2, orig_scheme, host2, path2
    end

    if location:sub(1, 1) == "/" then
        return orig_scheme .. "://" .. orig_host .. location, orig_scheme, orig_host, location
    end

    ngx.log(ngx.ERR, "[RedirectFollow] Unable to parse location: " .. location)
    return nil, nil, nil, nil
end

local function determine_host_header(location_host, redirect_follow_host, original_request_host)
    if not redirect_follow_host or redirect_follow_host == "" then
        return location_host
    end

    if redirect_follow_host == "$host" then
        return original_request_host
    end

    return redirect_follow_host
end

local function strip_port(host)
    return tostring(host or ""):gsub(":%d+$", "")
end

local function should_use_loopback_follow(host_header, original_request_host, cache_key)
    return cache_key and cache_key ~= ""
        and string.lower(strip_port(host_header)) == string.lower(strip_port(original_request_host))
end

local function build_loopback_host(location_host)
    local port = tostring(location_host or ""):match(":(%d+)$")
    if port then
        return LOOPBACK_FOLLOW_HOST .. ":" .. port
    end

    return LOOPBACK_FOLLOW_HOST
end

local function replace_url_host(url, replacement_host)
    return (url:gsub("^(https?://)[^/]+", "%1" .. replacement_host, 1))
end

local function append_args(url, args)
    if not args or args == "" then
        return url
    end

    if url:find("?", 1, true) then
        return url .. "&" .. args
    end

    return url .. "?" .. args
end

local function build_headers(host_header, keep_headers, original_headers, follow_cache_config, follow_cache_key)
    local headers = {}

    if keep_headers and original_headers then
        for k, v in pairs(original_headers) do
            local lower_k = k:lower()
            if lower_k ~= "host"
                and lower_k ~= "connection"
                and lower_k ~= "content-length"
                and lower_k ~= "transfer-encoding"
                and lower_k ~= string.lower(FOLLOW_CACHE_CONFIG_HEADER)
                and lower_k ~= string.lower(FOLLOW_CACHE_KEY_HEADER) then
                headers[k] = v
            end
        end
    end

    headers["Host"] = host_header
    if follow_cache_config and follow_cache_config ~= "" then
        headers[FOLLOW_CACHE_CONFIG_HEADER] = follow_cache_config
    end
    if follow_cache_key and follow_cache_key ~= "" then
        headers[FOLLOW_CACHE_KEY_HEADER] = follow_cache_key
    end
    return headers
end

local function is_redirect(status)
    return status == 301 or status == 302 or status == 307 or status == 308
end

function _M.follow(options)
    local location = options.location
    if not location or location == "" then
        return nil, "no redirect location"
    end

    local max_redirects = normalize_max_redirects(options.max_redirects)
    local keep_args = normalize_boolean(options.keep_args, true)
    local keep_headers = normalize_boolean(options.keep_headers, true)
    local redirect_follow_host = options.redirect_follow_host or ""
    local original_scheme = options.orig_scheme or "http"
    local original_host = options.orig_host or ngx.var.host
    local original_args = options.orig_args
    local original_headers = options.orig_headers or {}
    local method = options.method or ngx.req.get_method()
    local follow_cache_config = options.cache_config
    local follow_cache_key = options.cache_key

    ngx.log(ngx.INFO, "[RedirectFollow] Start follow route=" .. tostring(options.route_lane)
        .. " status=" .. tostring(options.status)
        .. " location=" .. location
        .. " max=" .. max_redirects
        .. " keep_args=" .. tostring(keep_args)
        .. " keep_headers=" .. tostring(keep_headers)
        .. " host_cfg=" .. redirect_follow_host)

    local current_location = location
    local current_scheme = original_scheme
    local current_host = original_host
    local redirect_count = 0
    local last_redirect_status = tonumber(options.status) or 302

    while redirect_count < max_redirects do
        local url, scheme, host = parse_location(current_location, current_scheme, current_host)
        if not url then
            return {
                kind = "redirect",
                status = last_redirect_status,
                location = current_location,
                followed = redirect_count,
            }
        end

        if keep_args and original_args and original_args ~= "" then
            url = append_args(url, original_args)
        end

        local host_header = determine_host_header(host, redirect_follow_host, original_host)
        local request_url = url
        local request_cache_config
        local request_cache_key
        if should_use_loopback_follow(host_header, original_host, follow_cache_key) then
            request_url = replace_url_host(url, build_loopback_host(host))
            request_cache_config = follow_cache_config
            request_cache_key = follow_cache_key
        end
        local headers = build_headers(host_header, keep_headers, original_headers, request_cache_config, request_cache_key)

        redirect_count = redirect_count + 1
        ngx.log(ngx.INFO, "[RedirectFollow] Redirect #" .. redirect_count
            .. " -> " .. request_url
            .. " Host=" .. host_header
            .. " cache_key_override=" .. tostring(request_cache_key ~= nil))

        local httpc = http.new()
        httpc:set_timeout(30000)

        local res, err = httpc:request_uri(request_url, {
            method = method,
            headers = headers,
            ssl_verify = false,
        })

        if not res then
            return nil, "request failed: " .. tostring(err)
        end

        if is_redirect(res.status) then
            local new_location = res.headers["Location"] or res.headers["location"]
            if not new_location then
                return nil, "redirect without location header"
            end

            current_location = new_location
            current_scheme = scheme
            current_host = host
            last_redirect_status = res.status
            ngx.log(ngx.INFO, "[RedirectFollow] Continue -> " .. new_location)
        else
            ngx.log(ngx.INFO, "[RedirectFollow] Final status=" .. res.status .. " after " .. redirect_count .. " redirect(s)")
            return {
                kind = "response",
                status = res.status,
                headers = res.headers or {},
                body = res.body,
                followed = redirect_count,
                cache_meta = {
                    cache_status = res.headers["AGC-Cache-Status"] or res.headers["agc-cache-status"],
                    cache_policy = res.headers["AGC-Cache-Policy"] or res.headers["agc-cache-policy"],
                    age = res.headers["Age"] or res.headers["age"],
                    cache_time = res.headers["AGC-Swift-CacheTime"] or res.headers["agc-swift-cachetime"],
                    save_time = res.headers["AGC-Swift-SaveTime"] or res.headers["agc-swift-savetime"],
                }
            }
        end
    end

    ngx.log(ngx.WARN, "[RedirectFollow] Exceeded max redirects (" .. max_redirects .. ")")
    return {
        kind = "redirect",
        status = last_redirect_status,
        location = current_location,
        followed = redirect_count,
    }
end

return _M