local _ENV = mkmodule('plugins.buildingplan.mechanisms')

local itemselection = require('plugins.buildingplan.itemselection')

local overlay = require('plugins.overlay')

local utils = require('utils')

local widgets = require('gui.widgets')

local view_sheets = df.global.game.main_interface.view_sheets

--------------------------------

-- MechanismOverlay

--

MechanismOverlay = defclass(MechanismOverlay, overlay.OverlayWidget)

MechanismOverlay.ATTRS{

desc='Adds mechanism selection capabilities to the link lever/pressure plate screens.',

default_pos={x=5,y=5},

default_enabled=true,

viewscreens='dwarfmode/ViewSheets/BUILDING/LinkingLever',

frame={w=57, h=13},

}

local function get_label()

local bld = dfhack.gui.getSelectedBuilding(true)

if bld and bld:getType() == df.building_type.Trap and

bld:getSubtype() == df.trap_type.PressurePlate

then

return 'Plate:'

end

return 'Lever:'

end

function MechanismOverlay:init()

self:addviews{

widgets.BannerPanel{

frame={t=5, l=1, r=1, h=1},

subviews={

widgets.Label{

frame={t=0, l=1},

text='Mechanism safety:'

},

widgets.CycleHotkeyLabel{

view_id='safety_lever',

frame={t=0, l=20, w=15},

key='CUSTOM_G',

label=get_label,

options={

{label='Any', value=0},

{label='Magma', value=2, pen=COLOR_RED},

{label='Fire', value=1, pen=COLOR_LIGHTRED},

},

initial_option=0,

on_change=self:callback('choose_mechanism', 'lever', true),

},

widgets.CycleHotkeyLabel{

view_id='safety_target',

frame={t=0, l=38, w=16},

key='CUSTOM_SHIFT_G',

label='Target:',

options={

{label='Any', value=0},

{label='Magma', value=2, pen=COLOR_RED},

{label='Fire', value=1, pen=COLOR_LIGHTRED},

},

initial_option=0,

on_change=self:callback('choose_mechanism', 'target', true),

},

}

},

widgets.HotkeyLabel{

frame={t=7, l=8, w=49, h=2},

key='CUSTOM_M',

label=function()

return itemselection.get_item_description(view_sheets.linking_lever_mech_lever_id,

nil,

self.subviews.safety_lever:getOptionLabel())

end,

auto_height=false,

enabled=function() return view_sheets.linking_lever_mech_lever_id ~= -1 end,

on_activate=self:callback('choose_mechanism', 'lever', false),

},

widgets.HotkeyLabel{

frame={t=10, l=8, w=49, h=2},

key='CUSTOM_SHIFT_M',

label=function()

return itemselection.get_item_description(view_sheets.linking_lever_mech_target_id,

nil,

self.subviews.safety_target:getOptionLabel())

end,

auto_height=false,

enabled=function() return view_sheets.linking_lever_mech_target_id ~= -1 end,

on_activate=self:callback('choose_mechanism', 'target', false),

},

}

end

local item_selection_dlg

local function reset_dlg()

if item_selection_dlg then

if item_selection_dlg:isActive() then

item_selection_dlg:dismiss()

end

item_selection_dlg = nil

end

end

local function is_reachable(bld_pos, item_id)

if not bld_pos then return false end

local item = df.item.find(item_id)

if not item then return false end

return dfhack.maps.canWalkBetween(xyz2pos(dfhack.items.getPosition(item)), bld_pos)

end

local function sort_by_safety(item_ids, safety)

if safety == 2 then return item_ids end

local safety_ids = {[safety]=utils.invert(item_ids)}

for sort_safety=safety+1,2 do

local safety_item_ids = require('plugins.buildingplan').getAvailableItemsByHeat(

df.building_type.Trap, df.trap_type.Lever, -1, 0, sort_safety)

safety_ids[sort_safety] = utils.invert(safety_item_ids)

end

local sorted_ids = {}

for _,item_id in ipairs(item_ids) do

for safety_idx=2,safety,-1 do

if safety_ids[safety_idx][item_id] then

table.insert(ensure_key(sorted_ids, safety_idx), item_id)

break

end

end

end

local ret = {}

for safety_id=safety,2 do

if sorted_ids[safety_id] then

table.move(sorted_ids[safety_id], 1, #sorted_ids[safety_id], #ret+1, ret)

end

end

return ret

end

-- returns mechanisms ordered by heat safety, so the least heat safe mechanisms are used first

-- when heat safety isn't important

local function get_available_items(safety, other_mechanism)

local item_ids = require('plugins.buildingplan').getAvailableItemsByHeat(

df.building_type.Trap, df.trap_type.Lever, -1, 0, safety)

local bld = dfhack.gui.getSelectedBuilding(true)

local bld_pos = bld and xyz2pos(bld.centerx, bld.centery, bld.z)

for idx=#item_ids,1,-1 do

local item_id = item_ids[idx]

if item_id == other_mechanism or not is_reachable(bld_pos, item_id) then

table.remove(item_ids, idx)

end

end

return sort_by_safety(item_ids, safety)

end

function MechanismOverlay:save_id(which, item_id)

local saved_id = ('saved_%s_id'):format(which)

local ui_id = ('linking_lever_mech_%s_id'):format(which)

item_id = item_id or -1

view_sheets[ui_id] = item_id

self[saved_id] = item_id

end

function MechanismOverlay:choose_mechanism(which, autoselect)

local widget_id = 'safety_' .. which

local safety = self.subviews[widget_id]:getOptionValue()

local ui_other_id = ('linking_lever_mech_%s_id'):format(which == 'lever' and 'target' or 'lever')

local available_item_ids = get_available_items(safety, view_sheets[ui_other_id])

if autoselect then

self:save_id(which, available_item_ids[1])

return

end

-- to integrate with ItemSelection's last used sorting

df.global.buildreq.building_type = df.building_type.Trap

local desc = self.subviews[widget_id]:getOptionLabel()

if desc ~= 'Any' then

desc = desc .. ' safe'

end

desc = desc .. ' mechanism'

item_selection_dlg = item_selection_dlg or itemselection.ItemSelectionScreen{

get_available_items_fn=function() return available_item_ids end,

desc=desc,

quantity=1,

autoselect=false,

on_cancel=reset_dlg,

on_submit=function(chosen_ids)

self:save_id(which, chosen_ids[1] or available_item_ids[1])

reset_dlg()

end,

}:show()

end

function MechanismOverlay:onInput(keys)

if MechanismOverlay.super.onInput(self, keys) then

return true

end

if keys._MOUSE_L then

if self:getMousePos() then

-- don't let clicks bleed through the panel

return true

end

-- don't allow the lever to be linked if mechanisms are not set

return view_sheets.linking_lever_mech_lever_id == -1 or

view_sheets.linking_lever_mech_target_id == -1

end

end

function MechanismOverlay:onRenderFrame(dc, rect)

MechanismOverlay.super.onRenderFrame(self, dc, rect)

if self.saved_lever_id ~= view_sheets.linking_lever_mech_lever_id then

self:choose_mechanism('lever', true)

end

if self.saved_target_id ~= view_sheets.linking_lever_mech_target_id then

self:choose_mechanism('target', true)

end

end

return _ENV