local helpers = require "spec.helpers"
local cjson = require "cjson.safe"

local REDIS_HOST = helpers.redis_host
local REDIS_PORT = helpers.redis_port
local REDIS_DATABASE = 1
local upstream_port = helpers.get_available_port()

local function flush_redis()
    local redis = require "resty.redis"
    local red = redis:new()
    red:set_timeout(2000)
    assert(red:connect(REDIS_HOST, REDIS_PORT))
    assert(red:select(REDIS_DATABASE))
    assert(red:flushall())
    red:close()
end

local function upstream_fixture(port)
    return [[
        server {
            listen ]] .. port .. [[;
            location /hold {
                content_by_lua_block {
                    ngx.sleep(0.8)
                    ngx.status = 200
                    ngx.say("ok")
                }
            }
            location /hold-past-lease-ttl {
                content_by_lua_block {
                    ngx.sleep(2.2)
                    ngx.status = 200
                    ngx.say("ok")
                }
            }
            location /fast {
                content_by_lua_block {
                    ngx.status = 200
                    ngx.say("ok")
                }
            }
        }
    ]]
end

local function request_headers(user, org, accept)
    local headers = {
        ["Host"] = "investibook-sse-limit.test",
        ["x-test-user"] = user,
        ["Accept"] = accept or "text/event-stream"
    }
    if org then
        headers["x-test-org"] = org
    end
    return headers
end

local function request(path, user, org, accept)
    local client = helpers.proxy_client()
    local headers = request_headers(user, org, accept)
    local res = assert(client:send({
        method = "GET",
        path = path,
        headers = headers
    }))
    local status = res.status
    local body = res:read_body()
    client:close()
    return status, body
end

for _, strategy in helpers.each_strategy() do
    describe("Plugin: investibook-sse-limit [#" .. strategy .. "]", function()
        local proxy_client

        lazy_setup(function()
            flush_redis()
            local bp = helpers.get_db_utils(strategy, {"routes", "services", "plugins"})
            local service = bp.services:insert({
                protocol = "http",
                host = "127.0.0.1",
                port = upstream_port
            })
            local route = bp.routes:insert({
                protocols = {"http"},
                hosts = {"investibook-sse-limit.test"},
                service = service
            })
            bp.plugins:insert({
                route = {id = route.id},
                name = "investibook-sse-limit",
                config = {
                    identity_header_name = "x-test-user",
                    org_header_name = "x-test-org",
                    user_concurrency = 1,
                    org_concurrency = 1,
                    global_concurrency = 2,
                    connections_per_minute = 3,
                    lease_ttl_seconds = 1,
                    retry_after_seconds = 5,
                    key_prefix = "investibook:sse:test",
                    fault_tolerant = false,
                    redis_host = REDIS_HOST,
                    redis_port = REDIS_PORT,
                    redis_database = REDIS_DATABASE
                }
            })

            assert(helpers.start_kong({
                database = strategy,
                nginx_conf = "spec/fixtures/custom_nginx.template",
                plugins = "investibook-sse-limit",
                nginx_http_include = upstream_fixture(upstream_port)
            }))
        end)

        lazy_teardown(function()
            helpers.stop_kong()
        end)

        before_each(function()
            flush_redis()
            proxy_client = helpers.proxy_client()
        end)

        after_each(function()
            if proxy_client then
                proxy_client:close()
            end
        end)

        it("limits concurrent streams per user and releases the lease", function()
            local first = ngx.thread.spawn(function()
                return request("/hold", "user-concurrency", "org-a")
            end)
            ngx.sleep(0.15)

            local rejected = assert(proxy_client:send({
                method = "GET",
                path = "/hold",
                headers = request_headers("user-concurrency", "org-a")
            }))
            assert.res_status(429, rejected)
            assert.equal("5", rejected.headers["Retry-After"])
            local body = cjson.decode(assert.res_body(rejected))
            assert.equal("Too many streaming connections. Please retry later.", body.detail)

            local ok, status = ngx.thread.wait(first)
            assert.is_true(ok)
            assert.equal(200, status)

            local after_release = assert(proxy_client:send({
                method = "GET",
                path = "/fast",
                headers = request_headers("user-concurrency", "org-a")
            }))
            assert.res_status(200, after_release)
        end)

        it("skips non-SSE requests without consuming limiter capacity", function()
            for _ = 1, 5 do
                local status = request("/fast", "mixed-user", "mixed-org",
                                       "application/json")
                assert.equal(200, status)
            end

            for _ = 1, 3 do
                local status = request("/fast", "mixed-user", "mixed-org")
                assert.equal(200, status)
            end
            local status = request("/fast", "mixed-user", "mixed-org")
            assert.equal(429, status)
        end)

        it("recognizes text/event-stream in a combined Accept header", function()
            for _ = 1, 3 do
                local status = request("/fast", "combined-accept-user", nil,
                                       "application/json, Text/Event-Stream, */*")
                assert.equal(200, status)
            end
            local status = request("/fast", "combined-accept-user", nil,
                                   "application/json, Text/Event-Stream, */*")
            assert.equal(429, status)
        end)

        it("renews the lease while a stream remains active", function()
            local first = ngx.thread.spawn(function()
                return request("/hold-past-lease-ttl", "heartbeat-user",
                               "heartbeat-org")
            end)
            ngx.sleep(1.3)

            local status, body = request("/fast", "heartbeat-user",
                                         "heartbeat-org")
            assert.equal(429, status)
            assert.equal("Too many streaming connections. Please retry later.",
                         cjson.decode(body).detail)

            local ok, first_status = ngx.thread.wait(first)
            assert.is_true(ok)
            assert.equal(200, first_status)

            local after_release = request("/fast", "heartbeat-user",
                                          "heartbeat-org")
            assert.equal(200, after_release)
        end)

        it("limits concurrent streams per org", function()
            local first = ngx.thread.spawn(function()
                return request("/hold", "org-user-a", "shared-org")
            end)
            ngx.sleep(0.15)

            local status, body = request("/fast", "org-user-b", "shared-org")
            assert.equal(429, status)
            assert.equal("Too many streaming connections. Please retry later.",
                         cjson.decode(body).detail)

            local ok, first_status = ngx.thread.wait(first)
            assert.is_true(ok)
            assert.equal(200, first_status)
        end)

        it("limits concurrent streams globally", function()
            local first = ngx.thread.spawn(function()
                return request("/hold", "global-user-a", "global-org-a")
            end)
            local second = ngx.thread.spawn(function()
                return request("/hold", "global-user-b", "global-org-b")
            end)
            ngx.sleep(0.15)

            local status = request("/fast", "global-user-c", "global-org-c")
            assert.equal(429, status)

            local ok_first, first_status = ngx.thread.wait(first)
            local ok_second, second_status = ngx.thread.wait(second)
            assert.is_true(ok_first)
            assert.is_true(ok_second)
            assert.equal(200, first_status)
            assert.equal(200, second_status)
        end)

        it("limits connection attempts per minute", function()
            for _ = 1, 3 do
                local status = request("/fast", "rate-user", nil)
                assert.equal(200, status)
            end
            local status, body = request("/fast", "rate-user", nil)
            assert.equal(429, status)
            assert.equal("Too many streaming connections. Please retry later.",
                         cjson.decode(body).detail)
        end)
    end)
end