a jobs module for creating/assigning jobs

jobs.lua

--a library for creation and modification of jobs
--TODO: workshop jobs/reactions and building specific jobs in general
local _ENV = mkmodule('hack.scripts.jobs')

local buildings = require('dfhack.buildings')
local tile_attrs = df.tiletype.attrs

--[=[ helper functions ]=]
local function findRef(vector,ref_type)
    for k,v in ipairs(vector) do
        if v:getType()==ref_type then
            return v,k
        end
    end
end
local function makeset(tbl)
    local ret={}
    for k,v in pairs(tbl) do
        ret[v]=true
    end
    return ret
end
--[=[ low level stuff ]=]
local function getLastJobLink()
    local st=df.global.world.job_list
    while st.next~=nil do
        st=st.next
    end
    return st
end
local function addNewJob(job)
    local lastLink=getLastJobLink()
    local newLink=df.job_list_link:new()
    newLink.prev=lastLink
    lastLink.next=newLink
    newLink.item=job
    job.list_link=newLink
    
    job.id=df.global.job_next_id
    df.global.job_next_id=df.global.job_next_id+1
end


local function removeJobFromUnit(unit)
    local job=unit.job.current_job
    unit.job.current_job=nil
    return job
end
local function removeUnitFromJob(job)
    local ref,id=findRef(job.general_refs,df.general_ref_type.UNIT_WORKER)
    if ref then
        job.general_refs:erase(id)
        local unit=df.unit.find(ref.unit_id)
        ref:delete()
        return unit
    end
end
local function performAssigns(args)
    if args.assigns then
        for _,assignFunction in ipairs(args.assigns) do
            assignFunction(args)
        end
    end
    if args.unit then
        assignJobToUnit(args.job,args.unit,args.from_pos)
    end
    local force=false
    if args.items then
        force=args.items.force
    end
    if #args.job.job_items>0 or force then
        setJobItems(args)
    end
end
local function makeJobImpl(args,isNewJob)
    
    local job=args.job
    job.job_type=args.job_type
    job.completion_timer=-1
    
    if args.pos then
        job.pos:assign(args.pos)
    end
    
    performAssigns(args)
    --very last thing:
    if isNewJob then
        addNewJob(job) --if we fail after this, we need a lot more cleanup!
    end
        
    return job
end
--[=[ some predicates for checks (well they are more error-predicate-thingies or sth... ]=]
local function predWieldsItemWithSkill(item_skill)
    local pred=function(args)
        if args.unit==nil then
            return --no error, will not assign unit, so df will find unit with pick
        end
        local inv=args.unit.inventory
        for k,v in pairs(inv) do
            if v.mode==1 and v.item:getMeleeSkill()==item_skill and args.unit.body.weapon_bp==v.body_part_id then
                return
            end
        end
        qerror(string.format("This job needs item with skill:%s equiped on your correct body-part",df.job_skill[item_skill]))
    end
    return pred
end
local function isWall(args)
    local tt=dfhack.maps.getTileType(args.pos)
    if tile_attrs[tt].shape~=df.tiletype_shape.WALL then
        qerror("Job needs a wall on target position")
    end
end
local function isHardMaterial(args)
    local tt=dfhack.maps.getTileType(args.pos)
    local mat=tile_attrs[tt].material
    local hard_materials=makeset{df.tiletype_material.STONE,df.tiletype_material.FEATURE,
        df.tiletype_material.LAVA_STONE,df.tiletype_material.MINERAL,df.tiletype_material.FROZEN_LIQUID,}
    if not hard_materials[mat] then
        qerror("Job needs a hard material (e.g. stone, ice...) at target position")
    end
end
local function isStairs(args)
    local tt=dfhack.maps.getTileType(args.pos)
    local shape=tile_attrs[tt].shape
    if not(shape==df.tiletype_shape.STAIR_UP or shape==df.tiletype_shape.STAIR_DOWN or 
        shape==df.tiletype_shape.STAIR_UPDOWN or shape==df.tiletype_shape.RAMP) then
        qerror("Job needs stairs or ramps at target position")
    end
end
local function isFloor(args)
    local tt=dfhack.maps.getTileType(args.pos)
    local shape=tile_attrs[tt].shape
    if not(shape==df.tiletype_shape.FLOOR or shape==df.tiletype_shape.BOULDER or shape==df.tiletype_shape.PEBBLES) then
        qerror("Job needs floor at target position")
    end
end
local function isTree(args)
    local tt=dfhack.maps.getTileType(args.pos)
    if tile_attrs[tt].shape~=df.tiletype_shape.TREE then
        qerror "Job needs tree at target position"
    end
end
local function isPlant(args)
    local tt=dfhack.maps.getTileType(args.pos)
    if tile_attrs[tt].shape~=df.tiletype_shape.SHRUB then
        qerror "Job needs plant at target position"
    end
end
local function isUnit(args)
    if args.trg_unit then
        return 
    end
    local pos=args.pos
    for k,v in pairs(df.global.world.units.active) do
        if v.pos.x==pos.x and v.pos.y==pos.y and v.pos.z==pos.z then
            return
        end
    end
    qerror "Unit must be present at job target position"
end
local function isSameSquare(args) --todo check this out... might be a problem
    local pos1=args.pos
    local pos2=copyall(args.unit)
    if pos1.x~=pos2.x or pos1.y~=pos2.y or pos1.z==pos2.z then
       qerror "Job can only be performed on same square a the unit square"
    end
end
local function isNotConstruct(args)
    local tt=dfhack.maps.getTileType(args.pos)
    if tile_attrs[tt].material==df.tiletype_material.CONSTRUCTION or dfhack.buildings.findAtTile(args.pos)~=nil then
        qerror "Job needs no constructions or buildings at target pos"
    end
end
local function isNonConstructedBuilding(args)
    local bld=args.building or dfhack.buildings.findAtTile(args.pos)
    if not bld or bld:getMaxBuildStage()==bld:getBuildStage()  then
        qerror "Job needs non-finished building"
    end
end
local function isBuilding(args)
    local bld=args.building or dfhack.buildings.findAtTile(args.pos)
    if bld==nil then
        qerror "Job needs a building in target position"
    end
end
local function isConstruct(args)
    local tt=dfhack.maps.getTileType(args.pos)
    if tile_attrs[tt].material~=df.tiletype_material.CONSTRUCTION then
        qerror "Job needs construction at target position"
    end
end
--[=[ end of predicate things ]=]
local function setEmptyNear(args)
    if args.from_pos then --user already set this prob...
        return 
    end
    local arrDx={-1,0,0,1}
    local arrDy={0,-1,1,0}
    local sAttr=df.tiletype_shape.attrs
    local pos
    for i,dx in ipairs(arrDx) do
        local dy=arrDy[i]
        pos={x=args.pos.x+dx,y=args.pos.y+dy,z=args.pos.z}
        local tt=dfhack.maps.getTileType(pos)
        if tt then
            --print(tt)
            --printall(tile_attrs[tt])
            --printall(sAttr[tile_attrs[tt].shape])
            if sAttr[tile_attrs[tt].shape].walkable then
                args.from_pos=pos
                return 
            end
        end
    end
    print("Warning empty tile near job not found")
    --emit warning?
end
local function setBuildingRef(args)
    local bld=args.building or dfhack.buildings.findAtTile(args.pos)
    args.job.general_refs:insert("#",{new=df.general_ref_building_holderst,building_id=bld.id})
    bld.jobs:insert("#",args.job)
    return true
end
local function setBuildingFilters(args)
    if args.job.job_items>0 then --already set
        return
    end
    for k,v in pairs(buildings.getFiltersByType(args, --TODO this could clash
        args.building:getType(),args.building:getSubtype(),args.building:getCustomType())) do
        args.job.job_items:insert("#",v)
    end
end
local function addItem(tbl,item,recurse,skip_add)
    if not skip_add then
        table.insert(tbl,item)
    end
    if recurse then
        local subitems=dfhack.items.getContainedItems(item)
        if subitems~=nil then
            for k,v in pairs(subitems) do
                AddItem(tbl,v,recurse)
            end
        end
    end
end
local function enumItems(args)
    local ret=args.table or {}
    if args.all then
        for k,v in pairs(df.global.world.items.all) do
            if v.flags.on_ground then
                addItem(ret,v,args.deep)
            end
        end
    elseif args.pos~=nil then
        for k,v in pairs(df.global.world.items.all) do
            if v.pos.x==args.pos.x and v.pos.y==args.pos.y and v.pos.z==args.pos.z and v.flags.on_ground then
                addItem(ret,v,args.deep)
            end
        end
    end
    if args.unit~=nil then
        for k,v in pairs(args.unit.inventory) do
            if args.inv[v.mode] then
                addItem(ret,v.item,args.deep)
            elseif args.deep then
                addItem(ret,v.item,args.deep,true)
            end
        end
    end
    return ret
end
-- HUGE FUNCTION WARNING..., can't use the built-in because it fails at many cases
function isSuitableItem(job_item,item)
    --todo butcher test
    if job_item.item_type~=-1 then
        if item:getType()~= job_item.item_type then
            return false, "type"
        elseif job_item.item_subtype~=-1 then
            if item:getSubtype()~=job_item.item_subtype then
                return false,"subtype"
            end
        end
    end
    
    if job_item.mat_type~=-1 then
        if item:getActualMaterial()~= job_item.mat_type then --unless we would want to make hist-fig specific reactions
            return false, "material"
        elseif job_item.mat_index~=-1 then
            if item:getActualMaterialIndex()~=job_item.mat_index then
                return false,"material index"
            end
        end
    end
    if job_item.flags1.sand_bearing and not item:isSandBearing() then
        return false,"not sand bearing"
    end
    if job_item.flags1.butcherable and not (item:getType()== df.item_type.CORPSE or item:getType()==df.item_type.CORPSEPIECE) then
        return false,"not butcherable"
    end
    local matinfo=dfhack.matinfo.decode(item)
    --print(matinfo:getCraftClass())
    --print("Matching ",item," vs ",job_item)
    if not matinfo:matches(job_item) then
        return false,"matinfo"
    end
    -- some bonus checks:
    if job_item.flags2.building_material and not item:isBuildMat() then
        return false,"non-build mat"
    end
    -- *****************
    --print("--Matched")
    --reagen_index?? reaction_id??
    if job_item.metal_ore~=-1 and not item:isMetalOre(job_item.metal_ore) then
        return false,"metal ore"
    end
    if job_item.min_dimension~=-1 then
    end
    if #job_item.contains~=0 then
    end
    if job_item.has_tool_use~=-1 then
        if not item:hasToolUse(job_item.has_tool_use) then
            return false,"tool use"
        end
    end
    if job_item.has_material_reaction_product~="" then
        local ok=false
        for k,v in pairs(matinfo.material.reaction_product.id) do
            if v.value==job_item.has_material_reaction_product then
                ok=true
                break
            end
        end
        if not ok then
            return false, "no material reaction product"
        end
    end
    if job_item.reaction_class~="" then
        local ok=false
        for k,v in pairs(matinfo.material.reaction_class) do
            if v.value==job_item.reaction_class then
                ok=true
                break
            end
        end
        if not ok then
            return false, "no material reaction class"
        end
    end
    return true
end
-- lists all not at the job position items
function getItemsUncollected(job)
    local ret={}
    for id,jitem in pairs(job.items) do
        local x,y,z=dfhack.items.getPosition(jitem.item)
        if x~=job.pos.x or y~=job.pos.y or z~=job.pos.z then
            table.insert(ret,jitem)
        end
    end
    return ret
end
function setJobItems(args)
    if args.items == nil or args.items.auto then --use df default logic and hope that it would work
        return
    end
    local job=args.job
    local items=args.items
    -- first find items that you want to use for the job, if items are not supplied by user
    if #args.items ==0 then
        local items_pos=args.items.pos 
        if items_pos==true then items_pos=args.from_pos or copyall(args.unit.pos) end
        
        items=EnumItems{pos=items_pos,unit=args.unit,
            inv={[df.unit_inventory_item.T_mode.Hauled]=args.items.unit,[df.unit_inventory_item.T_mode.Worn]=args.items.unit,
                 [df.unit_inventory_item.T_mode.Weapon]=args.items.unit,},deep=args.items.unit_deep or false}
    end
    -- figure out what is needed
    local item_counts={}
    for job_id, trg_job_item in ipairs(job.job_items) do
        item_counts[job_id]=trg_job_item.quantity
    end
    -- assign all the items to all the slots
    local used_item_id={}
    for job_id, trg_job_item in ipairs(job.job_items) do
        for item_number,cur_item in pairs(items) do 
            if not used_item_id[cur_item.id] then
                local item_suitable,msg
                if args.items.force and args.items.force[item_number] then -- skip checks, it's all good >:)
                    item_suitable=true
                else
                    item_suitable,msg=isSuitableItem(trg_job_item,cur_item)
                end
                
                if (item_counts[job_id]>0 and item_suitable) or args.items.force==true then 
                    --cur_item.flags.in_job=true --TODO: this would be nice, but if you cancel job it's not getting cleared and you are stuck
                    job.items:insert("#",{new=true,item=cur_item,role=df.job_item_ref.T_role.Reagent,job_item_idx=job_id})
                    item_counts[job_id]=item_counts[job_id]-cur_item:getTotalDimension()
                    used_item_id[cur_item.id]=true
                    --print(string.format("item added, job_item_id=%d, item %s, quantity left=%d",job_id,tostring(cur_item),item_counts[job_id]))
                end
            end
        end
    end
    -- fail if not all slots are filled
    if not args.items.force then
        for job_id, trg_job_item in ipairs(job.job_items) do
            if item_counts[job_id]>0 then
                qerror "Not enough items for job"
            end
        end
    end
    --check if we need to fetch first
    local uncollected = getItemsUncollected(job)
    if #uncollected == 0 then
        job.flags.working=true
    else
        job.flags.fetching=true
        uncollected[1].is_fetching=1
    end

    return true
end

local function performChecks(args)
    local check_callback=args.check_callback or 
        function(msg)
            qerror(msg)
        end
    if args.checks then
        for _,check in ipairs(args.checks) do
            local status,msg=pcall(check,args)
            if not status then
                check_callback(msg,args)
            end
        end
    end
    for _,check in ipairs(args.gen_checks or {}) do
        local status,msg=pcall(check,args)
        if not status then
            check_callback(msg,args)
        end
    end
end
local function getBuildJob(building)
    for idx,job in pairs(building.jobs) do
        if job.job_type==df.job_type.ConstructBuilding then
            return job
        end
    end
end
local function jobFilter(filter_table)
    return function (args)
        args.job.job_items:insert("#",filter_table)
    end
end
-- generates needed job checks and function that assign stuff
local function setupChecksAndAssigns(args)
    local data={ --todo maybe move this outside? could improve performance
        
        [df.job_type.CarveFortification    ]={"Carve Fortification"   ,{isWall,isHardMaterial},{setEmptyNear}},
        [df.job_type.DetailWall            ]={"Detail Wall"           ,{predWieldsItemWithSkill(df.job_skill.MINING),isWall,isHardMaterial},{setEmptyNear}},
        [df.job_type.DetailFloor           ]={"Detail Floor"          ,{predWieldsItemWithSkill(df.job_skill.MINING),isFloor,isHardMaterial,isSameSquare}},
        [df.job_type.CarveTrack            ]={"Carve Track"           ,{predWieldsItemWithSkill(df.job_skill.MINING),isFloor,isHardMaterial},{setCarveDir,setEmptyNear}},--TODO
        [df.job_type.Dig                   ]={"Dig"                   ,{predWieldsItemWithSkill(df.job_skill.MINING),isWall},{setEmptyNear}},
        [df.job_type.FellTree              ]={"Fell Tree"             ,{predWieldsItemWithSkill(df.job_skill.AXE),isTree},{setEmptyNear}},
        [df.job_type.RemoveConstruction    ]={"Remove Construction"   ,{isConstruct},{setEmptyNear}},
        
        [df.job_type.DestroyBuilding       ]={"Remove Building"       ,{isBuilding}},
        [df.job_type.CarveUpwardStaircase  ]={"Carve Upward Staircase" ,{predWieldsItemWithSkill(df.job_skill.MINING),isWall}},
        [df.job_type.CarveDownwardStaircase]={"Carve Downward Staircase",{predWieldsItemWithSkill(df.job_skill.MINING)}},
        [df.job_type.CarveUpDownStaircase  ]={"Carve Up/Down Staircase" ,{predWieldsItemWithSkill(df.job_skill.MINING)}},
        [df.job_type.CarveRamp             ]={"Carve Ramp"            ,{predWieldsItemWithSkill(df.job_skill.MINING),isWall}},
        [df.job_type.DigChannel            ]={"Dig Channel"           ,{predWieldsItemWithSkill(df.job_skill.MINING)}},
        [df.job_type.Fish                  ]={"Fish"                  ,{--[[isWater]]}}, --todo figure this out...
        [df.job_type.TameAnimal            ]={"Tame Animal"           ,{isUnit},{setCreatureRef}}, --add item?
        [df.job_type.GatherPlants          ]={"Gather Plants"         ,{isPlant,isSameSquare}},
        [df.job_type.RemoveStairs          ]={"Remove Stairs"         ,{isStairs,isNotConstruct}},
        [df.job_type.Clean                 ]={"Clean"                 ,{}},
        [df.job_type.CollectWebs           ]={"Gather Webs"           ,{--[[HasWeb]]},{setWebRef}},
        [df.job_type.ConstructBuilding     ]={"Build"                 ,{isNonConstructedBuilding},{setBuildingRef,setBuildingFilters}}, -- a special case, construct planned building
        --{"Diagnose Patient"    ,df.job_type.DiagnosePatient,{isUnit},{setPatientRef}},
        --{"Surgery"             ,df.job_type.Surgery,{isUnit},{setPatientRef}},
        --{"HandleLargeCreature" ,df.job_type.HandleLargeCreature,{isUnit},{SetCreatureRef}},
        --{"Build"               ,AssignJobToBuild,{NoConstructedBuilding}},
        --{"BuildLast"           ,BuildLast,{NoConstructedBuilding}}, <- convenience job?
        --[[
            a lot of reactions, etc from other stuff like:
                building <- needs items
                in building reactions/jobs <- needs items
                custom reactions <- needs reaction name too
        ]]
    }
    
    local job_type=args.job_type
    if data[job_type]==nil then
        qerror("Sorry unsupported job type")
    end
    args.name=data[job_type][1]
    args.gen_checks=data[job_type][2]
    args.assigns=data[job_type][3]
    --special cases:
    if job_type==df.job_type.ConstructBuilding then
        --continue existing job plz
        local building=args.building or dfhack.buildings.findAtTile(args.pos)
        args.job=getBuildJob(building) --TODO remove all these find at tile, they are slow
        args.building=building
    end
end
--[=[ public functions ]=]
-- lists possible jobs in the building.
function enumBuildingJobs(building)
    local data={
        [df.building_type.Farm]={
        {[df.job_type.PlantSeeds] = {"Plant Seeds",{},{jobFilter{new=true,df.item_type.SEEDS}}},
        --{[df.job_type.GatherPlants] = {"Plant Seeds",{},}
        }
        
    }
end
-- removes unit from job and job from unit
function unassignJob(job_or_unit)
    if df.job:is_instance(job_or_unit) then
        local unit=removeUnitFromJob(job_or_unit)
        removeJobFromUnit(unit)
    else
        local job=removeJobFromUnit(job_or_unit)
        removeUnitFromJob(job)
    end
end
--assigns job to unit, also reasigns if already has a job
function assignJobToUnit(job,unit,from_pos)
    if unit.job.current_job ~= nil then
        unassignJob(unit)
    end
    job.general_refs:insert("#",{new=df.general_ref_unit_workerst,unit_id=unit.id})
    unit.job.current_job=job
    unit.path.dest:assign(from_pos or unit.pos)
end
--[[
    makes a fully functional job. Arguments:
        job_type
        building - a building for work to be performed, some jobs don't need this (e.g. collect plants, or build construction)
        pos - job target position, could be autofilled if not set in some cases
        from_pos - a location for unit to be standing when performing job
        unit -  a job performing unit, if not set will leave for df to assign
        trg_unit - a target unit for some jobs (tries to fill it out from pos)
        no_reassign - fails if unit already has a job
        items - a table to set up items for job (can be not set, then it used default df logic)
            auto - same as items=nil
            unit - use items from unit
            unit_deep - recurse into units containers
            pos - at this position or if ==true then at from_pos
            force - ignore checks, can be set as items.force=true or as e.g. items.force={[1]=true,[3]=true} for 1st and 3rd forced item
                    also ignores missing items (i.e. 2 boulders for reaction needing 3) !! MAY CRASH THE GAME !!
            
        checks - an optional checks table, to check if everything is valid before anything
        no_checks - perform no checks and do it! (usefull for cheating)
        only_checks - do not create job, just check stuff
        check_callback - called with message and arguments table, if set this does not qerror out (useful for gui pre-checks, status showing)
]]
function makeJob(args)
    if args.job_type==nil then
        qerror("No job_type specified")
    end
    if args.no_reassign then
        if args.unit.job.current_job~=nil then
            qerror("Unit already has a job")
        end
    end
    setupChecksAndAssigns(args)
    --[=[ checking stuff out ]=]
    if not args.no_checks then
        performChecks(args)
    end
    if args.only_checks then
        return
    end
    --[=[ real creation ]=]
    if args.job==nil then
        local newJob=df.job:new()
        args.job=newJob
        return dfhack.call_with_finalizer(1,false,df.delete,newJob,makeJobImpl,args,true)
    else
        makeJobImpl(args,false)
    end
end

-- some convienience functions
-- make build, where it designates, and assigns job
-- needs to know either items or filters or material (see dfhack.buildings)
-- special arguments: skip_make, to skip the job logic, but then you could just use dfhack.buildings
function planBuilding(args)
    local bld,err=buildings.constructBuilding(args)
    if bld==nil then
        qerror(err)
    end
    args.building=bld
    -- we actually don't create the job, just assign unit
    for idx,job in pairs(bld.jobs) do
        if job.job_type==df.job_type.ConstructBuilding then
            args.job=job
            break
        end
    end
    
    if args.job==nil then
        qerror("No job created on building")
    end
    args.job_type=args.job.job_type
    if not args.skip_make then
        makeJob(args)-- need to assign all the items and other stuff
    end
    return bld
end
return _ENV