#!/usr/bin/env python

import itertools
import os
import pathlib
import re
import shlex
import subprocess
import time

file_ver1 = "./Brave-stable-v1.73.104-x86_64.AppImage"
file_ver2 = "./Brave-stable-v1.73.105-x86_64.AppImage"

block_sizes = ["128k", "512k", "1M"]
comps_squashfs = {
    "gzip": {"-Xcompression-level": ["1", "6", "9"]},
    "lz4": {},
    "xz": {
        "-Xdict-size": ["100%"],
    },
    "zstd": {
        "-Xcompression-level": ["1", "15", "22"],
    },
}

output = open("results.csv", "w")
output.write(
    "Compression Method,Compression Parameters,Block Size,Compression Time,"
    "FS Size,Startup Time (average of 5 runs),zsync download size"
)
output.flush()

try:
    os.mkdir("zsyncout")
except FileExistsError:
    pass

if not os.path.exists("ver1.AppDir"):
    subprocess.Popen([file_ver1, "--appimage-extract"]).wait()
    os.rename("squashfs-root", "ver1.AppDir")
if not os.path.exists("ver2.AppDir"):
    subprocess.Popen([file_ver2, "--appimage-extract"]).wait()
    os.rename("squashfs-root", "ver2.AppDir")

with open("uruntime-appimage-x86_64", "rb") as f:
    runtime = f.read()

for comp in comps_squashfs:
    # product all possible combinations of compression parameters
    comp_params = list(
        itertools.product(*list(list([i, j] for j in comps_squashfs[comp][i]) for i in comps_squashfs[comp]))
    )
    for comp_param in comp_params:
        for block_size in block_sizes:
            output.write("\n")
            output.write(f"SquashFS {comp},{shlex.join(sum(comp_param, []))},{block_size},")
            output.flush()

            print(
                "Preparing AppImages for compression method %s with parameters %s and block size %s"
                % (comp, sum(comp_param, []), block_size)
            )
            # Prepare the AppImages
            command = ["mksquashfs", "ver1.AppDir", "file_ver1.squashfs", "-root-owned", "-noappend", "-comp", comp]
            command.extend(["-b", block_size])
            for i in comp_param:
                command.extend(i)
            # file_ver1 is only the base file, so do not measure efficiency
            p = subprocess.Popen(command)
            if p.wait() != 0:
                print("Error in compression")
                continue
            with open("file_ver1.squashfs", "rb") as f:
                squashfs = f.read()
            os.remove("file_ver1.squashfs")
            with open("file_ver1.AppImage", "wb") as f:
                f.write(runtime)
                f.write(squashfs)
            del squashfs
            try:
                os.chmod("file_ver1.AppImage", 0o755)
            except PermissionError:
                pass

            # Compression time
            command[1] = "ver2.AppDir"
            command[2] = "file_ver2.squashfs"
            start = time.time()
            p = subprocess.Popen(command)
            if p.wait() != 0:
                print("Error in compression")
                continue
            output.write(f"{(time.time() - start) * 1000:.1f}ms,")
            output.flush()
            with open("file_ver2.squashfs", "rb") as f:
                squashfs = f.read()
            os.remove("file_ver2.squashfs")
            with open("file_ver2.AppImage", "wb") as f:
                f.write(runtime)
                f.write(squashfs)
            output.write(f"{len(squashfs)},")
            output.flush()
            print("SquashFS size: %d" % len(squashfs))
            del squashfs
            try:
                os.chmod("file_ver2.AppImage", 0o755)
            except PermissionError:
                pass

            # Startup time
            print("Measuring startup time, 10 rounds, only last 5 rounds count")
            for _ in range(2):
                startup_times = []
                for i in range(5):
                    start = time.time()
                    p = subprocess.Popen(["./file_ver2.AppImage"])
                    p.wait()
                    used_time = time.time() - start
                    startup_times.append(used_time)
                    print("Startup time round %d: %f" % (i + 1, used_time))
            output.write(f"{sum(startup_times) * 1000 / 5:.1f} ms,")
            output.flush()

            # zsync download size
            print("Creating zsync file")
            p = subprocess.Popen(["zsyncmake", "file_ver2.AppImage"])
            if p.wait() != 0:
                print("Error in zsync")
                continue
            print("Calculating zsync download size")
            p = subprocess.Popen(
                ["zsync", "-i", "file_ver1.AppImage", "file_ver2.AppImage.zsync", "-o", "./zsyncout"],
                stderr=subprocess.PIPE,
            )
            zsync_output = p.stderr.read().decode("utf-8")
            percentage = re.search(r"(\d+\.\d+)%", zsync_output).group(1)
            p.wait()
            try:
                os.remove("zsyncout.part")
            except FileNotFoundError:
                pass
            output.write(f"{(1 - float(percentage) / 100) * os.path.getsize('file_ver2.AppImage'):.0f}")
            output.flush()

output.close()