[DFHack][Lua] Dwarf Job Monitor Panel - Dwarf Activity/Job Heuristics Monitor

jobmonitor.lua

-- Dwarf Job Monitor - A floating resizable panel that shows dwarf job statistics
-- Author: Louis
-- Version: 1.0

local gui = require('gui')
local widgets = require('gui.widgets')
local utils = require('utils')

-- Main Screen Class
DwarfJobMonitor = defclass(DwarfJobMonitor, gui.ZScreen)
DwarfJobMonitor.ATTRS{
    pass_movement_keys = true,
    pass_mouse_clicks = true,
    defocusable = true,
}

function DwarfJobMonitor:init()
    -- Initialize pin state
    self.is_pinned = false
    
    -- Initialize data structures
    self.job_stats = {}
    self.job_units = {} -- Track units for each job type
    self.job_cycle_index = {} -- Track current cycle position for each job
    self.last_refresh = 0
    
    -- Do preliminary data gathering for sizing (without UI updates)
    self:gatherJobData()
    
    -- Calculate optimal frame size based on initial content
    local optimal_size = self:calculateOptimalFrameSize()
    
    -- Create the main resizable panel
    self:addviews{
        widgets.Window{
            view_id = 'main_window',
            frame = {t=5, l=5, w=optimal_size.w, h=optimal_size.h},
            frame_title = 'Dwarf Job Monitor',
            frame_style = gui.WINDOW_FRAME,
            frame_background = gui.CLEAR_PEN,
            draggable = true,
            resizable = true,
            resize_min = {w=35, h=15},
            
            subviews = {
                widgets.Panel{
                    view_id = 'header_panel',
                    frame = {t=0, l=0, r=0, h=3},
                    subviews = {
                        widgets.Label{
                            frame = {t=0, l=0},
                            text = {
                                {text='Total Dwarfs: ', pen=COLOR_CYAN},
                                {text=function() return tostring(self:getTotalDwarvesCount()) end, pen=COLOR_WHITE}
                            }
                        },
                        widgets.Label{
                            frame = {t=1, l=0},
                            text = {
                                {text='Idle: ', pen=COLOR_YELLOW},
                                {text=function() return tostring(self:getIdleDwarvesCount()) end, pen=COLOR_WHITE},
                                {text=' | Working: ', pen=COLOR_YELLOW},
                                {text=function() return tostring(self:getWorkingDwarvesCount()) end, pen=COLOR_WHITE}
                            }
                        },
                        widgets.HotkeyLabel{
                            frame = {t=2, l=0},
                            label = 'Refresh',
                            key = 'CUSTOM_R',
                            auto_width = true,
                            on_activate = function() self:refreshData() end,
                        },
                        widgets.HotkeyLabel{
                            frame = {t=2, l=15},
                            label = function() return self.is_pinned and 'Unpin' or 'Pin' end,
                            key = 'CUSTOM_P',
                            auto_width = true,
                            on_activate = function() self:togglePin() end,
                        },
                        widgets.HotkeyLabel{
                            frame = {t=2, r=0},
                            label = 'Close',
                            key = 'LEAVESCREEN',
                            auto_width = true,
                            on_activate = function() self:dismiss() end,
                        },
                    }
                },
                  widgets.Panel{
                    view_id = 'content_panel',
                    frame = {t=4, l=0, r=0, b=0},
                    subviews = {
                        widgets.List{
                            view_id = 'job_list',
                            frame = {t=0, l=0, r=0, b=0},
                            scroll_keys = widgets.STANDARDSCROLL,
                        }
                    }
                }
            }
        }
    }
    
    -- Now that widgets are created, refresh the data and update the UI
    self:refreshData()
end

function DwarfJobMonitor:gatherJobData()
    -- This method gathers job data without updating UI (used for initial sizing)
    self.job_stats = {}
    self.job_units = {}
    local citizens = dfhack.units.getCitizens(true)
    
    if not citizens then
        return
    end
    
    for _, unit in ipairs(citizens) do
        if unit and dfhack.units.isActive(unit) and not dfhack.units.isDead(unit) then
            local category, activity = self:categorizeUnit(unit)
            
            -- Group similar activities
            local key = activity or "Unknown"
            
            -- Simplify some common job names for better grouping
            if category == "job" and key then
                if key:find("Haul") then
                    if key:find("Stone") then key = "Hauling (Stone)"
                    elseif key:find("Wood") then key = "Hauling (Wood)"
                    elseif key:find("Food") then key = "Hauling (Food)"
                    elseif key:find("Item") then key = "Hauling (Items)"
                    elseif key:find("Body") then key = "Hauling (Bodies)"
                    elseif key:find("Refuse") then key = "Hauling (Refuse)"
                    elseif key:find("Furniture") then key = "Hauling (Furniture)"
                    elseif key:find("Animal") then key = "Hauling (Animals)"
                    else key = "Hauling (Other)"
                    end
                elseif key:find("Eat") or key:find("Drink") then
                    key = "Eating/Drinking"
                elseif key:find("Sleep") or key:find("Rest") then
                    key = "Resting"
                elseif key:find("Construction") or key:find("Build") then
                    key = "Construction"
                elseif key:find("Mine") or key:find("Dig") then
                    key = "Mining/Digging"
                elseif key:find("Craft") or key:find("Make") then
                    key = "Crafting"
                end
            end
            
            if not self.job_stats[key] then
                self.job_stats[key] = {count = 0, category = category}
                self.job_units[key] = {}
            end
            self.job_stats[key].count = self.job_stats[key].count + 1
            table.insert(self.job_units[key], unit.id)
        end
    end
end

function DwarfJobMonitor:calculateOptimalFrameSize()
    -- Base size for header and controls
    local base_width = 50
    local base_height = 8 -- Header (3) + padding (2) + minimum list space (3)
    
    -- Calculate content-based dimensions
    local num_jobs = 0
    local max_activity_length = 0
    
    for activity, data in pairs(self.job_stats) do
        num_jobs = num_jobs + 1
        max_activity_length = math.max(max_activity_length, string.len(activity))
    end
    
    -- If no jobs, use minimum size
    if num_jobs == 0 then
        return {w = math.max(base_width, 35), h = math.max(base_height, 15)}
    end
    
    -- Calculate optimal width based on content
    -- Account for: count (3) + "x " (2) + activity name + padding (2) + " [" (2) + progress bar (15) + "]" (1)
    local content_width = 3 + 2 + max_activity_length + 2 + 2 + 15 + 1
    local optimal_width = math.max(base_width, content_width + 4) -- +4 for window borders/padding
    
    -- Calculate optimal height based on number of jobs
    local list_height = math.min(num_jobs, 20) -- Cap at 20 visible items to prevent huge windows
    local optimal_height = base_height + list_height
    
    -- Ensure minimum size constraints
    optimal_width = math.max(optimal_width, 35)
    optimal_height = math.max(optimal_height, 15)
      -- Cap maximum size to reasonable screen limits
    optimal_width = math.min(optimal_width, 80)
    optimal_height = math.min(optimal_height, 40)
    
    return {w = optimal_width, h = optimal_height}
end

function DwarfJobMonitor:getTotalDwarvesCount()
    local citizens = dfhack.units.getCitizens(true)
    return #citizens
end

function DwarfJobMonitor:getIdleDwarvesCount()
    local count = 0
    for _, unit in ipairs(dfhack.units.getCitizens(true)) do
        if not unit.job.current_job and #unit.social_activities == 0 and 
           #unit.individual_drills == 0 and unit.military.squad_id == -1 then
            count = count + 1
        end
    end
    return count
end

function DwarfJobMonitor:getWorkingDwarvesCount()
    return self:getTotalDwarvesCount() - self:getIdleDwarvesCount()
end

function DwarfJobMonitor:getJobName(job)
    if not job then return nil end
    return dfhack.job.getName(job) or df.job_type[job.job_type] or "Unknown Job"
end

function DwarfJobMonitor:getActivityName(activity)
    if not activity then return nil end
    if not activity.type then return "Unknown Activity" end
    
    local activityNames = {
        [df.activity_entry_type.TrainingSession] = "Training Session",
        [df.activity_entry_type.IndividualSkillDrill] = "Individual Skill Drill",
        [df.activity_entry_type.FillServiceOrder] = "Fill Service Order",
        [df.activity_entry_type.StoreObject] = "Store Object",
        [df.activity_entry_type.Conversation] = "Conversation",
        [df.activity_entry_type.Worship] = "Worship",
        [df.activity_entry_type.Prayer] = "Prayer",
        [df.activity_entry_type.SocializeWParty] = "Socialize at Party",
        [df.activity_entry_type.Guardduty] = "Guard Duty",
        [df.activity_entry_type.Research] = "Research",
        [df.activity_entry_type.PonderTopic] = "Ponder Topic",
        [df.activity_entry_type.DiscussIdea] = "Discuss Idea",
        [df.activity_entry_type.ReadBook] = "Read Book",
        [df.activity_entry_type.WriteBook] = "Write Book",
        -- Additional activity types commonly seen
        [df.activity_entry_type.CombatTraining] = "Combat Training",
        [df.activity_entry_type.Socialize] = "Socializing",
        [df.activity_entry_type.Sparring] = "Sparring",
        [df.activity_entry_type.SkillDemonstration] = "Skill Demonstration",
        [df.activity_entry_type.RangedPractice] = "Ranged Practice",
        [df.activity_entry_type.Harassment] = "Harassment",
        [df.activity_entry_type.Encounter] = "Encounter",
        [df.activity_entry_type.Conflict] = "Conflict",
        [df.activity_entry_type.MakeBelieve] = "Make Believe",
        [df.activity_entry_type.Performance] = "Performance",
        [df.activity_entry_type.PlayWithToy] = "Playing with Toy",
        [df.activity_entry_type.Play] = "Playing",
        [df.activity_entry_type.Reunion] = "Reunion",
        [df.activity_entry_type.CopyWrittenContent] = "Copying Written Content",
        [df.activity_entry_type.TeachTopic] = "Teaching Topic",
    }
    
    -- Try to get a friendly name first
    local name = activityNames[activity.type]
    if name then
        return name
    end
    
    -- Fallback to the raw enum name, but clean it up
    local raw_name = df.activity_entry_type[activity.type]
    if raw_name then
        -- Convert CamelCase to readable format (e.g., "SomeActivity" -> "Some Activity")
        local cleaned = raw_name:gsub("([a-z])([A-Z])", "%1 %2")
        return cleaned
    end
    
    -- Final fallback
    return "Unknown Activity"
end

function DwarfJobMonitor:categorizeUnit(unit)
    -- Safely check if unit has a current job
    if unit.job and unit.job.current_job then
        local job_name = self:getJobName(unit.job.current_job)
        return "job", job_name or "Unknown Job"
    end
      -- Check social activities
    if unit.social_activities and #unit.social_activities > 0 then
        local activity = dfhack.units.getMainSocialActivity(unit)
        if activity then
            local success, activity_name = pcall(function() 
                return self:getActivityName(activity) 
            end)
            if success and activity_name then
                return "activity", activity_name
            else
                -- If getActivityName fails, try to get some basic info
                local basic_name = "Social Activity"
                if activity.type then
                    local raw_name = df.activity_entry_type[activity.type]
                    if raw_name then
                        basic_name = raw_name:gsub("([a-z])([A-Z])", "%1 %2")
                    end
                end
                return "activity", basic_name
            end
        else
            return "activity", "Social Activity"
        end
    end
    
    -- Check individual drills
    if unit.individual_drills and #unit.individual_drills > 0 then
        return "training", "Individual Training"
    end
    
    -- Check military squad
    if unit.military and unit.military.squad_id ~= -1 then
        local squad = df.squad.find(unit.military.squad_id)
        if squad and (squad.orders and #squad.orders > 0 or squad.activity ~= -1) then
            return "military", "Military Duty"
        end
    end
    
    -- Check if sleeping
    if unit.counters and unit.counters.unconscious > 0 then
        return "rest", "Sleeping"
    end
    
    -- Check if resting
    if unit.counters and (unit.counters.winded > 0 or unit.counters.stunned > 0) then
        return "rest", "Resting"
    end
    
    -- Default to idle
    return "idle", "Idle"
end

function DwarfJobMonitor:refreshData()
    self.job_stats = {}
    self.job_units = {} -- Track units for each job type
    local citizens = dfhack.units.getCitizens(true)
    
    if not citizens then
        return
    end
      for _, unit in ipairs(citizens) do
        if unit and dfhack.units.isActive(unit) and not dfhack.units.isDead(unit) then
            local category, activity = self:categorizeUnit(unit)
            
            -- Group similar activities
            local key = activity or "Unknown"
            
            -- Simplify some common job names for better grouping
            if category == "job" and key then
                if key:find("Haul") then
                    if key:find("Stone") then key = "Hauling (Stone)"
                    elseif key:find("Wood") then key = "Hauling (Wood)"
                    elseif key:find("Food") then key = "Hauling (Food)"
                    elseif key:find("Item") then key = "Hauling (Items)"
                    elseif key:find("Body") then key = "Hauling (Bodies)"
                    elseif key:find("Refuse") then key = "Hauling (Refuse)"
                    elseif key:find("Furniture") then key = "Hauling (Furniture)"
                    elseif key:find("Animal") then key = "Hauling (Animals)"
                    else key = "Hauling (Other)"
                    end
                elseif key:find("Eat") or key:find("Drink") then
                    key = "Eating/Drinking"
                elseif key:find("Sleep") or key:find("Rest") then
                    key = "Resting"
                elseif key:find("Construction") or key:find("Build") then
                    key = "Construction"
                elseif key:find("Mine") or key:find("Dig") then
                    key = "Mining/Digging"
                elseif key:find("Craft") or key:find("Make") then
                    key = "Crafting"
                end
            end
            
            if not self.job_stats[key] then
                self.job_stats[key] = {count = 0, category = category}
                self.job_units[key] = {}
            end
            self.job_stats[key].count = self.job_stats[key].count + 1
            table.insert(self.job_units[key], unit.id)
        end
    end
    
    self:updateJobList()
    self.last_refresh = dfhack.getTickCount()
end

function DwarfJobMonitor:adjustFrameSize()
    -- Calculate new optimal size based on current content
    local optimal_size = self:calculateOptimalFrameSize()
    local current_frame = self.subviews.main_window.frame
    
    -- Only resize if the content requires a significantly different size
    local width_diff = math.abs(optimal_size.w - current_frame.w)
    local height_diff = math.abs(optimal_size.h - current_frame.h)
    
    -- Resize if difference is more than 5 units to avoid constant small adjustments
    if width_diff > 5 or height_diff > 5 then
        self.subviews.main_window.frame.w = optimal_size.w
        self.subviews.main_window.frame.h = optimal_size.h
        self.subviews.main_window:updateLayout()
    end
end

function DwarfJobMonitor:getCategoryColor(category)
    local colors = {
        job = COLOR_LIGHTGREEN,
        activity = COLOR_LIGHTBLUE,
        training = COLOR_YELLOW,
        military = COLOR_LIGHTRED,
        rest = COLOR_LIGHTMAGENTA,
        idle = COLOR_WHITE
    }
    return colors[category] or COLOR_GREY
end

function DwarfJobMonitor:updateJobList()
    local choices = {}
    
    -- Convert stats to sorted list
    local sorted_stats = {}
    for activity, data in pairs(self.job_stats) do
        table.insert(sorted_stats, {activity = activity, count = data.count, category = data.category})
    end
    
    -- Sort by count (descending)
    table.sort(sorted_stats, function(a, b) return a.count > b.count end)
    
    -- Calculate the maximum activity name length for alignment
    local max_name_length = 0
    for _, stat in ipairs(sorted_stats) do
        max_name_length = math.max(max_name_length, string.len(stat.activity))
    end
    
    for _, stat in ipairs(sorted_stats) do
        local color = self:getCategoryColor(stat.category)
        local count_text = string.format("%2d", stat.count)
        local activity_text = stat.activity
        
        -- Pad the activity name to align the progress bars
        local padding = string.rep(" ", max_name_length - string.len(activity_text))
        
        -- Create progress bar representation
        local max_count = sorted_stats[1] and sorted_stats[1].count or 1
        local bar_width = math.min(15, math.floor((stat.count / max_count) * 15))
        local bar = string.rep("#", bar_width) .. string.rep("-", 15 - bar_width)
        
        table.insert(choices, {
            text = {
                {text = count_text, pen = color},
                {text = "x ", pen = COLOR_GREY},
                {text = activity_text, pen = color},
                {text = padding, pen = COLOR_GREY},
                {text = " [", pen = COLOR_GREY},
                {text = bar, pen = color},
                {text = "]", pen = COLOR_GREY}
            },
            activity = stat.activity,
            count = stat.count,
            category = stat.category
        })
    end
    
    if #choices == 0 then
        table.insert(choices, {
            text = {{text = "No active dwarfs found", pen = COLOR_GREY}},
            activity = "none",
            count = 0,
            category = "none"
        })
    end
    
    self.subviews.job_list:setChoices(choices)
    
    -- Set up click handler for cycling through units
    self.subviews.job_list.on_submit = function(idx, choice)
        if choice and choice.activity and choice.activity ~= "none" then
            self:cycleToNextUnit(choice.activity)
        end
    end
end

function DwarfJobMonitor:cycleToNextUnit(job_name)
    local units = self.job_units[job_name]
    if not units or #units == 0 then
        -- No units left in this category, do nothing
        return
    end
    
    -- Reset refresh timer when user clicks for cycling
    self.last_refresh = dfhack.getTickCount()
    
    -- Initialize cycle index if it doesn't exist
    if not self.job_cycle_index[job_name] then
        self.job_cycle_index[job_name] = 1
    else
        -- Cycle to next unit
        self.job_cycle_index[job_name] = self.job_cycle_index[job_name] + 1
        if self.job_cycle_index[job_name] > #units then
            self.job_cycle_index[job_name] = 1
        end
    end
    
    -- Try to find a valid unit, skipping over non-existent ones
    local attempts = 0
    local max_attempts = #units
    local unit = nil
    
    while attempts < max_attempts do
        local unit_id = units[self.job_cycle_index[job_name]]
        unit = df.unit.find(unit_id)
        
        if unit then
            -- Found a valid unit, break out of the loop
            break
        else
            -- Current unit doesn't exist, skip to the next one
            self.job_cycle_index[job_name] = self.job_cycle_index[job_name] + 1
            if self.job_cycle_index[job_name] > #units then
                self.job_cycle_index[job_name] = 1
            end
            attempts = attempts + 1
        end
    end
    
    if not unit then
        -- No valid units found after checking all, refresh data to clean up
        self:refreshData()
        return
    end
    
    -- Zoom to unit
    dfhack.gui.revealInDwarfmodeMap(
        xyz2pos(dfhack.units.getPosition(unit)), true, true)
    
    -- Follow unit if in fortress mode
    if dfhack.world.isFortressMode() then
        df.global.plotinfo.follow_item = -1
        df.global.plotinfo.follow_unit = unit.id
        pcall(function()
            -- if spectate is available, add the unit to the follow history
            local spectate = require('plugins.spectate')
            spectate.spectate_addToHistory(unit.id)
        end)
    end
end

function DwarfJobMonitor:togglePin()
    self.is_pinned = not self.is_pinned
    -- Update the window title to show pin status
    local title = self.is_pinned and 'Dwarf Job Monitor [PINNED]' or 'Dwarf Job Monitor'
    self.subviews.main_window.frame_title = title
end

function DwarfJobMonitor:onInput(keys)
    if keys.LEAVESCREEN or keys._MOUSE_R then
        -- Only dismiss if not pinned
        if not self.is_pinned then
            self:dismiss()
            return true
        end
    end
    
    -- Auto-refresh every 5 seconds
    local now = dfhack.getTickCount()
    if now - self.last_refresh > 500 then -- ~5 seconds at 100 ticks/second
        self:refreshData()
    end
    
    return DwarfJobMonitor.super.onInput(self, keys)
end

function DwarfJobMonitor:onRenderBody(dc)
    DwarfJobMonitor.super.onRenderBody(self, dc)
end

-- Main script entry point
if dfhack_flags.module then
    return DwarfJobMonitor
end

-- Global variable to store the screen instance
if not dfhack.isWorldLoaded() then
    qerror("This script requires a loaded world")
end

-- Check if there's already a running instance
if screen and not screen._native then
    screen = nil
end

if screen then
    -- Try to check if screen is still valid
    local success, is_dismissed = pcall(function() return screen:isDismissed() end)
    if not success or is_dismissed then
        screen = nil
    else
        -- If already showing and valid, bring to front
        screen:raise()
        return
    end
end

-- Create and show new instance
screen = DwarfJobMonitor{}:show()