#!/usr/bin/env python3

import os
import shutil
import sys


def write_files(source, target_dir, template, num_copies):
    master_path = os.path.join(target_dir, template.format(0))

    print("Writing {}".format(master_path))
    with open(master_path, "wb") as out_stream:
        out_stream.write(source)

    for n in range(1, num_copies):
        copy_path = os.path.join(target_dir, template.format(n))
        print("Copying to {}".format(copy_path))
        shutil.copy(master_path, copy_path)


BIG_SIZE = 30 * 1024 * 1024 + 5 # 30 MB + 5 bytes
SMALL_SIZE = 3 * 1024 * 1024 + 5 # 3 MB + 5 bytes


def main():
    target_dir = sys.argv[1]
    os.makedirs(target_dir, exist_ok=True)

    buf = bytearray(b"A" * BIG_SIZE)
    write_files(buf, target_dir, "all_a_{}", 5)

    # change a value in the middle to ensure dupes aren't re-unified when
    # identical blocks follow differing blocks
    buf[BIG_SIZE >> 1] ^= 0x0F
    write_files(buf, target_dir, "diff_middle_{}", 5)

    # have one file that differs earlier, splitting it off into its own group
    # of 1 (and therefore not a dupe set)
    buf[BIG_SIZE >> 2] ^= 0x0F
    write_files(buf, target_dir, "diff_singleton_{}", 1)
    buf[BIG_SIZE >> 2] ^= 0x0F
    buf[BIG_SIZE >> 1] ^= 0x0F

    # now at the end. this will test the algorithm to keep the files open
    # as long as possible
    buf[-1] ^= 0x0F
    write_files(buf, target_dir, "diff_end_{}", 5)
    buf[-1] ^= 0x0F

    # now at the start
    buf[0] ^= 0x0F
    write_files(buf, target_dir, "diff_start_{}", 5)
    buf[0] ^= 0x0F

    # sanity (?) something the same size but all different
    buf = bytearray(b"B" * BIG_SIZE)
    write_files(buf, target_dir, "all_b_{}", 5)

    # smaller files, but stress the number that we have open at once
    buf = bytearray(b"A" * SMALL_SIZE)
    write_files(buf, target_dir, "open_stress_a_{}", 320)

    buf[-512] ^= 0x0F
    write_files(buf, target_dir, "open_stress_b_{}", 260)

    # some zero files
    write_files(b"", target_dir, "zero_{}", 5)

    # some 1-byte files
    write_files(b"1", target_dir, "one_{}", 2)


if __name__ == "__main__":
    main()

@@ -109,24 +109,67 @@
 um_copies):%0A
+    os.makedirs(target_dir, exist_ok=True)%0A
     master_p
@@ -659,22 +659,20 @@
 ():%0A    
-target
+base
 _dir = s
@@ -690,45 +690,94 @@
     
-os.makedirs(target_dir, exist_ok=True
+root_a = os.path.join(base_dir, %22root_a%22)%0A    root_b = os.path.join(base_dir, %22root_b%22
 )%0A%0A 
@@ -829,34 +829,30 @@
 _files(buf, 
-target_dir
+root_a
 , %22all_a_%7B%7D%22
@@ -1026,34 +1026,30 @@
 _files(buf, 
-target_dir
+root_a
 , %22diff_midd
@@ -1228,34 +1228,30 @@
 _files(buf, 
-target_dir
+root_a
 , %22diff_sing
@@ -1464,34 +1464,30 @@
 _files(buf, 
-target_dir
+root_a
 , %22diff_end_
@@ -1578,26 +1578,22 @@
 es(buf, 
-target_dir
+root_a
 , %22diff_
@@ -1743,26 +1743,22 @@
 es(buf, 
-target_dir
+root_a
 , %22all_b
@@ -1892,34 +1892,30 @@
 _files(buf, 
-target_dir
+root_a
 , %22open_stre
@@ -1973,26 +1973,22 @@
 es(buf, 
-target_dir
+root_a
 , %22open_
@@ -2050,26 +2050,22 @@
 es(b%22%22, 
-target_dir
+root_a
 , %22zero_
@@ -2123,31 +2123,195 @@
 1%22, 
-target_dir, %22one
+root_a, %22one_%7B%7D%22, 2)%0A    %0A    # test only-across-roots option%0A    buf = bytearray(b%22R%22 * SMALL_SIZE)%0A    write_files(buf, root_a, %22oar_a_%7B%7D%22, 0)%0A    write_files(buf, root_b, %22oar_b
 _%7B%7D%22, 
-2
+0
 )%0A%0A%0A

test --only-across-roots