local computer = require('computer')
local robot = require("robot");
local component = require("component")
local inv = component.inventory_controller
local geo = component.geolyzer
local chunk = component.chunkloader

chunk.setActive(true)

local branchLength = 32
local turnRight = true

local toolSlotCount = 2

function newSet(arr)
  local set = {}
  for i, v in ipairs(arr) do
    set[v] = true
  end
  return set
end

local slots = {
  {'charger', 1},
  {'lever', 1},
  {'battery', 1},
  {'fortune', 1},
  {'tool', 2}
}

local toolIds = newSet({
  'minecraft:stone_pickaxe',
  'minecraft:iron_pickaxe',
  'minecraft:diamond_pickaxe'
})

local validItemIds = {
  charger = newSet({'opencomputers:charger'}),
  lever = newSet({'minecraft:lever'}),
  battery = newSet({'thermalexpansion:cell'}),
  fortune = toolIds,
  tool = toolIds
}

slotIdRange = {}

local itemSlotCount = 0
local currentSlot = robot.inventorySize()

for i = #slots, 1, -1 do
  slotIdRange[slots[i][1]] = {currentSlot - (slots[i][2] - 1), currentSlot}
  itemSlotCount = itemSlotCount + slots[i][2]
  currentSlot = currentSlot - slots[i][2]
end

local validTools = {
  'minecraft:stone_pickaxe',
  'minecraft:iron_pickaxe',
  'minecraft:diamond_pickaxe'
}

local ores = newSet({
  'minecraft:gold_ore',
  'minecraft:iron_ore',
  'minecraft:coal_ore',
  'minecraft:lapis_ore',
  'minecraft:diamond_ore',
  'minecraft:redstone_ore',
  'minecraft:emerald_ore',
  'thaumcraft:ore_cinnabar',
})

local fortuneOres = newSet({
  'minecraft:diamond_ore',
  'minecraft:emerald_ore',
  'thermalfoundation:ore',
})

local bottom = 0
local top = 1
local back = 2
local front = 3
local right = 4
local left = 5

function forward()
  for i = 1, 5, 1 do
    if robot.forward() then
      break
    else
      swingIfTool()
    end
  end
end

function back()
  for i = 1, 5, 1 do
    if robot.back() then
      break
    else
      robot.turnAround()
      swingIfTool()
      robot.turnAround()
    end
  end
end

function up()
  for i = 1, 5, 1 do
    if robot.up() then
      break
    else
      swingUpIfTool()
    end
  end
end

function down()
  for i = 1, 5, 1 do
    if robot.down() then
      break
    else
      swingDownIfTool()
    end
  end
end

function place()
  for i = 1, 5, 1 do
    if robot.place() then
      break
    else
      swingIfTool()
    end
  end
end

function placeUp()
  for i = 1, 5, 1 do
    if robot.placeUp() then
      break
    else
      swingUpIfTool()
    end
  end
end

function placeDown()
  for i = 1, 5, 1 do
    if robot.placeDown() then
      break
    else
      swingDownIfTool()
    end
  end
end

function fortuneSwap()
  if slotIdRange['fortune'] then
    local item = inv.getStackInInternalSlot(slotIdRange['fortune'][1])
    if isValid('fortune', item) then
      equipSlot(slotIdRange['fortune'][1])
    end
  end
end

function isValid(kind, item)
  if item == nil then
    return false
  end
  return validItemIds[kind][item.name]
end

function isOre(block)
  return ores[block.name] or false
end

function isFortuneOre(block)
  return fortuneOres[block.name] or false
end

function mineOre()
  for i = 0, 5, 1 do
    local block = geo.analyze(i)
    if isOre(block) then
      local isFortune = isFortuneOre(block)
      if isFortune then
        fortuneSwap()
      end
      mineInDir(i)
      if isFortune then
        fortuneSwap()
      end
      mineOre()
      undoMineInDir(i)
    end
  end
end

function mineInDir(dir)
  if dir == 0 then
    swingDownIfTool()
    down()
  elseif dir == 1 then
    swingUpIfTool()
    up()
  elseif dir == 2 then
    robot.turnAround()
    swingIfTool()
    forward()
  elseif dir == 3 then
    swingIfTool()
    forward()
  elseif dir == 4 then
    robot.turnRight()
    swingIfTool()
    forward()
  elseif dir == 5 then
    robot.turnLeft()
    swingIfTool()
    forward()
  else
    error('ERROR: invalid dir')
  end
end

function undoMineInDir(dir)
  if dir == 0 then
    up()
  elseif dir == 1 then
    down()
  elseif dir == 2 then
    robot.turnAround()
    forward()
  elseif dir == 3 then
    back()
  elseif dir == 4 then
    back()
    robot.turnLeft()
  elseif dir == 5 then
    back()
    robot.turnRight()
  else
    error('ERROR: invalid dir')
  end
end

function turnInBranchDir()
  if turnRight then
    robot.turnRight()
  else
    robot.turnLeft()
  end
end

function swingIfTool(dontRecharge)
  for i = 1, 5, 1 do
    if robot.durability() == nil then
      equip('tool')
    end
    robot.swing()
    raiseIfInventoryFull()
    if not dontRecharge then
      rechargeIfNeeded()
    end
    if not robot.detect() then
      break
    end
  end
end

function swingUpIfTool(dontRecharge)
  for i = 1, 5, 1 do
    if robot.durability() == nil then
      equip('tool')
    end
    robot.swingUp()
    raiseIfInventoryFull()
    if not dontRecharge then
      rechargeIfNeeded()
    end
    if not robot.detectUp() then
      break
    end
  end
end

function swingDownIfTool(dontRecharge)
  for i = 1, 5, 1 do
    if robot.durability() == nil then
      equip('tool')
    end
    robot.swingDown()
    raiseIfInventoryFull()
    if not dontRecharge then
      rechargeIfNeeded()
    end
    if not robot.detectDown() then
      break
    end
  end
end

function hasAny(kind)
  for i = slotIdRange[kind][1], slotIdRange[kind][2], 1 do
    if isValid(kind, inv.getStackInInternalSlot(i)) then
      return true
    end
  end
  return false
end

function hasChargedBattery()
  for i = slotIdRange['battery'][1], slotIdRange['battery'][2], 1 do
    local stack = inv.getStackInInternalSlot(i)
    if isValid('battery', inv.getStackInInternalSlot(i)) then
      return stack.Energy and stack.Energy > 1000
    end
  end
  return false
end

function raiseIfInventoryFull()
  local lastSlot = robot.inventorySize() - itemSlotCount
  if inv.getStackInInternalSlot(lastSlot) ~= nil then
    error('ERROR: inventory full')
  end
end

function includes(arr, item)
  for i, v in ipairs(arr) do
    if v == item then
      return true
    end
  end
  return false
end

function equip(kind)
  local itemFound = false
  for i = slotIdRange[kind][1], slotIdRange[kind][2], 1 do
    if isValid(kind, inv.getStackInInternalSlot(i)) then
      equipSlot(i)
      itemFound = true
    end
  end
  if not itemFound then
    error('ERROR: no item \''..kind..'\' found')
  end
end

function selectItem(kind)
  local itemFound = false
  for i = slotIdRange[kind][1], slotIdRange[kind][2], 1 do
    if isValid(kind, inv.getStackInInternalSlot(i)) then
      robot.select(i)
      itemFound = true
    end
  end
  if not itemFound then
    error('ERROR: no item \''..kind..'\' found')
  end
end

function selectBlankItemSlot(kind)
  local blankSlotFound = false
  for i = slotIdRange[kind][1], slotIdRange[kind][2], 1 do
    if inv.getStackInInternalSlot(i) == nil then
      robot.select(i)
      blankSlotFound = true
    end
  end
  if not blankSlotFound then
    error('ERROR: no blank slot found for item \''..kind..'\'')
  end
end

function equipSlot(slotIndex)
  robot.select(slotIndex)
  inv.equip()
  robot.select(1)
end

function fixPositionOf(kind)
  for i = 1, robot.inventorySize() - itemSlotCount, 1 do
    if isValid(kind, inv.getStackInInternalSlot(i)) then
      for j = slotIdRange[kind][1], slotIdRange[kind][2], 1 do
        if not isValid(kind, inv.getStackInInternalSlot(j)) then
          robot.select(i)
          robot.transferTo(j)
          robot.select(1)
          break
        end
      end
    end
  end
end

function putAway()
  foundSlot = false
  for i = slotIdRange['tool'][1], slotIdRange['tool'][2], 1 do
    if not isValid('tool', inv.getStackInInternalSlot(i)) then
      equipSlot(i)
      foundSlot = true
      break
    end
  end
  if not foundSlot then
    error('ERROR: not open tool slot found')
  end
end

local high = false

function mineForward()
  swingIfTool()
  forward()
  if not high then
    placeDownIfAir()
  end
  mineOre()
  if high then
    swingDownIfTool()
    down()
    placeDownIfAir()
  else
    swingUpIfTool()
    up()
  end
  high = not high
  mineOre()
end

function placeDownIfAir()
  if not robot.detectDown() then
    robot.placeDown()
  end
  if not robot.detectDown() then
    error("ERROR: place failed!")
  end
end

function branch()
  turnInBranchDir()
  for i = 1, branchLength, 1 do
    mineForward()
  end
  robot.turnAround()
  for i = 1, branchLength, 1 do
    forward()
  end
  turnInBranchDir()
end

function rechargeIfNeeded()
  local energyLow = computer.energy() / computer.maxEnergy() < 0.15
  if energyLow and hasChargedBattery() then
    recharge()
  end
end

function recharge()
  robot.turnAround()
  swingIfTool(true)
  swingUpIfTool(true)
  forward()
  swingUpIfTool(true)
  swingIfTool(true)
  forward()
  swingUpIfTool(true)
  robot.turnAround()
  forward()
  selectItem('battery')
  place()
  back()
  selectItem('charger')
  place()
  up()
  selectItem('lever')
  place()
  robot.use()
  down()
  while computer.energy() / computer.maxEnergy() < .99 do
    os.sleep(3)
  end
  up()
  swingIfTool(true)
  down()
  selectBlankItemSlot('charger')
  swingIfTool(true)
  forward()
  selectBlankItemSlot('battery')
  swingIfTool(true)
  robot.select(1)
  forward()
end

while true do
  for i = 1, 4, 1 do
    mineForward()
  end
  branch()
end