
local function printUsage()
	print( "Usages:" )
	print( "label get" )
	print( "label get <drive>" )
	print( "label set <text>" )
	print( "label set <drive> <text>" )
	print( "label clear" )
	print( "label clear <drive>" )
end

local function checkDrive( sDrive )
    if peripheral.getType( sDrive ) == "drive" then
        -- Check the disk exists
        local bData = disk.hasData( sDrive )
        if not bData then
            print( "No disk in "..sDrive.." drive" )
            return false
        end
    else
	    print( "No disk drive named "..sDrive )
        return false
	end
	return true
end

local function get( sDrive )
	if sDrive ~= nil then
	    if checkDrive( sDrive ) then
            local sLabel = disk.getLabel( sDrive )
            if sLabel then
                print( "Disk label is \""..sLabel.."\"" )
            else
                print( "No Disk label" )
            end
        end
	else
		local sLabel = os.getComputerLabel()
		if sLabel then
			print( "Computer label is \""..sLabel.."\"" )
		else
			print( "No Computer label" )
		end
	end
end

local function set( sDrive, sText )
	if sDrive ~= nil then
	    if checkDrive( sDrive ) then
            disk.setLabel( sDrive, sText )
            local sLabel = disk.getLabel( sDrive )
            if sLabel then
                print( "Disk label set to \""..sLabel.."\"" )
            else
                print( "Disk label cleared" )
            end
        end
	else
		os.setComputerLabel( sText )
		local sLabel = os.getComputerLabel()
		if sLabel then
			print( "Computer label set to \""..sLabel.."\"" )
		else
			print( "Computer label cleared" )
		end
	end
end

local tArgs = { ... }
local sCommand = tArgs[1]
if sCommand == "get" then
    -- Get a label
    if #tArgs == 1 then
        get( nil )
    elseif #tArgs == 2 then
        get( tArgs[2] )
    else
        printUsage()
    end
elseif sCommand == "set" then
    -- Set a label
    if #tArgs == 2 then
        set( nil, tArgs[2] )
    elseif #tArgs == 3 then
        set( tArgs[2], tArgs[3] )
    else
        printUsage()
    end
elseif sCommand == "clear" then
    -- Clear a label
    if #tArgs == 1 then
        set( nil, nil )
    elseif #tArgs == 2 then
        set( tArgs[2], nil )
    else
        printUsage()
    end
else
    printUsage()
end
