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

local M = {}

local function client_get_desk_idx(c)
    if c.first_tag then
        local desk_idx = utils.tag_desk_idx(c.first_tag)
        if desk_idx ~= nil and desk_idx >= 0 then
            return desk_idx
        else
            utils.warn('UrgentWgt', '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_add(c)
    function handler_client_show(c)
        if not c.urgent then
            return
        end

        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
    function handler_client_hide(c)
        if self._clients[c] then
            self:_desk_client_remove(self._clients[c], c)
        end
    end
    function handler_client_remove(c)
        c:disconnect_signal('untagged',         handler_client_hide)
        c:disconnect_signal('tagged',           handler_client_show)
        c:disconnect_signal('unmanage',         handler_client_remove)
        c:disconnect_signal('property::urgent', handler_client_remove)

        handler_client_hide(c)
        self._clients[c] = nil
    end

    -- handle client state changes:
    -- client moved: hide it on its previous desktop, show on the new one
    c:connect_signal('untagged',         handler_client_hide)
    c:connect_signal('tagged',           handler_client_show)
    -- unurgent or destroyed: forget about client
    c:connect_signal('property::urgent', handler_client_remove)
    c:connect_signal('unmanage',         handler_client_remove)

    -- show the client on its desk right now
    handler_client_show(c)
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