-- run on a computer with a wired network

local pretty = require("cc.pretty")
local ledger_file = "ledger.json"

local function panic(msg)
  term.setTextColor(colors.pink)
  write(msg)
  term.setTextColor(colors.white)
  read()
end

local function loadLedger()
  if not fs.exists(ledger_file) then
    return {}
  end

  local file = fs.open(ledger_file, "r")
  local data = file.readAll()
  file.close()
  assert(data)
  local table = textutils.unserializeJSON(data)
  assert(table)
  return table
end

local ledger = loadLedger()

local function saveLedger()
  local file = fs.open(ledger_file, "w")
  file.writeLine(textutils.serializeJSON(ledger))
  file.close()
end

local function getChestContent(number)
  local name = "minecraft:chest_" .. number
  local chest = peripheral.wrap(name)

  local res = {}
  for index,item in pairs(chest.list()) do
    if res[item.name] == nil then
      res[item.name] = {count = 0, where = {}}
    end
    res[item.name].count = res[item.name].count + item.count
    table.insert(res[item.name].where, index)
  end

  return res
end

local function pprint(obj)
  pretty.pretty_print(obj)
end

local function verifyLedger(from_chest, item_name)
  local content = getChestContent(from_chest)
  local content_item = content[item_name]
  if content_item == nil then
    if ledger[item_name].count == 0 then
      print(string.format("we ran out of %s. it's supposed to be in chest %d, but there's 0 of it", item_name, from_chest))
    else
      print(string.format("no %s in chest %d. ledger must be wrong...", item_name, from_chest))
    end
    if ledger[item_name].old_chests and #ledger[item_name].old_chests > 0 then
      local new_chest = table.remove(ledger[item_name].old_chests)
      ledger[item_name].chest = new_chest
      term.setTextColor(colors.cyan)
      print(string.format("using old chest %d!", new_chest))
      term.setTextColor(colors.white)
      saveLedger()

      return verifyLedger(new_chest, item_name)
    end

    pprint(ledger[item_name])
    panic("press any key to remove that item form the ledger")
    ledger[item_name] = nil
    saveLedger()
    return {[item_name] = {where = {}}}
  end
  return content
end

local function countToCategory(count)
  if count <= 1 then return 0 -- todo
  elseif count <= 16 then return 0
  elseif count <= 64 then return 1
  elseif count <= 64*3 then return 2
  elseif count <= 64*9 then return 5
  elseif count <= 64*9*2 then return 6
  elseif count <= 64*9*3 then return 3
  else return 4 end
end
local MAX_CATEGORY = countToCategory(math.huge)
local CATEGORY_COUNT = 6

local function getConnectedChests()
  local names = peripheral.getNames()
  local numbers = {}
  for _,name in pairs(names) do
    if name:sub(1,16) == "minecraft:chest_" then
      local number = tonumber(name:sub(17,-1))
      table.insert(numbers, number)
    end
  end

  table.sort(numbers)
  return numbers
end

local function isChestReserved(chest)
  return chest >= 87 and chest < 94
end

local function getUnallocatedChests()
  local chests = getConnectedChests()
  local chests_bool = {}
  for _,chest in pairs(chests) do
    chests_bool[chest] = true
  end

  local function excludeCategory(category_name)
    if ledger[category_name] then
      for _,chest in pairs(ledger[category_name]) do
        chests_bool[chest] = false
      end
    end
  end

  -- exclude chests that already have a category
  for i=0,CATEGORY_COUNT do
    excludeCategory("_category" .. i)
  end
  excludeCategory("_category_reserved")

  -- exclude reserved chests >= 87
  for _,chest in pairs(chests) do
    if isChestReserved(chest) then
      chests_bool[chest] = false
    end
  end

  local res = {}
  for chest,bool in pairs(chests_bool) do
    if bool then table.insert(res, chest) end
  end

  table.sort(res)
  return res
end

local function getChestFillLevel(number)
  local name = "minecraft:chest_" .. number
  local chest = peripheral.wrap(name)
  if chest == nil then
    print(name .. " is not contented!")
  end

  local list = chest.list()
  return #list, chest.size()
end

local function isChestFull(number)
  local filled, total = getChestFillLevel(number)
  return filled == total
end

local function isChestEmpty(number)
  local filled, _ = getChestFillLevel(number)
  return filled == 0
end

local function allocateChestPos(cat, max)
  local res = 1 + cat*10
  if res > max then res = max end
  return res
end

local function addOldChest(item_name)
  if ledger[item_name] == nil then return end
  if ledger[item_name].old_chests == nil then ledger[item_name].old_chests = {} end
  table.insert(ledger[item_name].old_chests, ledger[item_name].chest)
end

local function addToReserved(chest)
  print("chest not empty", chest)
  if ledger._category_reserved == nil then ledger._category_reserved = {} end
  table.insert(ledger._category_reserved, chest)
end

local function findChest(cat, needed_space)
  needed_space = needed_space or 0
  local category_name = "_category" .. cat
  if not ledger[category_name] then ledger[category_name] = {} end

  for _,v in pairs(ledger[category_name]) do
    local filled, total = getChestFillLevel(v)
    local empty = total - filled
    if filled < total then
      if (cat ~= MAX_CATEGORY and empty >= needed_space) or filled == 0 then
        return v
      end
    end
  end

  while true do
    local unallocated_chests = getUnallocatedChests()
    local chest_to_allocate = unallocated_chests[allocateChestPos(cat, #unallocated_chests)]
    if not isChestEmpty(chest_to_allocate) then
      addToReserved(chest_to_allocate)
      saveLedger()
    else
      table.insert(ledger[category_name], chest_to_allocate)
      return chest_to_allocate
    end
  end
end

local function transferItems(from_chest, from_slot, to_chest)
  local name1 = "minecraft:chest_" .. from_chest
  local name2 = "minecraft:chest_" .. to_chest
  local chest1 = peripheral.wrap(name1)

  return chest1.pushItems(name2, from_slot)
end

local function transferAllItems(from_chest, item_name, to_chest, content)
  content = content or getChestContent(from_chest)

  local total = 0
  for _,slot in pairs(content[item_name].where) do
    total = total + transferItems(from_chest, slot, to_chest)
  end

  return total, string.format("moved %d of %s [%d -> %d]", total, item_name, from_chest, to_chest)
end

local function ensureLedgerStruct(item_name, new_chest)
  if ledger[item_name] == nil then
    ledger[item_name] = {count = 0, chest = new_chest}
  else
    ledger[item_name].chest = new_chest
  end
end

local function changeCategory(from_chest, item_name, new_cat)
  local count = ledger[item_name].count
  local needed_space = math.ceil(count / 64) + 1
  -- `+1` is a hack, because sometimes it can get stuck in a loop switching back and forth

  local new_chest = findChest(new_cat, needed_space)
  local content = verifyLedger(from_chest, item_name)
  if ledger[item_name] then from_chest = ledger[item_name].chest end
  local total, _ = transferAllItems(from_chest, item_name, new_chest, content)

  -- here we might lose track of some items, if `new_chest` actually doesn't have enough space
  -- but with `needed_space` it should be fine
  if total < count then
    print(total, "<", count)
    panic("couldn't move all the items to new chest...")
    addOldChest(item_name)
  end

  ensureLedgerStruct(item_name, new_chest)
  saveLedger()
end

local function handleNoSpace(item_name)
  local cat = countToCategory(ledger[item_name].count)

  if cat == MAX_CATEGORY then
    print("adding a new chest for " .. item_name)
    addOldChest(item_name)
    ledger[item_name].chest = findChest(cat)
    saveLedger()
  else
    print("moving " .. item_name .. " to a new chest in the same category")
    changeCategory(ledger[item_name].chest, item_name, cat)
  end
end

local function whereDoIPutThat(item_name, item_count)
  if ledger[item_name] and ledger[item_name].count > 0 then
    local cat1 = countToCategory(ledger[item_name].count)
    local cat2 = countToCategory(ledger[item_name].count + item_count)
    if cat1 == cat2 then
      if isChestFull(ledger[item_name].chest) then
        handleNoSpace(item_name)
      end
      return ledger[item_name].chest
    else
      print(string.format("category for %s changed form %d to %d", item_name, cat1, cat2))
      changeCategory(ledger[item_name].chest, item_name, cat2)
      return ledger[item_name].chest
    end
  else
    local cat = countToCategory(item_count)
    local chest = findChest(cat)

    ledger[item_name] = {count = 0, chest = chest}
    return chest
  end
end

local function scanInputChest(number)
  local content = getChestContent(number)
  pprint(content)

  for k,v in pairs(content) do
    local chest = whereDoIPutThat(k, v.count)
    print(k, "->", chest)
    for _,slot in pairs(v.where) do
      local res = transferItems(number, slot, chest)
      ledger[k].count = ledger[k].count + res
      saveLedger()
    end
  end
end

local function importMany(args)
  for i=2,#args do
    scanInputChest(tonumber(args[i]))
  end
  saveLedger()
end

local function usage()
  print("usage: chests [import|export|whereis|shell] [chest] [chests...|item]")
end

local function boolToNumber(value)
  return value and 1 or 0
end

local function levenshteinDistance_impl(d, stride, a, b, i, j)
  if i == 0 then return j end
  if j == 0 then return i end
  local index = (i-1) + (j-1) * stride
  if d[index] ~= nil then return d[index] end

  local r1 = levenshteinDistance_impl(d,stride,a,b, i-1, j) + 1
  local r2 = levenshteinDistance_impl(d,stride,a,b, i, j-1) + 1
  local r3 = levenshteinDistance_impl(d,stride,a,b, i-1, j-1) + boolToNumber(a[i] ~= b[j])

  local res = math.min(r1, r2, r3)
  d[index] = res
  return res
end

local function stringToCharArray(str)
  local t = {}
  str:gsub(".",function(c) table.insert(t,c) end)
  return t
end

local function levenshteinDistance(a, b)
  local d = {}
  return levenshteinDistance_impl(d, a:len(), stringToCharArray(a), stringToCharArray(b), a:len(), b:len())
end

local function findItemId(search)
  if ledger[search] then return search end

  local candidate = nil
  local min_distance = 9999
  for k,v in pairs(ledger) do
    if k:find(search) and v.count > 0 then
      local distance = levenshteinDistance(search, k)
      if min_distance > distance then
        min_distance = distance
        candidate = k
      end
    end
  end

  return candidate
end

local function whereis(search)
  local name = findItemId(search)
  if name then
    term.setTextColor(colors.lightGray)
    print(name)
    term.setTextColor(colors.white)
    pprint(ledger[name])
  else
    print(string.format("no '%s' found", search))
  end
end

local function whereisMany(args)
  for i=2,#args do
    whereis(args[i])
    if i ~= #args then print() end
  end
end

local function export(destination, search, limit)
  limit = limit or (64*9)
  local item_name = findItemId(search)
  if item_name == nil then
    print(string.format("couldn't find anything in the system that matched '%s'", search))
    return
  end
  local from_chest = ledger[item_name].chest
  local content = verifyLedger(from_chest, item_name)
  if ledger[item_name] then from_chest = ledger[item_name].chest end

  local total = 0
  for _,slot in pairs(content[item_name].where) do
    local delta = transferItems(from_chest, slot, destination)
    ledger[item_name].count = ledger[item_name].count - delta
    saveLedger()

    limit = limit - delta
    total = total + delta
    if limit <= 0 then break end
  end

  write("exported ")
  term.setTextColor(colors.lightGray)
  write(item_name)
  term.setTextColor(colors.white)
  write(" to chest ")
  term.setTextColor(colors.lime)
  write(destination)
  term.setTextColor(colors.white)
  print(" (" .. total .. " items total)")
end

local function exportMany(args)
  local chest = tonumber(args[2])
  for i=3,#args do
    export(chest, args[i])
  end
end

local function startsWith(haystack, needle)
  return haystack:sub(1, needle:len()) == needle
end

local function runShellCommand(command, chest)
  ledger = loadLedger()
  if command == "import" then
    scanInputChest(chest)
  elseif startsWith(command, "import ") and command:len() > 7 then
    scanInputChest(tonumber(command:sub(8, -1)))
  else
    export(chest, command)
  end
  saveLedger()
end

local function shellPrompt()
  term.setTextColor(colors.lime)
  write(">>> ")
  term.setTextColor(colors.white)
end

local function exportShell(chest)
  term.setTextColor(colors.lightGray)
  print("started export shell to chest " .. chest)
  print("(empty string to exit)")
  term.setTextColor(colors.white)

  local function shell_completion(text)
    local res = {}
    if text:len() == 0 then return res end

    if startsWith("import", text) then
      table.insert(res, ("import"):sub(text:len()+1, -1))
    end

    for k,v in pairs(ledger) do
      if k:sub(1,1) == "_" then
      elseif v.count == 0 then
      elseif startsWith(k, text) then
        table.insert(res, k:sub(text:len()+1, -1))
      elseif not text:find(":") and k:find(":") then
        local index = k:find(":")
        local new_k = k:sub(index+1, -1)
        if startsWith(new_k, text) then
          table.insert(res, new_k:sub(text:len()+1, -1))
        end
      end
    end
    return res
  end

  if ledger._history == nil then ledger._history = {} saveLedger() end
  while true do
    shellPrompt()
    local res = read(nil, ledger._history, shell_completion)
    ledger = loadLedger()
    if res == "" then return end
    runShellCommand(res, chest)
    table.insert(ledger._history, res)
    saveLedger()
  end
end

local function setSensibleCategoryLocations()
  ledger._category0 = {60, 59, 82, 10}
  ledger._category1 = {11, 12, 13, 14}
  ledger._category2 = {15, 16, 17, 18}
  ledger._category3 = {19, 20, 21, 22}
  ledger._category4 = {65, 66, 67, 68, 77, 78, 79, 80, 76, 75, 74, 73, 72}
  ledger._category5 = {94, 95, 96, 97}
  ledger._category6 = {110, 111, 112, 113}
end

local function emptyChest(number)
  local content = getChestContent(number)
  for item_name,v in pairs(content) do
    local cat = countToCategory(v.count)
    changeCategory(ledger[item_name].chest, item_name, cat)
  end
end

local function fuckAroundAndFindOut(args)
  print("fucking around and finding out!!")
  setSensibleCategoryLocations()
  for i=2,#args do
    emptyChest(tonumber(args[i]))
  end
  saveLedger()
end

local function safeExport(destination, item_name, limit)
  if ledger[item_name] and ledger[item_name].count > 0 then
    export(destination, item_name, limit)
  else
    print(string.format("can't export %s because we don't have any", item_name))
  end
end

local function smelt(args)
  -- should probably be in a separate file or something
  local smelt_ore_chest = 92
  local smelt_fuel_chest = 93
  local smelt_ingot_chest = 91

  local to_smelt = args
  if #args == 1 then
    to_smelt = {"NULL", "minecraft:raw_copper", "minecraft:raw_gold", "minecraft:raw_iron"}
  end

  safeExport(smelt_fuel_chest, "minecraft:lava_bucket")
  for i=2,#to_smelt do
    safeExport(smelt_ore_chest, to_smelt[i])
  end
  scanInputChest(smelt_ingot_chest)
  scanInputChest(87)
  scanInputChest(89)
  scanInputChest(90)
  saveLedger()
end

local function listenForMessages(chest)
  local modem = peripheral.find("modem")
  modem.open(1)
  while true do
    local event, side, channel, replyChannel, message, distance = os.pullEvent("modem_message")
    if channel == 1 then
      print("received message: " .. message)
      runShellCommand(message, chest)
      term.setTextColor(colors.lightBlue)
      print("finished importing!")
      term.setTextColor(colors.white)
      shellPrompt()
    end
  end
end

local function shellAndMessages(chest)
  parallel.waitForAny(
    function() listenForMessages(chest) end,
    function() exportShell(chest) end
  )
end

local function excludeEmptyChests(chests)
  local res = {}

  for _,chest in pairs(chests) do
    if not isChestEmpty(chest) then
      table.insert(res, chest)
    end
  end

  return res
end

local function rescanChests()
  for key,value in pairs(ledger) do
    if key:find(":") then
      ledger[key] = nil
    end
  end

  local chests = getConnectedChests()
  local chests_has_stuff = {}

  for _,chest in pairs(chests) do
    if isChestReserved(chest) then
      print("skipping chest " .. chest)
    else
      local content = getChestContent(chest)
      for k,v in pairs(content) do
        chests_has_stuff[chest] = true
        if ledger[k] == nil then
          ledger[k] = {count = v.count, chest = chest}
        else
          if ledger[k].count < 64*6*9 then
            print(k .. " is in more than one chest at once")
          end
          addOldChest(k)
          ledger[k].chest = chest
          ledger[k].count = ledger[k].count + v.count
          if ledger[k].count > 64*6*9 and ledger[k].count - v.count < 64*6*9 then
            term.setTextColor(colors.lightBlue)
            print("assuming category " .. MAX_CATEGORY .. " for " .. k)
            term.setTextColor(colors.white)
          end
        end
      end
    end
  end

  local unallocated_chests = getUnallocatedChests()
  for _,chest in pairs(unallocated_chests) do
    if chests_has_stuff[chest] then
      addToReserved(chest)
    end
  end

  ledger._category_reserved = excludeEmptyChests(ledger._category_reserved)
  ledger._category4 = excludeEmptyChests(ledger._category4)

  saveLedger()
end

local function excludeFullChests(item_name, chests)
  local count = 0
  local res = {}

  for _,chest in pairs(chests) do
    local content = getChestContent(chest)
    if content[item_name] and content[item_name].count < 64*6*9 then
      count = count + content[item_name].count
      table.insert(res, chest)
    end
  end

  return res, count
end

local function consolidateChests()
  for key,value in pairs(ledger) do
    if key:find(":") and value.old_chests ~= nil then
      local original_chest = value.chest
      local all_chests = {original_chest, table.unpack(value.old_chests)}
      local category = countToCategory(value.count)
      local non_full_chests, count = excludeFullChests(key, all_chests)
      local needed_space = math.ceil(count / 64)

      if #non_full_chests <= 1 then
        print("no need to consolidate " .. key)
      elseif needed_space <= 6*9 then
        print("consolidating item " .. key)

        local new_chest = findChest(category, needed_space)
        local total = 0
        for _,chest in pairs(non_full_chests) do
          local count_moved, msg = transferAllItems(chest, key, new_chest)
          total = total + count_moved
          print(msg)
        end

        ensureLedgerStruct(key, new_chest)
        if total < count then
          term.setTextColor(colors.orange)
          print("couldn't move all the items to new chest...")
          term.setTextColor(colors.white)
          table.insert(ledger[key].old_chests, original_chest)
        else
          ledger[key].old_chests = nil
        end
      else
        -- todo: also handle two almost-full chests
        term.setTextColor(colors.orange)
        print("not enough space for consolidation of item " .. key)
        print(needed_space .. " slots needed")
        term.setTextColor(colors.white)

        local new_chest = non_full_chests[1]
        if original_chest ~= new_chest then
          print(string.format("switching main chest %d -> %d", original_chest, new_chest))
          ensureLedgerStruct(key, new_chest)
          table.insert(ledger[key].old_chests, original_chest)
        end
      end
    end
  end

  saveLedger()
end

local function tableContains(table, value)
  for i,v in pairs(table) do
    if value == v then return i end
  end
  return nil
end

local function detectWrongCategories()
  for key,value in pairs(ledger) do
    if key:find(":") and value.count > 0 then
      local chest = value.chest
      local category = countToCategory(value.count)
      local category_array = ledger["_category" .. category]
      if not tableContains(category_array, chest) then
        local index = tableContains(ledger._category_reserved, chest)
        if index then
          print(string.format("moving chest %d to category %d (it was reserved)", chest, category))
          table.remove(ledger._category_reserved, index)
          table.insert(category_array, chest)
        else
          -- todo: move the item
          print(string.format("item %s in wrong category (chest %d is not in category %d)", key, chest, category))
        end
      end
    end
  end

  saveLedger()
end

local function recoverLedger()
  if next(ledger._history) == nil then ledger._history = nil end
  if next(ledger) ~= nil then
    print("ledger not empty")
    panic("press any key to rescan chests anyway...")
    rescanChests()
    consolidateChests()
    detectWrongCategories()
    return
  end

  setSensibleCategoryLocations()
  rescanChests()
end

local function checkChestSizes(args)
  local start_index = tonumber(args[2] or 1)
  local end_index = tonumber(args[3] or 130)
  for i=start_index,end_index do
    local chest = peripheral.wrap("minecraft:chest_" .. i)
    if chest then print(i, chest.size()) end
  end
end

local function main(args)
  if #args < 1 then
    usage()
    print("got the wrong number of args...")
  elseif args[1] == "import" then
    importMany(args)
  elseif args[1] == "fuck_with_ledger" then
    fuckAroundAndFindOut(args)
  elseif args[1] == "print_ledger" then
    pprint(ledger)
    write("getUnallocatedChests() ==> ")
    pprint(getUnallocatedChests())
  elseif args[1] == "smelt" then
    smelt(args)
  elseif args[1] == "check_size" then
    checkChestSizes(args)
  elseif args[1] == "export" then
    exportMany(args)
  elseif args[1] == "shell" then
    shellAndMessages(tonumber(args[2]))
  elseif args[1] == "recover_ledger" then
    recoverLedger()
  elseif args[1] == "whereis" then
    whereisMany(args)
  else
    usage()
    print("unknown command...")
  end
end

local args = {...}
main(args)