local os        = require("os")
local robot     = require("robot")
local component = require("component")
local event     = require("event")
local nav       = require("nav")


local magnet = component.tractor_beam

local depth, width = ...

if not depth or not width then
    print('Usage: lumberjack <depth> <width>')
    print('Operates in a depth x width grid of trees with 4 blocks between them')
    print('Depth is forward, width is right')
    print('Place stone as the 8th block above ground where a tree would grow')
    print('Starting location of the robot is directly facing the first tree')
    return
end

depth = tonumber(depth)
width = tonumber(width)

local dir_forward = true
local dir_right = true

local pos_forward = 1
local pos_right = 1

local running = true

local function mine_tree()
    robot.swing()
    nav.forward()
    robot.swingDown()
    for _ = 1, 5 do
        robot.swingUp()
        nav.up()
    end
    nav.down(5)
    robot.placeDown()
    nav.back()
end

local function move_depth()
    if dir_forward then
        if pos_forward == depth then
            dir_forward = false
        end
    else
        if pos_forward == 1 then
            dir_forward = true
        end
    end

    if dir_forward then
        robot.turnRight()
        nav.forward()
        robot.turnLeft()
        nav.forward(5)
        robot.turnLeft()
        nav.forward()
        robot.turnRight()
        pos_forward = pos_forward + 1
    else
        robot.turnRight()
        nav.forward()
        robot.turnRight()
        nav.forward(5)
        robot.turnRight()
        nav.forward()
        robot.turnRight()
        pos_forward = pos_forward - 1
    end
end

local function collect_items()
    while magnet.suck() do end
end

local function move_width()
    if dir_right then
        if pos_right == width then
            dir_right = false
            move_depth()
        else
            robot.turnRight()
            nav.forward(5)
            robot.turnLeft()
            pos_right = pos_right + 1
        end
    else
        if pos_right == 1 then
            dir_right = true
            move_depth()
        else
            robot.turnLeft()
            nav.forward(5)
            robot.turnRight()
            pos_right = pos_right - 1
        end
    end
end

local function handle_interrupt()
    if event.pull(0.5, "interrupted") then
        print("Interrupted. Returning to original location...")
        robot.setLightColor(0xff0000)

        -- return to original location
        robot.turnLeft()
        while pos_right > 1 do
            nav.forward(5)
            pos_right = pos_right - 1
        end
        nav.back()
        robot.turnRight()
        while pos_forward > 1 do
            nav.back(5)
            pos_forward = pos_forward - 1
        end
        robot.turnLeft()
        nav.forward()
        robot.turnRight()

        running = false
    end
end

robot.setLightColor(0x00ff00)

while running do
    if robot.detect() then
        mine_tree()
    end
    move_width()
    collect_items()
    handle_interrupt()
end