#!/usr/bin/python3

import sys
import tempfile
import subprocess
import os
import shutil
import struct

def check_squashfs_header(data, offset):
    block_size = struct.unpack_from("<I", data, offset + 12)[0]
    if block_size not in (4096, 8192, 16384, 32768, 65536, 131072, 262144, 524288, 1048576):
        return False
    compression_type = struct.unpack_from("<H", data, offset + 20)[0]
    if (compression_type < 1) or (compression_type > 7):
        return False
    block_log_size = struct.unpack_from("<H", data, offset + 22)[0]
    if block_size != (1 << block_log_size):
        return False
    return True

def get_squashfs_offset(appimage):
    magic_number = b"hsqs"
    with open(appimage, "rb") as f:
        data= f.read()
    offset = 0
    while True:
        offset2 = data.find(magic_number, offset)
        if offset2 == -1:
            raise ValueError("SquashFS magic number not found")
        if check_squashfs_header(data, offset2):
            return offset2
        offset = offset2 + len(magic_number)

def extract_file_from_appimage(appimage, file_path, output_path):
    offset = get_squashfs_offset(appimage)
    command = ["unsquashfs", "-o", str(offset), "-f", "-d", output_path, appimage, file_path]
    subprocess.run(command, check=True)

def get_data(data, offset):
    pos = data.find(b'"', offset)
    pos2 = data.find(b'"', pos + 1)
    if pos == -1 or pos2 == -1:
        raise ValueError("Invalid SVG file")
    return data[pos + 1:pos2]

def convert_svg_to_png(svg_path, png_path, output_size):
    with open(svg_path, "rb") as f:
        data = f.read()
    if b"<svg" not in data:
        raise ValueError("Invalid SVG file")
    pos_width = data.find(b'width="')
    pos_height = data.find(b'height="')
    pos_viewbox = data.find(b'viewBox="')
    if pos_width != -1 and pos_height != -1:
        width = int(float(get_data(data, pos_width)))
        height = int(float(get_data(data, pos_height)))
    else:
        if pos_viewbox == -1:
            raise ValueError("SVG file does not contain width/height or viewBox attributes")
        viewbox_values = get_data(data, pos_viewbox).split()
        if len(viewbox_values) != 4:
            raise ValueError("Invalid viewBox attribute in SVG file")
        width = int(float(viewbox_values[2]))
        height = int(float(viewbox_values[3]))
    if width == height:
        output_width = output_size
        output_height = output_size
    else:
        aspect_ratio = width / height
        if aspect_ratio > 1:
            output_width = output_size
            output_height = int(output_size / aspect_ratio)
        else:
            output_width = int(output_size * aspect_ratio)
            output_height = output_size
    subprocess.run(["inkscape", "--export-type=png", "-w", str(output_width), "-h", str(output_height), "-o", png_path, svg_path], check=True)

def main():
    appimage = sys.argv[1]
    outpath = sys.argv[2]
    output_size = int(sys.argv[3])

    tmp = tempfile.TemporaryDirectory(prefix="appimage-thumbnailer-")
    icon_name = ".DirIcon"

    while True:
        extract_file_from_appimage(appimage, icon_name, tmp.name)
        icon = os.path.join(tmp.name, icon_name)

        # the icon may be a symlink, so we need to resolve it recursively
        if not os.path.islink(icon):
            break
        icon_name = os.readlink(icon)
    if icon.endswith(".svg"):
        convert_svg_to_png(icon, outpath, output_size)
    else:
        shutil.copy(icon, outpath)

if __name__ == "__main__":
    main()
