-- Based on https://github.com/Klaessen/openresty-loadbalancers/blob/main/sticky-balancer.lua

local _M = {}

local balancer = require "ngx.balancer"
local cookie_name = "INVIDIOUS_SERVER_ID"
local servers
local weighted_servers
local domain

local function set_headers(server_index)
    -- https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie#partitioned
    ngx.header["Set-Cookie"] = cookie_name .. "=" .. server_index .. "; domain=" .. domain .. "; Path=/; HttpOnly; SameSite=None; Secure; Partitioned"
    ngx.header["X-Server-Id"] = server_index
end

-- Generate a weighted server list based on weights
local function generate_weighted_server_list(servers)
  local weighted_servers = {}
  for _, server in ipairs(servers) do
    for i = 1, server.weight do
      table.insert(weighted_servers, server)
    end
  end
  return weighted_servers
end

-- Select server based on cookie or assign a new one
local function select_server()
  local cookie = ngx.var["cookie_" .. cookie_name]
  local server_index

  math.randomseed(os.time())

  if cookie then
    server_index = tonumber(cookie)
    ngx.header["X-Server-Id"] = server_index
  else
    server_index = math.random(#servers)
    set_headers(server_index)
  end

  local server = weighted_servers[server_index]
  return server
end

local function check_args(args)
  for key, server_index in pairs(args) do
    if key == "backend" then
      server_index = tonumber(server_index)

      -- To redirect to another backend if user inputs a backend that doesn't exists
      -- Ex: ?backend=4 will give you X-Server-Id=1 (Backend 1)
      local val = val % #servers
      if val == -1 then
        val = #servers
      end

      ok, err = balancer.set_current_peer(servers[val][0], servers[val][2])

      if not ok then
        return ngx.exit(501)
      end

      set_headers()
    end
  end
end

function _M.run(upstreams)
  domain = ".nadeko.net"
  local host = ngx.req.get_headers()["Host"]

  -- TOR Support
  if string.match(host, ".onion") then
    domain = host
  end

  -- I2P Support
  if string.match(host, ".i2p") then
    domain = host
  end

  servers = upstreams
  weighted_servers = generate_weighted_server_list(servers)
  local ok, err

  local args = ngx.req.get_uri_args()
  if args then
    check_args(args)
  end

  local server = select_server()
  if not server then
    ngx.log(ngx.ERR, "Failed to set the current peer: ", err)
    return ngx.exit(502)
  end

  -- Unix socket support
  if string.match(server[1], 'unix:') then
    ok, err = balancer.set_current_peer(server[1])
  else
    ok, err = balancer.set_current_peer(server[1], server[2])
  end

  if not ok then
    ngx.log(ngx.ERR, "Failed to set the current peer: ", err)
    return ngx.exit(500)
  end

  -- https://github.com/openresty/lua-resty-core/blob/master/lib/ngx/balancer.md#enable_keepalive
  ok, err = balancer.enable_keepalive(60, 1000)
  if not ok then
    ngx.log(ngx.ERR, "Failed to set keepalive: ", err)
    return
  end
end

return _M