local websocket = require "http.websocket"
local copas = require "copas"
local cjson = require "cjson"

WebSocketClient = {}
WebSocketClient.__index = WebSocketClient

function WebSocketClient.new(conf)
    local self = setmetatable({}, WebSocketClient)
    local conf = conf or {}
    self.uri = conf.uri or nil
    self.ws = nil

    -- Set custom callback functions if provided
    self.on_connected = conf.on_connected or self.on_connected
    self.on_disconnected = conf.on_disconnected or self.on_disconnected
    self.on_received = conf.on_received or self.on_received

    self:connect()
    copas.addthread(function() self:receive() end)
    
    return self
end

function WebSocketClient:connect()
    self.ws = websocket.new_from_uri(self.uri)
    assert(self.ws:connect())
    self:on_connected()
end

function WebSocketClient:disconnect()
    if self.ws then
        self.ws:close()
        self:on_disconnected()
    end
end

function WebSocketClient:send(data)
    data = data or {}

    local commandAPI = data.command or 'error'
    local dataAPI = data.body or nil
    local finalData = nil

    if type(dataAPI) == 'table' then
        finalData = cjson.encode(dataAPI)
    else
        finalData = dataAPI or 'Empty'
    end

    -- print(commandAPI)
    -- print(dataAPI)

    local message = {
        event = commandAPI,
        data = finalData
    }

    local message_json = cjson.encode(message)

    if self.ws then
        assert(self.ws:send(message_json))
    end
end

function WebSocketClient:receive()
    while true do
        local data, opcode, is_fin = self.ws:receive()
        if data then
            self:on_received(data)
            if data == "close_connection" then
                self:disconnect()
                break
            end
        else
            self:on_disconnected()
            break
        end
        copas.sleep(0) 
    end
end

function WebSocketClient:on_connected()
    print("Connected to the server.")
end

function WebSocketClient:on_disconnected()
    print("Disconnected from the server.")
    while true do
        print("Attempting to reconnect...")
        local success, err = pcall(self.connect, self)
        if success then
            print("Reconnected to the server.")
            copas.addthread(function() self:receive() end)
            break
        else
            print("Failed to reconnect, error: " .. tostring(err))
            print("Retrying in 5 seconds...")
            copas.sleep(5) -- Wait for 5 seconds before attempting to reconnect
        end
    end
end

function WebSocketClient:on_received(data)
    print("Received data:", data)
end

return WebSocketClient