-- SPDX-FileCopyrightText: © 2022—2026 Vladimir Zorin <vladimir@deviant.guru>
-- SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later
-- Licensed under OWL v1.0+. See LICENSE.

--[==[
Async Redis client using the lev event loop.
Must be called within `lev.run()`.
]==]

local lev = require("lev")
local resp = require("redis.resp")

local DEFAULT_POOL_SIZE = 20
local POOL_IDLE_THRESHOLD = 5
local POOL_MAX_IDLE = 60

local connection_pool = {}
local pool_size = tonumber(os.getenv("REDIS_CLIENT_POOL_SIZE")) or DEFAULT_POOL_SIZE
if pool_size < 1 then
	pool_size = DEFAULT_POOL_SIZE
end

local force_close = function(client)
	local state = client.__state
	if state.closed then
		return
	end
	state.closed = true
	state.in_pool = false
	state.socket:close()
end

local mark_broken = function(client)
	client.__state.broken = true
	force_close(client)
end

local prune_pool = function(pool_key)
	local pool = connection_pool[pool_key]
	if not pool then
		return
	end

	local now = lev.now()
	local kept = {}
	for _, client in ipairs(pool) do
		local state = client.__state
		local idle = now - (state.pooled_at or state.last_used or 0)
		if state.broken or state.closed or idle >= POOL_MAX_IDLE then
			force_close(client)
		else
			kept[#kept + 1] = client
		end
	end
	connection_pool[pool_key] = kept
end

---! Execute a Redis command and return the parsed response
---@ client:cmd(`...`) -> `result`, `err`{.str .err}
local redis_command = function(self, ...)
	local cmd = resp.encode_command(...)
	local ok, err = self.__state.socket:send(cmd)
	if not ok then
		mark_broken(self)
		return nil, err
	end
	local recv_line = self.__state.recv_line
	local recv_bytes = self.__state.recv_bytes
	local result, rerr = resp.read_response(recv_line, recv_bytes)
	if result then
		if result.type == "error" then
			return nil, result.value
		end
		self.__state.last_used = lev.now()
		return result.value
	end
	if rerr ~= "not found" then
		mark_broken(self)
	end
	return nil, rerr
end

---! Execute multiple Redis commands in a single round-trip (pipelining)
---@ client:pipeline(`commands`{.tbl}) -> `results`{.tbl .or_nil}, `err`{.str .err}
--[===[
Sends all commands at once and reads all responses. Each entry in
`commands` is an array of arguments (same as `cmd` varargs).
Returns an array of result tables `{ value, err }` in the same order.
]===]
local redis_pipeline = function(self, commands)
	if not commands or #commands == 0 then
		return {}
	end
	local parts = {}
	for _, cmd_args in ipairs(commands) do
		parts[#parts + 1] = resp.encode_command(unpack(cmd_args))
	end
	local ok, err = self.__state.socket:send(table.concat(parts))
	if not ok then
		mark_broken(self)
		return nil, err
	end
	local recv_line = self.__state.recv_line
	local recv_bytes = self.__state.recv_bytes
	local results = {}
	for i = 1, #commands do
		local result, rerr = resp.read_response(recv_line, recv_bytes)
		if not result then
			if rerr == "not found" then
				-- NULL bulk string (e.g. GET on non-existent key)
				results[i] = { nil }
			else
				mark_broken(self)
				return nil, rerr or ("failed to read response " .. i)
			end
		elseif result.type == "error" then
			results[i] = { nil, result.value }
		else
			results[i] = { result.value }
		end
	end
	self.__state.last_used = lev.now()
	return results
end

---! Read a raw RESP response from the connection
---@ client:read() -> `response`{.tbl .or_nil}, `err`{.str .err}
local read = function(self)
	local result, err = resp.read_response(self.__state.recv_line, self.__state.recv_bytes)
	if result then
		self.__state.last_used = lev.now()
	elseif err ~= "not found" then
		mark_broken(self)
	end
	return result, err
end

---! Close the Redis connection
---@ client:close(`no_keepalive`{.bool .opt}) -> `ok`{.bool}
--[===[
Closes the connection. By default the underlying socket is returned
to the connection pool for reuse. Pass `true` to force-close the
socket instead of pooling it. The socket is also force-closed when
the pool exceeds its size limit.
]===]
local close = function(self, no_keepalive)
	local state = self.__state
	local pool_key = state.pool_key
	if state.in_pool then
		return true
	end
	if not pool_key or no_keepalive or state.broken or state.closed then
		force_close(self)
		return true
	end
	prune_pool(pool_key)
	if #(connection_pool[pool_key] or {}) >= pool_size then
		force_close(self)
		return true
	end
	connection_pool[pool_key] = connection_pool[pool_key] or {}
	state.pooled_at = lev.now()
	state.in_pool = true
	table.insert(connection_pool[pool_key], self)
	return true
end

---! Connect to a Redis server asynchronously
---@ connect(`config`{.tbl .str .opt}) -> `client`{.tbl .or_nil}, `err`{.str .err}
--[===[
Establishes an async connection to a Redis server, returning a client
object. Must be called within `lev.run()`.

`config`, when provided, can be either a string URL, or a table.
]===]
local connect = function(config)
	local conf, conf_err = resp.normalize_config(config)
	if not conf then
		return nil, conf_err
	end

	local pool_key = conf.host .. ":" .. conf.port .. "/" .. (conf.db or "0") .. (conf.ssl and "+ssl" or "")

	-- Try reuse from pool
	prune_pool(pool_key)
	local pool = connection_pool[pool_key]
	while pool and #pool > 0 do
		local client = table.remove(pool)
		local state = client.__state
		state.in_pool = false
		if state.broken or state.closed then
			force_close(client)
		elseif lev.now() - (state.last_used or 0) < POOL_IDLE_THRESHOLD then
			return client
		else
			local pong, perr = client:cmd("PING")
			if pong == "PONG" then
				return client
			end
			mark_broken(client)
		end
	end
	connection_pool[pool_key] = connection_pool[pool_key] or {}

	local tls_opts = nil
	if conf.ssl then
		tls_opts = { mode = "client" }
	end

	local sock, err = lev.connect(conf.host, conf.port, {
		timeout = conf.timeout,
		tls = tls_opts,
	})
	if not sock then
		return nil, err
	end
	sock:setsockopt("nodelay", 1)

	-- Async I/O callbacks for RESP parser
	local recv_line = function()
		local data, lerr = sock:recv_until("\r\n")
		if not data then
			return nil, lerr
		end
		-- Strip trailing \r\n
		if data:sub(-2) == "\r\n" then
			data = data:sub(1, -3)
		end
		return data
	end
	local recv_bytes = function(n)
		return sock:recv_exactly(n)
	end

	local obj = {
		cfg = conf,
		__state = {
			socket = sock,
			pool_key = pool_key,
			recv_line = recv_line,
			recv_bytes = recv_bytes,
		},
		cmd = redis_command,
		pipeline = redis_pipeline,
		read = read,
		close = close,
	}

	-- AUTH
	if conf.auth then
		local aok, aerr = obj:cmd("AUTH", conf.auth.user, conf.auth.pass)
		if not aok then
			sock:close()
			return nil, aerr
		end
	end

	-- SELECT db
	if conf.db then
		local sok, serr = obj:cmd("SELECT", conf.db)
		if not sok then
			sock:close()
			return nil, serr
		end
	end

	return obj
end

return { connect = connect }