-- Copyright (C) Yichun Zhang (agentzh)


-- FIXME: this library is very rough and is currently just for testing
--        the websocket client.


local wbproto = require "nginx.websocket.protocol"
local bit = require "bit"


local _recv_frame = wbproto.recv_frame
local _send_frame = wbproto.send_frame
local new_tab = wbproto.new_tab
local tcp = ngx.socket.tcp
local re_match = ngx.re.match
local re_find  = ngx.re.find
local encode_base64 = ngx.encode_base64
local concat = table.concat
local char = string.char
local str_find = string.find
local rand = math.random
local rshift = bit.rshift
local band = bit.band
local setmetatable = setmetatable
local type = type
local debug = ngx.config.debug
local ngx_log = ngx.log
local ngx_DEBUG = ngx.DEBUG
local assert = assert
local ssl_support = true

if not ngx.config
    or not ngx.config.ngx_lua_version
    or ngx.config.ngx_lua_version < 9011
then
    ssl_support = false
end

local _M = new_tab(0, 13)
_M._VERSION = '0.13'


local mt = { __index = _M }


function _M.new(self, opts)
    local sock, err = tcp()
    if not sock then
        return nil, err
    end

    local max_payload_len, send_unmasked, timeout
    local max_recv_len, max_send_len
    if opts then
        max_payload_len = opts.max_payload_len
        max_recv_len = opts.max_recv_len
        max_send_len = opts.max_send_len

        send_unmasked = opts.send_unmasked
        timeout = opts.timeout

        if timeout then
            sock:settimeout(timeout)
        end
    end

    max_payload_len = max_payload_len or 65535
    max_recv_len = max_recv_len or max_payload_len
    max_send_len = max_send_len or max_payload_len

    return setmetatable({
        sock = sock,
        max_recv_len = max_recv_len,
        max_send_len = max_send_len,
        send_unmasked = send_unmasked,
    }, mt)
end


function _M.connect(self, uri, opts)
    local sock = self.sock
    if not sock then
        return nil, "not initialized"
    end

    local is_unix = false

    local m, err
    if re_find(uri, "unix:") then
        is_unix = true
        m, err = re_match(uri, [[^(wss?)://(unix:[^:]+):()(.*)]], "jo")
    else
        m, err = re_match(uri, [[^(wss?)://([^:/]+)(?::(\d+))?(.*)]], "jo")
    end

    if not m then
        if err then
            return nil, "failed to match the uri: " .. err
        end

        return nil, "bad websocket uri"
    end

    local scheme = m[1]
    local addr = m[2]
    local port = m[3]
    local path = m[4]

    -- ngx.say("host: ", host)
    -- ngx.say("port: ", port)

    local ssl = scheme == "wss"
    if ssl and not ssl_support then
        return nil, "ngx_lua 0.9.11+ required for SSL sockets"
    end

    if not port then
        port = ssl and 443 or 80
    end

    if path == "" then
        path = "/"
    end

    local ssl_verify, server_name, headers, proto_header, origin_header
    local sock_opts = {}
    local client_cert, client_priv_key
    local header_host
    local key

    if opts then
        local protos = opts.protocols
        if protos then
            if type(protos) == "table" then
                proto_header = "\r\nSec-WebSocket-Protocol: "
                               .. concat(protos, ",")

            else
                proto_header = "\r\nSec-WebSocket-Protocol: " .. protos
            end
        end

        local origin = opts.origin
        if origin then
            origin_header = "\r\nOrigin: " .. origin
        end

        if opts.pool then
            sock_opts.pool = opts.pool
        end
        --pool_size specify the size of the connection pool. If omitted and no backlog option was provided, no pool will be created.
        if opts.pool_size then
            sock_opts.pool_size = opts.pool_size
        end
        if opts.backlog then
            sock_opts.backlog = opts.backlog
        end


        client_cert = opts.client_cert
        client_priv_key = opts.client_priv_key

        if client_cert then
            assert(client_priv_key,
                   "client_priv_key must be provided with client_cert")
        end

        ssl_verify = opts.ssl_verify

        server_name = opts.server_name
        if server_name ~= nil and type(server_name) ~= "string" then
            return nil, "SSL server_name must be a string"
        end

        if opts.headers then
            headers = opts.headers
            if type(headers) ~= "table" then
                return nil, "custom headers must be a table"
            end
        end

        header_host = opts.host
        if header_host ~= nil and type(header_host) ~= "string" then
            return nil, "custom host header must be a string"
        end

        key = opts.key
        if key ~= nil and type(key) ~= "string" then
            return nil, "custom Sec-WebSocket-Key must be a string"
        end
    end

    local ok, err
    if is_unix then
        ok, err = sock:connect(addr, sock_opts)
    else
        ok, err = sock:connect(addr, port, sock_opts)
    end
    if not ok then
        return nil, "failed to connect: " .. err
    end

    -- check for connections from pool:
    local reused_count, err = sock:getreusedtimes()
    if not reused_count then
        return nil, "failed to get reused times: " .. tostring(err)
    end

    if reused_count > 0 then
        -- being a reused connection (must have done handshake)
        return 1, nil, "connection reused"
    end

    if ssl then
        if client_cert then
            ok, err = sock:setclientcert(client_cert, client_priv_key)
            if not ok then
                return nil, "failed to set TLS client certificate: " .. err
            end
        end

        server_name = server_name or header_host or addr

        ok, err = sock:sslhandshake(false, server_name, ssl_verify)
        if not ok then
            return nil, "ssl handshake failed: " .. err
        end
    end

    local custom_headers
    if headers then
        custom_headers = concat(headers, "\r\n")
        custom_headers = "\r\n" .. custom_headers
    end

    -- do the websocket handshake:

    if not key then
        local bytes = char(rand(256) - 1, rand(256) - 1, rand(256) - 1,
                           rand(256) - 1, rand(256) - 1, rand(256) - 1,
                           rand(256) - 1, rand(256) - 1, rand(256) - 1,
                           rand(256) - 1, rand(256) - 1, rand(256) - 1,
                           rand(256) - 1, rand(256) - 1, rand(256) - 1,
                           rand(256) - 1)

        key = encode_base64(bytes)
    end

    local host_header = header_host 
                        or (is_unix and "unix_sock" or addr .. ":" .. port)

    local req = "GET " .. path .. " HTTP/1.1\r\nUpgrade: websocket\r\nHost: "
                .. host_header
                .. "\r\nSec-WebSocket-Key: " .. key
                .. (proto_header or "")
                .. "\r\nSec-WebSocket-Version: 13"
                .. (origin_header or "")
                .. "\r\nConnection: Upgrade"
                .. (custom_headers or "")
                .. "\r\n\r\n"

    local bytes, err = sock:send(req)
    if not bytes then
        return nil, "failed to send the handshake request: " .. err
    end

    local header_reader = sock:receiveuntil("\r\n\r\n")
    -- FIXME: check for too big response headers
    local header, err, partial = header_reader()
    if not header then
        return nil, "failed to receive response header: " .. err
    end

    -- error("header: " .. header)

    -- FIXME: verify the response headers

    m, err = re_match(header, [[^\s*HTTP/1\.1\s+]], "jo")
    if not m then
        return nil, "bad HTTP response status line: " .. header
    end

    return 1, nil, header
end


function _M.set_timeout(self, time)
    local sock = self.sock
    if not sock then
        return nil, nil, "not initialized yet"
    end

    return sock:settimeout(time)
end


function _M.recv_frame(self)
    if self.fatal then
        return nil, nil, "fatal error already happened"
    end

    local sock = self.sock
    if not sock then
        return nil, nil, "not initialized yet"
    end

    local data, typ, err =  _recv_frame(sock, self.max_recv_len, false)
    if not data and not str_find(err, ": timeout", 1, true) then
        self.fatal = true
    end
    return data, typ, err
end


local function send_frame(self, fin, opcode, payload)
    if self.fatal then
        return nil, "fatal error already happened"
    end

    if self.closed then
        return nil, "already closed"
    end

    local sock = self.sock
    if not sock then
        return nil, "not initialized yet"
    end

    local bytes, err = _send_frame(sock, fin, opcode, payload,
                                   self.max_send_len,
                                   not self.send_unmasked)
    if not bytes then
        self.fatal = true
    end
    return bytes, err
end
_M.send_frame = send_frame


function _M.send_text(self, data)
    return send_frame(self, true, 0x1, data)
end


function _M.send_binary(self, data)
    return send_frame(self, true, 0x2, data)
end


local function send_close(self, code, msg)
    local payload
    if code then
        if type(code) ~= "number" or code > 0x7fff then
            return nil, "bad status code"
        end
        payload = char(band(rshift(code, 8), 0xff), band(code, 0xff))
                        .. (msg or "")
    end

    if debug then
        ngx_log(ngx_DEBUG, "sending the close frame")
    end

    local bytes, err = send_frame(self, true, 0x8, payload)

    if not bytes then
        self.fatal = true
    end

    self.closed = true

    return bytes, err
end
_M.send_close = send_close


function _M.send_ping(self, data)
    return send_frame(self, true, 0x9, data)
end


function _M.send_pong(self, data)
    return send_frame(self, true, 0xa, data)
end


function _M.close(self)
    if self.fatal then
        return nil, "fatal error already happened"
    end

    local sock = self.sock
    if not sock then
        return nil, "not initialized"
    end

    if not self.closed then
        local bytes, err = send_close(self)
        if not bytes then
            return nil, "failed to send close frame: " .. err
        end
    end

    return sock:close()
end


function _M.set_keepalive(self, ...)
    local sock = self.sock
    if not sock then
        return nil, "not initialized"
    end

    return sock:setkeepalive(...)
end


return _M