#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p python3Packages.requests

import base64
import json
import plistlib
import requests
import re
from typing import Dict, Any, List


def update_macos_extra():
    assets = plistlib.loads(
        requests.get(
            "https://mesu.apple.com/assets/com_apple_MobileAsset_Font7/com_apple_MobileAsset_Font7.xml"
        ).content
    )["Assets"]

    families: Dict[str, Any] = dict()
    all_assets: List[Dict[str, str]] = []
    for asset in assets:
        # Period used for extra PUA font
        names = set(
            [
                info["FontFamilyName"]
                for info in asset["FontInfo4"]
                if not info["FontFamilyName"].startswith(".")
            ]
        )
        name = names.pop()

        # There seems to be some support for things other than SHA-1 but I haven't seen it
        hsh = "sha1-" + base64.b64encode(asset["_Measurement"]).decode("utf-8")
        url = asset["__BaseURL"] + asset["__RelativePath"]
        asset_obj = {
            "url": url,
            "hash": hsh,
        }
        all_assets.append(asset_obj)
        version = asset["_CompatibilityVersion"]

        if len(names) > 1:
            print("Multiple family names found in asset:", ", ".join(names))
        if name in families:
            other_version = families[name]["version"]
            if version == other_version:
                families[name]["assets"].append(asset_obj)
                continue
            elif int(version) < int(other_version):
                print(f"Duplicate asset found for family {name}, skipping")
                continue

        families[name] = {
            "assets": [asset_obj],
            "version": version,
        }

    return {"extra": families, "all_extra": {"assets": all_assets, "version": "mixed"}}


def update_macos_ipsw():
    ipsws = plistlib.loads(
        requests.get(
            "https://mesu.apple.com/assets/macos/com_apple_macOSIPSW/com_apple_macOSIPSW.xml"
        ).content
    )["MobileDeviceSoftwareVersionsByVersion"]["1"]["MobileDeviceSoftwareVersions"]

    # Apple seems to be naming everything Mac now, use the newest one
    # As of writing everything is UniversalMac but do this anyway
    pattern = re.compile(r"Mac(\d+),(\d+)")
    max_ident = (0, 0)
    for ident in ipsws.keys():
        if match := pattern.fullmatch(ident):
            major, minor = match.groups()
            major, minor = int(major), int(minor)
            if major > max_ident[0] or major == max_ident[0] and minor > max_ident[1]:
                max_ident = major, minor

    ident = f"Mac{max_ident[0]},{max_ident[1]}"

    # Unknown has an extra level we don't want, ignore it
    build = max(ipsws[ident].keys(), key=lambda item: "" if item == "Unknown" else item)
    ipsw = ipsws[ident][build]["Restore"]
    version = ipsw["ProductVersion"]

    print(f"Using ipsw from macOS {version} (build {build})")
    return {
        "base": {
            "url": ipsw["FirmwareURL"],
            "hash": "sha1-"
            + base64.b64encode(
                base64.b16decode(ipsw["FirmwareSHA1"], casefold=True)
            ).decode("utf-8"),
            "version": version,
            "build": build,
        }
    }


with open("assets.json", "w") as f:
    json.dump(update_macos_ipsw() | update_macos_extra(), f)