local args = { ... }

local function getDrive(side)
    if side ~= nil then
        if peripheral.isPresent(side) and peripheral.getType(side) == "rp_drive" then
            return peripheral.wrap(side)
        end
    end
    return nil
end
local drive = getDrive(args[2])

local function printWithColors(...)
    local args = { ... }
    for i=1,#args,2 do
        if term.isColor() then
            term.setTextColor(args[i])
        end
        write(args[i+1])
    end
    if term.isColor() then
        term.setTextColor(colors.white)
    end
    print("")
end
local serNoColors = {
    colors.orange,
    colors.magenta,
    colors.lightBlue,
    colors.yellow,
    colors.pink,
    colors.cyan,
    colors.purple,
    colors.green
}

if args[1] == "read" then
    if drive == nil then
        error("drive not detected on side " .. (args[2] or "nil"))
    end
    if args[3] == nil then
        error("destination filename not provided")
    end
    if not drive.isReadable() then
        error("source disk is not readable")
    end
    local file = fs.open(args[3], "wb")
    if file == nil then
        error("could not open destination file")
    end
    write("Reading")
    for i=1,drive.getSectorCount() do
        drive.readSector(i)
        for j=1,drive.getBytesPerSector() do
            file.write(drive.getBufferByte(j))
        end
        write(".")
    end
    file.close()
    print(" OK!")
elseif args[1] == "write" then
    if drive == nil then
        error("drive not detected on side " .. (args[2] or "nil"))
    end
    if args[3] == nil then
        error("source filename not provided")
    end
    if not drive.isWritable() then
        error("destination disk is not writable")
    end
    local file = fs.open(args[3], "rb")
    if file == nil then
        error("could not open source file")
    end
    write("Writing")
    local isWriting = true
    local currentSector = 0
    while isWriting do
        local lastJ = 0
        for j=1,drive.getBytesPerSector() do
            local b = file.read()
            if b == nil then
                isWriting = false
                drive.setBufferByte(j, 0)
            else
                lastJ = j
                drive.setBufferByte(j, b)
            end
        end
        if lastJ > 0 then
            currentSector = currentSector + 1
            if currentSector > drive.getSectorSize() then
                error("file too big")
            end
            drive.writeSector(currentSector)
            write(".")
        end
    end
    file.close()
    print(" OK!")
elseif args[1] == "label" then
    if drive == nil then
        error("drive not detected on side " .. (args[2] or "nil"))
    end
    if #args < 3 then
        local oldLabel = drive.getLabel()
        if oldLabel == nil or #oldLabel <= 0 then
            error("label not provided")
        end
        -- Clear label
        if drive.setLabel("") then
            print("Label removed successfully.")
        else
            print("Label removal failed!")
        end
    else
        -- Set label
        local label = args[3]
        for i=4,#args do
            label = label .. " " .. args[i]
        end
        if drive.setLabel(label) then
            print("Label changed successfully.")
        else
            print("Label change failed!")
        end
    end
elseif args[1] == "copy" then
    if drive == nil then
        error("drive not detected on side " .. (args[2] or "nil"))
    end
    local destDrive = getDrive(args[3])
    if destDrive == nil then
        error("drive not detected on side " .. (args[3] or "nil"))
    end
    if not drive.isReadable() then
        error("source disk is not readable")
    end
    if not destDrive.isWritable() then
        error("destination disk is not writable")
    end
    if drive.getSectorCount() > destDrive.getSectorSize() then
        error("destination disk too small")
    end
    write("Copying")
    for i=1,drive.getSectorCount() do
        drive.readSector(i)
        for j=1,drive.getBytesPerSector() do
            destDrive.setBufferByte(j, drive.getBufferByte(j))
        end
        destDrive.writeSector(i)
        write(".")
    end
    print(" OK!")
elseif args[1] == "erase" then
    if drive == nil then
        error("drive not detected on side " .. (args[2] or "nil"))
    end
    if drive.erase() then
        print("Disk erase successful.")
    else
        print("Disk erase failed!")
    end
elseif args[1] == "about" then
    printWithColors(colors.red, "r", colors.yellow, "p", colors.lightBlue, "c", colors.white, "/", colors.lightGray, "drive", colors.white, " 0.1.0")
elseif args[1] == "info" then
    if drive == nil then
        error("drive not detected on side " .. (args[2] or "nil"))
    end
    if drive.isPresent() then
        local line1 = ""
        local line2 = ""
        local line1Color = colors.white
        local line2Color = colors.white
        local serNoColor = colors.white
        local line3 = ""
        local line4 = ""
        local drawSerNo = false
        local drawNotch = false

        if drive.getSerialNumber() ~= nil then
            local serialNumber = drive.getSerialNumber()
            drawSerNo = true

            line1 = "#" .. serialNumber
            local serNoValue = 0
            for i=1,#serialNumber do
                serNoValue = serNoValue + string.byte(serialNumber, i, i)
            end
            serNoColor = serNoColors[serNoValue % #serNoColors]
        else
            line1 = "Empty"
        end
        local label = drive.getLabel()
        if label ~= nil and #label > 0 then
            line1 = label .. " (" .. line1 .. ")"
        end

        if drive.isWritable() then
            drawNotch = true
            if drive.isReadable() then
                line2 = "Read/Write"
            else
                line2 = "Write-only"
            end
            line2Color = colors.lime
        elseif drive.isReadable() then
            line2 = "Read-only"
            line2Color = colors.red
        end

        local sectorSize = drive.getSectorSize()
        if sectorSize > 0 then
            local sizeKbStr = (sectorSize / 8)
            local sizeStr = sectorSize
            if drive.isWritable() then
                local sectorCount = drive.getSectorCount()
                sizeKbStr = (sectorCount / 8) .. "/" .. sizeKbStr
                sizeStr = sectorCount .. "/" .. sizeStr

                line3 = sizeStr .. " sectors written"
                line4 = "(" .. sizeKbStr .. " KBytes)"
            else
                line3 = sizeStr .. " sectors"
                line4 = "(" .. sizeKbStr .. " KBytes)"
            end
        end

        printWithColors(colors.lightGray, " _________")
        if drawSerNo and drawNotch then
            printWithColors(colors.lightGray, "| ", serNoColor, ".......", colors.lightGray, " |  ", line1Color, line1)
            printWithColors(colors.lightGray, "| ", serNoColor, "'''''''", colors.lightGray, "[   ", line2Color, line2)
        elseif drawSerNo and (not drawNotch) then
            printWithColors(colors.lightGray, "| ", serNoColor, ".......", colors.lightGray, " |  ", line1Color, line1)
            printWithColors(colors.lightGray, "| ", serNoColor, "'''''''", colors.lightGray, " |  ", line2Color, line2)
        elseif (not drawSerNo) and drawNotch then
            printWithColors(colors.lightGray, "|         |  ", line1Color, line1)
            printWithColors(colors.lightGray, "|        [   ", line2Color, line2)
        else
            printWithColors(colors.lightGray, "|         |  ", line1Color, line1)
            printWithColors(colors.lightGray, "|        [   ", line2Color, line2)
        end
        printWithColors(colors.lightGray, "|   / \\   |  ")
        printWithColors(colors.lightGray, "|   \\ /   |  ", colors.white, line3)
        printWithColors(colors.lightGray, "|    _    |  ", colors.lightGray, line4)
        printWithColors(colors.lightGray, "|___| |___|")
    else
        print("No disk detected.")
    end
else
    print("Usage:")
    print("  rpc/drive info [side]")
    printWithColors(colors.lightGray, "  Query drive status/disk information.")
    print("  rpc/drive label [side] [name]")
    printWithColors(colors.lightGray, "  Change disk label.")
    print("  rpc/drive read [side] [filename] ")
    printWithColors(colors.lightGray, "  Read disk contents to file.")
    print("  rpc/drive write [side] [filename]")
    printWithColors(colors.lightGray, "  Write file contents to disk.")
    print("  rpc/drive copy [side] [side]")
    printWithColors(colors.lightGray, "  Write disk contents to disk.")
    print("  rpc/drive erase [side]")
    printWithColors(colors.lightGray, "  Erase disk contents.")
    print("  rpc/drive about")
    printWithColors(colors.lightGray, "  Show about information.")
end