local awful     = require("awful")
local beautiful = require("beautiful")
local shape     = require("gears.shape")
local wibox     = require("wibox")

local M = {}

local function client_get_desk_idx(c)
    if c.first_tag then
        local desk_idx = tonumber('0x' .. c.first_tag.name:sub(1, 2))
        if desk_idx ~= nil and desk_idx >= 0 then
            return desk_idx
        else
            print('cannot get desk idx: ' .. c.first_tag.name)
        end
    end
    return false
end

local function urgent_widget_create(name)
    local textbox = wibox.widget.textbox(' ' .. name .. ' ')
    local cont_bg = wibox.container.background(textbox)
    cont_bg.bg = beautiful.bg_urgent
    cont_bg.fg = beautiful.fg_urgent

    cont_bg:set_shape(shape.rounded_rect, 5)
    cont_bg.shape_border_width = 1
    cont_bg.shape_border_color = beautiful.border_focus

    return cont_bg
end

local UrgentWidget = {}

function UrgentWidget:_desk_client_add(desk_idx, c)
    -- create the client list, if it does not exist already
    if self._desktops[desk_idx] == nil then
        self._desktops[desk_idx] = {}
    end

    self._desktops[desk_idx][c] = true

    if self._desktop_widgets[desk_idx] == nil then
        self._desktop_widgets[desk_idx] = urgent_widget_create(desk_idx)
    end

    for k, w in pairs(self._layout:get_children()) do
        if w == self._desktop_widgets[desk_idx] then
            return
        end
    end
    self._layout:add(self._desktop_widgets[desk_idx])
end

function UrgentWidget:_desk_client_remove(desk_idx, c)
    self._desktops[desk_idx][c] = nil

    -- remove the widget if there are no clients left
    local count = 0
    for c, v in pairs(self._desktops[desk_idx]) do
        if v then
            count = count + 1
        end
    end
    if count == 0 then
        self._layout:remove_widgets(self._desktop_widgets[desk_idx])
    end
end

function UrgentWidget:_client_remove(c)
    if self._clients[c] then
        self:_desk_client_remove(self._clients[c], c)
    end
    self._clients[c] = nil
end

function UrgentWidget:_client_add(c)
    local desk_idx   = client_get_desk_idx(c)
    self._clients[c] = desk_idx

    if desk_idx then
        self:_desk_client_add(desk_idx, c)
    end

    -- handle client state changes:
    -- lose the urgency flag
    c:connect_signal('property::urgent',
        function (c)
            if not c.urgent then
                self:_client_remove(c)
            end
        end)
    -- client destroyed
    c:connect_signal('unmanage',
        function (c)
            self:_client_remove(c)
        end)
    -- client moved
    c:connect_signal('untagged',
        function(c, t)
            local desk_idx = self._clients[c]
            if desk_idx then
                self:_desk_client_remove(desk_idx, c)
            end
            self._clients[c] = false
        end)
    c:connect_signal('tagged',
        function(c, t)
            local desk_idx   = client_get_desk_idx(c)
            self._clients[c] = desk_idx

            if desk_idx then
                self:_desk_client_add(desk_idx, c)
            end
        end)
end

function UrgentWidget:new()
    local ret = setmetatable({}, self)
    self.__index = self

    ret._layout = wibox.layout.fixed.horizontal()
    ret.widget  = awful.widget.only_on_screen(ret._layout, "primary")

    -- maps from clients to their desk num (or false, when no desk num
    -- can be determined)
    ret._clients  = {}
    -- contains a table of urgent clients per each desktop number
    ret._desktops = {}
    ret._desktop_widgets = {}

    -- signal handler for adding new clients
    client.connect_signal('property::urgent',
        function (c)
            -- only do something when urgency flag is on
            if c.urgent then
                ret:_client_add(c)
            end
        end)

    return ret
end

M.UrgentWidget = UrgentWidget

return M