-- Changes the beliefs (values) of units.
--@ module = true

--[[ Notes:
> When using these script functions in a script, use modtools/set-need's rebuildNeeds function to rebuild needs to reflect new beliefs.
> When using script functions, use numerical IDs for representing beliefs.
> The random assignment is placeholder for now. I don't understand how they work
]]

local utils = require('utils')

local validArgs = utils.invert({
  "all",
  "belief",
  "set",
  "tier",
  "modify",
  "step",
  "random",
  "default",
  "citizens",
  "unit",
  "help",
  "list",
  "noneed",
  "listunit"
})

rng = rng or dfhack.random.new(nil, 10)

-- These ranges represent the bounds at which the description for a belief changes
local tierRanges = {{min = -50, max = -41, weight = 4}, {min = -40, max = -26, weight = 20}, {min = -25, max = -11, weight = 85}, {min = -10, max = 10, weight = 780}, {min = 11, max = 25, weight = 85}, {min = 26, max = 40, weight = 20}, {min = 41, max = 50, weight = 4}}

---------------------
local function clamp(value, low, high)
  return math.min(high, math.max(low, value))
end

local function weightedRoll(weightedTable)
  local maxWeight = 0
  for index, result in ipairs(weightedTable) do
    maxWeight = maxWeight + result.weight
  end

  local roll = rng:random(maxWeight) + 1
  local currentNum = roll
  local result

  for index, currentResult in ipairs(weightedTable) do
    currentNum = currentNum - currentResult.weight
    if currentNum <= 0 then
      result = currentResult.id
      break
    end
  end

  return result
end

local function randomiseWithinBounds(min, max)
  local range = math.abs(min - max)
  local roll = rng:random(range+1)

  return min + roll
end

---------------------
-- Set unit's belief to a given value
-- Use this as the main method of changing a unit's beliefs!
function setUnitBelief(unit, belief, value, forcePersonal)
  -- Clamp value
  local value = clamp(value, -50, 50)

  -- Remove current belief (if exists)
  removeUnitBelief(unit, belief)

  -- We may want to avoid adding personal beliefs of the same strength as what's already represented by the unit's culture
  -- If the tiers match, forcePersonal must be true for the belief to be added
  local forcePersonal = forcePersonal or true
  local cultureTier = getBeliefTier(getUnitCultureBelief(unit, belief))
  local newBeliefTier = getBeliefTier(value)

  if cultureTier ~= newBeliefTier or forcePersonal then
    addUnitBelief(unit, belief, value)
  end
end

-- Returns true if the unit's belief of the given type is cultural, rather than a personal belief. Otherwise, returns false.
function isCultureBelief(unit, belief)
  return not isPersonalBelief(unit, belief)
end

-- Returns true if the unit's belief of the given type is personal, rather than their culture's default. Otherwise, returns false.
function isPersonalBelief(unit, belief)
  local found = false

  for index, beliefInstance in ipairs(unit.status.current_soul.personality.values) do
    if beliefInstance.type == belief then
      found = true
      break
    end
  end

  return found
end

-- Returns the strength of the unit's belief of the given type
function getUnitBelief(unit, belief)
  local found = false
  local strength = 0

  for index, beliefInstance in ipairs(unit.status.current_soul.personality.values) do
    if beliefInstance.type == belief then
      found = true
      strength = beliefInstance.strength
      break
    end
  end

  if not found then -- Is cultural belief
    strength = getUnitCultureBelief(unit, belief)
  end

  return strength
end

-- Removes the unit's personal belief of given type, if they have it. This means they'll revert to their culture's belief value instead.
-- Returns true if a belief of the type was found and removed, false if not.
function removeUnitBelief(unit, belief)
  local upers = unit.status.current_soul.personality
  local removed = false

  for index, beliefInstance in ipairs(upers.values) do
    if beliefInstance.type == belief then
      upers.values:erase(index)
      removed = true
      break
    end
  end

  return removed
end

-- Add a personal belief to the unit at a given value
-- Note: Doesn't do checking to ensure doesn't already exist!
function addUnitBelief(unit, belief, value)
  local upers = unit.status.current_soul.personality
  local value = clamp(value, -50, 50)

  upers.values:insert("#", {new = true, type = belief, strength = value})
end

-- Gives a list of all the unit's beliefs and their values.
-- Returns a table where the keys are the belief names and its values are its strength
-- If tiers is true, the value is the tier of the belief instead of the trait's strength
local function getUnitBeliefList(unit, tiers)
  local list = {}

  for id, beliefName in ipairs(df.value_type) do
    if beliefName ~= 'NONE' then
      local strength = getUnitBelief(unit, id)

      if tiers then
        list[beliefName] = getBeliefTier(strength)
      else
        list[beliefName] = strength
      end
    end
  end

  return list
end

-- Sets unit's belief to a random value within the given tier.
function setUnitBeliefTier(unit, belief, tier)
  local tier = clamp(tier, 1, 7)
  -- Randomly generate a strength value which falls between within the new tier, and set that as the unit's belief strength
  local strength = randomiseWithinBounds(tierRanges[tier].min, tierRanges[tier].max)
  setUnitBelief(unit, belief, strength)
end

-- Move the unit's belief in tier increments
function stepUnitBelief(unit, belief, modifier)
  local currentTier = getBeliefTier(getUnitBelief(unit, belief))
  local newTier = clamp(currentTier + modifier, 1, 7)

  -- Only do so if new tier is different from the current one
  if currentTier ~= newTier then
    setUnitBeliefTier(unit, belief, newTier)
  end
end

-- Modify the unit's belief by a given amount
function modifyUnitBelief(unit, belief, modifier)
  local currentStrength = getUnitBelief(unit, belief)
  local newStrength = clamp(currentStrength + modifier, -50, 50)

  setUnitBelief(unit, belief, newStrength)
end

-- Randomly assigns a tier based on default weight values and passes that on to setUnitBeliefTier
-- (Could alternatively use randomBeliefValue and setUnitBelief)
function randomiseUnitBelief(unit, belief)
  local randomTier = randomBeliefTier()

  setUnitBeliefTier(unit, belief, randomTier)
end

-- Get the "default" belief that the unit uses, based on their culture
-- Priority is Cultural Identity -> Civ's values -> Nothing
function getUnitCultureBelief(unit, belief)
  if belief == df.value_type.NONE then return 0 end
  local upers = unit.status.current_soul.personality

  if upers.cultural_identity ~= -1 then
    return df.cultural_identity.find(upers.cultural_identity).values[df.value_type[belief]]
  elseif upers.civ_id ~= -1 then
    return df.historical_entity.find(upers.civ_id).resources.values[df.value_type[belief]]
  else
    return 0 --outsiders have no culture
  end
end

-- Gets what tier of belief the given strength value falls into
-- 1: Lowest | 2: Very Low | 3: Low | 4: Neutral | 5: High | 6: Very High | 7: Highest
function getBeliefTier(strength)
  local range = 1
  for index, data in ipairs(tierRanges) do
    if strength >= data.min and strength <= data.max then
      range = index
      break
    end
  end

  return range
end

-- Currently a hack solution to randomisation.
-- Doesn't respect culture weights (if they have any impact)
function randomBeliefTier()
  local weightedTable = {} --as:{_type:table,id:number,weight:number}[]

  for rangeIndex, data in ipairs(tierRanges) do
    local addition = {}
    addition.id = rangeIndex
    addition.weight = data.weight
    table.insert(weightedTable, addition)
  end

  -- Determine the range tier
  local result = weightedRoll(weightedTable)

  return result
end

-- Uses above, but then gets a strength value within the tier's bounds
function randomBeliefValue()
  local randomTier = randomBeliefTier()

  -- Get a random value within that range
  return randomiseWithinBounds(tierRanges[randomTier].min, tierRanges[randomTier].max)
end

-- Print unit's beliefs into the dfhack console. Cultural beliefs are marked with a *
function printUnitBeliefs(unit)
  print("Beliefs for " .. dfhack.df2console(dfhack.units.getReadableName(unit)) .. ":")
  for id, name in ipairs(df.value_type) do
    if name ~= 'NONE' then
      local strength = getUnitBelief(unit, id)
      local cultural = isCultureBelief(unit, id)

      local out = name
      if cultural then
        out = out .. "*"
      end
      out = out .. " " .. strength

      print(out)
    end
  end
end

function main(...)
  local args = utils.processArgs({...}, validArgs)
  local setneed = dfhack.reqscript("modtools/set-need")

  if args.help then
    print(dfhack.script_help())
    return
  end

  if args.list then
    for index, valueName in ipairs(df.value_type) do
      if valueName ~= 'NONE' then
        print(index .. " (" .. valueName .. ")")
      end
    end
    return
  end

  -- Valid target check
  local unitsList = {} --as:df.unit[]
  if not args.citizens then
    -- Assume trying to target a unit
    local unit
    if args.unit then
        if tonumber(args.unit) then
            unit = df.unit.find(tonumber(args.unit))
        end
    end

    -- If unit ID wasn't provided / unit couldn't be found,
    -- Try getting selected unit
    if unit == nil then
        unit = dfhack.gui.getSelectedUnit(true)
    end

    if unit == nil then
        qerror("Couldn't find unit. If you don't want to target a specific unit, use -citizens.")
    else
        table.insert(unitsList, unit)
    end
  elseif args.citizens then
    -- Technically this will exclude insane citizens, but this is the
    -- easiest thing that dfhack provides

    -- Abort if not in Fort mode
    if not dfhack.world.isFortressMode() then
        qerror("-citizens argument only available in Fortress Mode.")
    end

    unitsList = dfhack.units.getCitizens(false, true)
  end

  -- Belief check
  local beliefsList = {} --as:number[]
  if not args.belief and not args.all and not args.listunit then
    qerror("Please specify a belief to edit using -belief <ID> or -all.")
  end
  if not args.all then -- One belief provided
    local beliefId
    if tonumber(args.belief) then
      beliefId = tonumber(args.belief)
    else
      beliefId = df.value_type[args.belief:upper()]
    end
    table.insert(beliefsList, beliefId)
  else -- To modify: all
    for id, keyName in ipairs(df.value_type) do
      if keyName ~= 'NONE' then
        table.insert(beliefsList, id)
      end
    end
  end

  -- Modifier check
  local modifyStyle
  local value
  if args.set then
    modifyStyle = "set"
    value = tonumber(args.set)
  elseif args.tier then
    modifyStyle = "tier"
    value = tonumber(args.tier)
  elseif args.modify then
    modifyStyle = "modify"
    value = tonumber(args.modify)
  elseif args.step then
    modifyStyle = "step"
    value = tonumber(args.step)
  elseif args.random then
    modifyStyle = "random"
  elseif args.default then
    modifyStyle = "default"
  elseif args.listunit then --A slight hack but it's fine
    modifyStyle = "list"
  else
    qerror("Please provide a valid modifier.")
  end

  -- Execute
  for index, unit in ipairs(unitsList) do
    for index, belief in ipairs(beliefsList) do
      if modifyStyle == "set" then
        setUnitBelief(unit, belief, value)
      elseif modifyStyle == "tier" then
        setUnitBeliefTier(unit, belief, value)
      elseif modifyStyle == "modify" then
        modifyUnitBelief(unit, belief, value)
      elseif modifyStyle == "step" then
        stepUnitBelief(unit, belief, value)
      elseif modifyStyle == "random" then
        randomiseUnitBelief(unit, belief)
      elseif modifyStyle == "default" then
        removeUnitBelief(unit, belief)
      end
    end

    if modifyStyle == "list" then
        printUnitBeliefs(unit)
    end
    if not args.noneed and not args.listunit then
      setneed.rebuildNeeds(unit)
    end
  end
end

if not dfhack_flags.module then
  main(...)
end
