-----------------------------------------------------------------------------
-- LuaSocket helper module
-- Author: Diego Nehab
-----------------------------------------------------------------------------
-----------------------------------------------------------------------------
-- Declare module and import dependencies
-----------------------------------------------------------------------------
local base = _G
local string = require("string")
local math = require("math")
local socket = require("socket.core")
local _M = socket

-----------------------------------------------------------------------------
-- Exported auxiliar functions
-----------------------------------------------------------------------------
function _M.connect4(address, port, laddress, lport)
	return socket.connect(address, port, laddress, lport, "inet")
end

function _M.connect6(address, port, laddress, lport)
	return socket.connect(address, port, laddress, lport, "inet6")
end

function _M.bind(host, port, backlog)
	if host == "*" then host = "0.0.0.0" end

	local addrinfo, err = socket.dns.getaddrinfo(host)

	if not addrinfo then return nil, err end

	local sock, res
	err = "no info on address"

	for i, alt in base.ipairs(addrinfo) do
		if alt.family == "inet" then
			sock, err = socket.tcp()
		else
			sock, err = socket.tcp6()
		end

		if not sock then return nil, err end

		sock:setoption("reuseaddr", true)
		res, err = sock:bind(alt.addr, port)

		if not res then
			sock:close()
		else
			res, err = sock:listen(backlog)

			if not res then sock:close() else return sock end
		end
	end

	return nil, err
end

_M.try = _M.newtry()

function _M.choose(table)
	return function(name, opt1, opt2)
		if base.type(name) ~= "string" then
			name, opt1, opt2 = "default", name, opt1
		end

		local f = table[name or "nil"]

		if not f then
			base.error("unknown key (" .. base.tostring(name) .. ")", 3)
		else
			return f(opt1, opt2)
		end
	end
end

-----------------------------------------------------------------------------
-- Socket sources and sinks, conforming to LTN12
-----------------------------------------------------------------------------
-- create namespaces inside LuaSocket namespace
local sourcet, sinkt = {}, {}
_M.sourcet = sourcet
_M.sinkt = sinkt
_M.BLOCKSIZE = 2048
sinkt["close-when-done"] = function(sock)
	return base.setmetatable(
		{
			getfd = function()
				return sock:getfd()
			end,
			dirty = function()
				return sock:dirty()
			end,
		},
		{
			__call = function(self, chunk, err)
				if not chunk then
					sock:close()
					return 1
				else
					return sock:send(chunk)
				end
			end,
		}
	)
end
sinkt["keep-open"] = function(sock)
	return base.setmetatable(
		{
			getfd = function()
				return sock:getfd()
			end,
			dirty = function()
				return sock:dirty()
			end,
		},
		{
			__call = function(self, chunk, err)
				if chunk then return sock:send(chunk) else return 1 end
			end,
		}
	)
end
sinkt["default"] = sinkt["keep-open"]
_M.sink = _M.choose(sinkt)
sourcet["by-length"] = function(sock, length)
	return base.setmetatable(
		{
			getfd = function()
				return sock:getfd()
			end,
			dirty = function()
				return sock:dirty()
			end,
		},
		{
			__call = function()
				if length <= 0 then return nil end

				local size = math.min(socket.BLOCKSIZE, length)
				local chunk, err = sock:receive(size)

				if err then return nil, err end

				length = length - string.len(chunk)
				return chunk
			end,
		}
	)
end
sourcet["until-closed"] = function(sock)
	local done
	return base.setmetatable(
		{
			getfd = function()
				return sock:getfd()
			end,
			dirty = function()
				return sock:dirty()
			end,
		},
		{
			__call = function()
				if done then return nil end

				local chunk, err, partial = sock:receive(socket.BLOCKSIZE)

				if not err then
					return chunk
				elseif err == "closed" then
					sock:close()
					done = 1
					return partial
				else
					return nil, err
				end
			end,
		}
	)
end
sourcet["default"] = sourcet["until-closed"]
_M.source = _M.choose(sourcet)
return _M