package.path = "/mod/?.lua;" .. package.path
local PriorityQueue = require "PriorityQueue"

function Set(list)
  local set = {}
  for _, l in ipairs(list) do set[l] = true end
  return set
end

local yieldTime = 0
function yield()
  if os.clock() - yieldTime > 2 then
    sleep(0)
    yieldTime = os.clock()
  end
end

function arr(n, init)
  local a = {}
  for _ = 1, n do
    table.insert(a, init)
  end
  return a
end
function arr2(x, y, init)
  local a = {}
  for _ = 1, y do
    table.insert(a, arr(x, init))
  end
  return a
end

local dirMap = {
  ["0,-1"] = 0,
  ["0,1"] = 0.5,
  ["1,0"] = 0.25,
  ["-1,0"] = -0.25,
}
local invDirMap = {
  [0] = vector.new(0, 0, -1),
  [0.5] = vector.new(0, 0, 1),
  [0.25] = vector.new(1, 0, 0),
  [-0.25] = vector.new(-1, 0, 0),

  [-0.5] = vector.new(0, 0, 1),
  [-0.75] = vector.new(1, 0, 0),
  [0.75] = vector.new(-1, 0, 0),
}
function direction(v)
  local d = dirMap[v.x..","..v.z]
  assert(d)
  return d
end

function findDistances(points)
  local dist = arr2(#points, #points, math.huge)

  for i = 1, #points do
    for j = 1, #points do
      if i ~= j then
        dist[i][j] = (points[i] - points[j]):length()
      end
    end
  end
  for i = 1, #points do
    dist[i][i] = 0
  end

  return dist
end

function groupOres(ores, dist, offset)
  if not offset then
    offset = 0
  end

  local groups = {}
  for i, o in ipairs(ores) do
    table.insert(groups, {
      ["type"] = o.name,
      ["blocks"] = {
        [i] = vector.new(o.x, o.y, o.z),
      },
    })
  end

  local didJoin = true
  while didJoin do
    didJoin = false
    for i, io in pairs(groups) do
      for x, a in pairs(io.blocks) do
        for j, jo in pairs(groups) do
          if i ~= j and jo.type == io.type then
            for y, b in pairs(jo.blocks) do
              -- yield()
              assert(a ~= b)
              if dist[offset + x][offset + y] == 1 then
                for z, p in pairs(jo.blocks) do
                  io.blocks[z] = p
                end
                groups[j] = nil
                didJoin = true
                break
              end
            end
          end
        end
      end
    end
  end

  local retGroups = {}
  for _, g in pairs(groups) do
    local v = vector.new(0, 0, 0)
    local retBlocks = {}
    for _, p in pairs(g.blocks) do
      v = v + p
      table.insert(retBlocks, p)
    end
    g.blocks = retBlocks
    g.center = v / #g.blocks
    assert(g.center)
    table.insert(retGroups, g)
  end

  return retGroups
end

-- nearest neighbour
function pathThrough(points, dist, s, e)
  local unvisited = {}
  local ui = nil
  for i = 1, #points do
    if not e or i ~= e then
      table.insert(unvisited, i)
    end
    if i == s then
      ui = #unvisited
    end
  end
  assert(ui)

  local visited = {}
  local route = {}
  while #unvisited ~= 1 do
    local u = unvisited[ui]
    table.remove(unvisited, ui)

    local cDist = math.huge
    local closest = nil
    for ii = 1, #unvisited do
      local i = unvisited[ii]
      assert(i ~= u)
      if dist[u][i] < cDist then
        cDist = dist[u][i]
        closest = ii
      end
    end
    assert(closest)
    table.insert(route, unvisited[closest])
    ui = closest
  end
  if e then
    table.insert(route, e)
  end

  return route
end

-- A*
function findPath(s, e, neighbors, limit)
  local points = {}
  function addPoint(p)
    local k = p:tostring()
    if not points[k] then
      points[k] = p
    end
    return k
  end
  s = addPoint(s)
  e = addPoint(e)

  function h(i)
    return (points[e] - points[i]):length()
  end

  local gScoreT = {}
  gScoreT[s] = 0
  function gScore(i)
    if gScoreT[i] then
      return gScoreT[i]
    else
      return math.huge
    end
  end

  local fScoreT = {}
  fScoreT[s] = h(s)
  function fScore(i)
    if fScoreT[i] then
      return fScoreT[i]
    else
      return math.huge
    end
  end

  local openSet = PriorityQueue(function(a, b)
    return fScore(a) < fScore(b)
  end)
  openSet:enqueue(s, s)

  local prev = {}
  local found = false
  local brokeLimit = false
  while not openSet:empty() do
    local current = openSet:dequeue()
    if current == e then
      found = true
      break
    elseif limit and (points[current] - points[s]):length() > limit then
      e = current
      brokeLimit = true
      break
    end

    local ns = neighbors(points[current])
    for i, n in ipairs(ns) do
      ns[i] = addPoint(n)
    end

    for _, n in ipairs(ns) do
      tentativeG = gScore(current) + 1
      if tentativeG < gScore(n) then
        prev[n] = current
        gScoreT[n] = tentativeG
        fScoreT[n] = tentativeG + h(n)
        if not openSet:contains(n) then
          openSet:enqueue(n, n)
        end
      end
    end
  end
  assert(found or brokeLimit)

  local path = {points[e]}
  local current = e
  while prev[current] do
    current = prev[current]
    table.insert(path, 1, points[current])
  end
  return path, brokeLimit
end

local ORES = Set{
  "minecraft:coal_ore",
  "minecraft:deepslate_coal_ore",
  "minecraft:iron_ore",
  "minecraft:deepslate_iron_ore",
  "minecraft:copper_ore",
  "minecraft:deepslate_copper_ore",
  "minecraft:gold_ore",
  "minecraft:deepslate_gold_ore",
  "minecraft:redstone_ore",
  "minecraft:deepslate_redstone_ore",
  "minecraft:emerald_ore",
  "minecraft:deepslate_emerald_ore",
  "minecraft:lapis_ore",
  "minecraft:deepslate_lapis_ore",
  "minecraft:diamond_ore",
  "minecraft:deepslate_diamond_ore",

  "minecraft:glowstone",
  "minecraft:nether_quartz_ore",
  "minecraft:ancient_debris",

  "create:zinc_ore",
  "create:deepslate_zinc_ore",
  "create_new_age:thorium_ore",
  "create_new_age:magnetite_block",

  "powah:deepslate_uraninite_ore_poor",
  "powah:deepslate_uraninite_ore",
  "powah:deepslate_uraninite_ore_dense",
  "powah:uraninite_ore_poor",
  "powah:uraninite_ore",
  "powah:uraninite_ore_dense",

  -- "minecraft:diamond_ore",
  -- "minecraft:netherrack",
}
local BLACKLIST = Set{
  "minecraft:diamond_pickaxe",
  "computercraft:wireless_modem_advanced",
  "advancedperipherals:geo_scanner",
  "enderchests:ender_chest",
  "functionalstorage:ender_drawer",
  "advancedperipherals:chunk_controller",
}
local FUEL = "minecraft:coal"
local PROTOCOL = "automine"
local PROTOCOL_CTRL = "automineCtrl"

settings.define("mine.name", {
  description = "Miner name",
  default = "DiggyBoi",
  type = "string",
})
settings.define("mine.stripWidth", {
  description = "Strip width",
  default = 16,
  type = "number",
})

function selectItem(item, nbt)
  for i = 1, 16 do
    info = turtle.getItemDetail(i)
    if (not item and not info) or (info and info.name == item and (not nbt or item.nbt == nbt)) then
      turtle.select(i)
      return i
    end
  end
end


Miner = {}
function Miner.new()
  local self = setmetatable({}, { __index = Miner })
  self.name = settings.get("mine.name")
  self.stripWidth = settings.get("mine.stripWidth")
  self.equipped = {
    ["left"] = {
      ["last"] = "NONE",
      ["lastNbt"] = nil,
      ["fun"] = turtle.equipLeft,
    },
    ["right"] = {
      ["last"] = "NONE",
      ["lastNbt"] = nil,
      ["fun"] = turtle.equipRight,
    },
  }

  -- Get upgrades to a known state by "equipping" nothing
  self:equip(nil, "left")
  self:equip(nil, "right")

  self:equip("advancedperipherals:geo_scanner", "right")
  self.scanner = peripheral.wrap("right")

  self:equip("computercraft:wireless_modem_advanced", "right")
  rednet.open("right")

  self:findDimension()
  self.pos = vector.new(0, 0, 0)
  self:doGPS(true)

  self:equip("advancedperipherals:chunk_controller", "left")

  self.i = 1
  return self
end

function Miner:equip(item, side, nbt)
  local e = self.equipped[side]
  if item == e.last and nbt == e.lastNbt then
    return
  end

  assert(selectItem(item, nbt))
  e.fun()
  e.last = item
  e.lastNbt = nbt
end

function Miner:sendMessage(type, msg)
  self:equip("computercraft:wireless_modem_advanced", "right")
  rednet.open("right")
  msg["type"] = type
  msg["name"] = self.name
  msg["dimension"] = self.dimension
  msg["pos"] = self.absolutePos
  rednet.broadcast(msg, PROTOCOL)
end

function Miner:dig(dir)
  self:equip("minecraft:diamond_pickaxe", "right")
  if dir == "up" then
    assert(turtle.digUp())
  elseif dir == "forward" then
    assert(turtle.dig())
  elseif dir == "down" then
    assert(turtle.digDown())
  end
end

function Miner:xchgItems(forceUnload)
  local slots = {}
  local emptySlots = 0
  for i = 1, 16 do
    local info = turtle.getItemDetail(i)
    if info then
      assert(info.name)
      if not BLACKLIST[info.name] then
        slots[i] = info
      end
    else
      emptySlots = emptySlots + 1
    end
  end

  if (emptySlots >= 2 and not forceUnload) and turtle.getFuelLevel() >= 1000 then
    return
  end

  if turtle.detectDown() then
    self:dig("down")
  end

  if emptySlots < 2 or forceUnload then
    selectItem("enderchests:ender_chest")
    assert(turtle.placeDown())

    for i, info in pairs(slots) do
      assert(turtle.select(i))
      assert(turtle.dropDown())
    end
    self:dig("down")
  end

  if turtle.getFuelLevel() < 1000 then
    selectItem("functionalstorage:ender_drawer")
    assert(turtle.placeDown())
    assert(turtle.suckDown())
    assert(turtle.refuel())
    self:dig("down")
  end
end

function Miner:safeNeighbors(p)
  local ns = {
    p + vector.new( 1,  0,  0),
    p + vector.new( 0,  0,  1),
    p + vector.new(-1,  0,  0),
    p + vector.new( 0,  0, -1),
  }
  -- only allow downwards if above / below bedrock
  if (self.absolutePos - self.pos + p).y - 1 > -60 then
    table.insert(ns, p + vector.new(0, -1, 0))
  end
  if self.dimension ~= "the_nether" or (self.absolutePos - self.pos + p).y + 1 < 123 then
    table.insert(ns, p + vector.new(0, 1, 0))
  end

  return ns
end

function Miner:safeMove(dir)
  if dir == "up" then
    while turtle.detectUp() do
      self:dig("up")
    end
    assert(turtle.up())
  elseif dir == "forward" then
    while turtle.detect() do
      self:dig("forward")
    end
    assert(turtle.forward())
  else
    if turtle.detectDown() then
      self:dig("down")
    end
    assert(turtle.down())
  end
end

function Miner:findDimension()
  self:equip("computercraft:wireless_modem_advanced", "right")
  local m = peripheral.wrap("right")

  local replyChannel = 666 + os.computerID()
  m.open(replyChannel)
  m.transmit(666, replyChannel)

  local event, side, channel, _, reply, distance = os.pullEvent("modem_message")
  assert(distance)
  assert(channel == replyChannel)
  self.dimension = reply
  m.close(replyChannel)

  print("Looks like my dimension is " .. self.dimension)
end

function Miner:doGPS(orientation)
  local x, y, z = gps.locate()
  assert(x)
  self.absolutePos = vector.new(x, y, z)

  if orientation then
    assert(turtle.forward())
    x, y, z = gps.locate()
    assert(x)
    self.dir = vector.new(x, y, z) - self.absolutePos
    assert(turtle.back())
  end
end

function Miner:findOres(radius)
  self:equip("advancedperipherals:geo_scanner", "right")
  local info, err = self.scanner.scan(radius)
  assert(info, err)

  local found = {}
  for _, b in ipairs(info) do
    -- only ores and avoid bedrock
    if ORES[b.name] and self.absolutePos.y + b.y > -60 and (self.dimension ~= "the_nether" or self.absolutePos.y + b.y < 123) then
      table.insert(found, b)
    end
  end
  return found
end

function Miner:faceDir(newDir)
  local deltaDir = direction(newDir) - direction(self.dir)
  if deltaDir == 0.25 or deltaDir == -0.75 then
    assert(turtle.turnRight())
    self.dir = newDir
  elseif deltaDir == -0.25 or deltaDir == 0.75 then
    assert(turtle.turnLeft())
    self.dir = newDir
  elseif deltaDir == 0.5 or deltaDir == -0.5 then
    assert(turtle.turnRight())
    assert(turtle.turnRight())
    self.dir = newDir
  elseif deltaDir == 0 then
    -- nothing to do :)
  else
    assert(false, "invalid delta dir "..deltaDir)
  end
end

function Miner:turnDir(delta)
  local newDirNum = direction(self.dir) + delta
  if newDirNum >= 1 then
    newDirNum = newDirNum - 1
  elseif newDirNum <= -1 then
    newDirNum = newDirNum + 1
  end

  self:faceDir(invDirMap[newDirNum])
end

function Miner:navigateThrough(path, limit)
  for _, p in ipairs(path) do
    -- print("move from to", self.pos, p)
    local delta = p - self.pos
    assert(delta:length() ~= 0)
    -- print("doing", delta)
    if delta:length() ~= 1 then
      -- print("path finding between points")
      local path, brokeLimit = findPath(self.pos, p, function(v) return self:safeNeighbors(v) end, limit)
      table.remove(path, 1)
      if not brokeLimit then
        table.remove(path, #path)
      end
      assert(#path ~= 0)

      self:navigateThrough(path)
      if brokeLimit then
        return false
      end
      delta = p - self.pos
    end
    assert(delta:length() == 1)

    local moveDir = nil
    if delta.y == 1 then
      moveDir = "up"
    elseif delta.y == -1 then
      moveDir = "down"
    else
      self:faceDir(delta)
      moveDir = "forward"
    end

    self:safeMove(moveDir)
    self.pos = p
    self.absolutePos = self.absolutePos + delta
  end
  return true
end

function Miner:checkRecall()
  local sender, msg = rednet.receive(PROTOCOL_CTRL, 0.5)
  if sender and msg.type == "recall" then
    -- HACK: Seems like we need this so the monitor can pick up the ack...
    sleep(0.5)
    self:sendMessage("ackRecall", msg)

    self:xchgItems(true)

    local to = vector.new(msg.to.x, msg.to.y, msg.to.z)
    print("Recalling to " .. to:tostring())
    while not self:navigateThrough({self.pos + (to - self.absolutePos)}, 64) do
      self:sendMessage("recallProgress", { to = msg.to })
      self:xchgItems()
    end

    self:xchgItems(true)
    return true
  end

  return false
end

function Miner:mineOres(radius)
  local startingDir = self.dir
  self.pos = vector.new(0, 0, 0)

  local ores = self:findOres(radius)
  local orePoints = {}
  for _, b in ipairs(ores) do
    table.insert(orePoints, vector.new(b.x, b.y, b.z))
  end
  local veins = groupOres(ores, findDistances(orePoints))
  print("Found "..#ores.." ores ("..#veins.." veins)")

  local veinsPoints = {
    self.pos,
    self.dir * radius * 2,
  }
  for _, v in ipairs(veins) do
    table.insert(veinsPoints, v.center)
  end

  local veinsPath = pathThrough(veinsPoints, findDistances(veinsPoints), 1, 2)
  -- strip out end point
  table.remove(veinsPath, #veinsPath)

  local veinsSummary = {}
  for _, i in ipairs(veinsPath) do
    local v = veins[i - 2]
    table.insert(veinsSummary, {
      pos = self.absolutePos + v.center,
      type = v.type,
      count = #v.blocks,
    })
  end
  self:sendMessage("iterationStart", {
    i = self.i,
    oreCount = #ores,
    veins = veinsSummary,
  })
  if self:checkRecall() then
    return false
  end

  for pi, i in ipairs(veinsPath) do
    self:xchgItems()

    local vein = veins[i - 2]
    local closest = nil
    local cDist = math.huge
    for _, b in ipairs(vein.blocks) do
      local d = (b - self.pos):length()
      if d < cDist then
        closest = b
        cDist = d
      end
    end
    assert(closest)

    -- Find path here so we can strip out the actual ore itself
    local pathToStart = findPath(self.pos, closest, function(v) return self:safeNeighbors(v) end)
    table.remove(pathToStart, 1)
    table.remove(pathToStart, #pathToStart)

    print("Moving to vein of " .. #vein.blocks .. " " .. vein.type .. " " .. cDist .. " blocks away starting at", closest)
    self:navigateThrough(pathToStart)

    local veinPoints = {self.pos}
    for _, b in ipairs(vein.blocks) do
      table.insert(veinPoints, b)
    end

    local veinPathI = pathThrough(veinPoints, findDistances(veinPoints), 1)
    local veinPath = {}
    for _, i in ipairs(veinPathI) do
      table.insert(veinPath, veinPoints[i])
    end

    print("Digging through vein (" .. #vein.blocks .. " blocks)")
    self:sendMessage("veinStart", {
      i = pi,
      total = #veins,
      oreType = vein.type,
      oreCount = #vein.blocks,
    })
    if self:checkRecall() then
      return false
    end

    self:navigateThrough(veinPath)
  end

  self:navigateThrough({veinsPoints[2]})
  self:faceDir(startingDir)
  self.i = self.i + 1
  return true
end

function Miner:loop()
  self:turnDir(0.25)

  local shouldRun = true
  while shouldRun do
    local i = self.i - 1
    if i > 1 and (i % self.stripWidth == 0 or i % self.stripWidth == 1) then
      local delta = -0.25 -- left
      if i % (self.stripWidth * 2) == 0 or i % (self.stripWidth * 2) == 1 then
        delta = 0.25 -- right
      end

      self:turnDir(delta)
    end

    shouldRun = self:mineOres(16)
  end
end

local miner = Miner.new()
miner:loop()