#!/usr/bin/env python3
import atexit
import hashlib
import os.path
import re
import shutil
import subprocess
import urllib.request
import zipfile


def calculate_md5(file_path):
    """计算文件的 MD5 值"""
    hash_md5 = hashlib.md5()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def download_file(url, file_path):
    """下载文件"""
    print(f"==> Downloading from {url}")
    urllib.request.urlretrieve(url, file_path)


def main():
    current_path = os.path.dirname(os.path.abspath(__file__))
    apk_path = os.path.join(current_path, "magisk.apk")
    unzip_path = os.path.join(current_path, "temp")
    overlay_path = os.path.join(current_path, "rootfs")
    overlay_magisk_path = os.path.join(overlay_path, "vendor", "etc", "init", "magisk")
    
    # Magisk 下载配置
    magisk_url = "https://github.com/topjohnwu/Magisk/releases/download/v30.2/Magisk-v30.2.apk"
    expected_md5 = "2691c30ccf059af2536cb0af803c787c"
    
    # 检查并下载 Magisk APK
    need_download = True
    if os.path.exists(apk_path):
        print("==> Checking existing Magisk APK...")
        actual_md5 = calculate_md5(apk_path)
        if actual_md5 == expected_md5:
            print("==> Magisk APK already exists and verified!")
            need_download = False
        else:
            print(f"==> MD5 mismatch. Expected: {expected_md5}, Got: {actual_md5}")
    
    if need_download:
        print("==> Downloading Magisk v30.2...")
        download_file(magisk_url, apk_path)
        actual_md5 = calculate_md5(apk_path)
        if actual_md5 != expected_md5:
            raise Exception(f"Downloaded file MD5 verification failed! Expected: {expected_md5}, Got: {actual_md5}")
        print("==> Download completed and verified!")

    shutil.rmtree(unzip_path, ignore_errors=True)
    os.makedirs(unzip_path, exist_ok=True)
    atexit.register(shutil.rmtree, unzip_path, ignore_errors=True)

    print("==> Extracting archive...")
    with zipfile.ZipFile(apk_path) as z:
        z.extractall(unzip_path)

    shutil.rmtree(overlay_path, ignore_errors=True)
    os.makedirs(overlay_magisk_path, exist_ok=True)

    print("==> Installing magisk now ...")
    # Extract arm64 binaries
    lib64_path = os.path.join(unzip_path, "lib", "arm64-v8a")
    for parent, dirnames, filenames in os.walk(lib64_path):
        for filename in filenames:
            so_path = os.path.join(lib64_path, filename)
            so_name = re.search(r"lib(.*)\.so", filename)
            if so_name:
                target_path = os.path.join(overlay_magisk_path, so_name.group(1))
                shutil.copyfile(so_path, target_path)
                subprocess.check_call(["chmod", "+x", target_path])

    # Extract arm32 magisk binary
    lib32_path = os.path.join(unzip_path, "lib", "armeabi-v7a")
    magisk32_src = os.path.join(lib32_path, "libmagisk32.so")
    magisk32_dst = os.path.join(overlay_magisk_path, "magisk32")
    if os.path.exists(magisk32_src):
        shutil.copyfile(magisk32_src, magisk32_dst)
        subprocess.check_call(["chmod", "+x", magisk32_dst])

    # Copy magisk.apk to target directory
    apk_dst = os.path.join(overlay_magisk_path, "magisk.apk")
    if os.path.exists(apk_path):
        shutil.copyfile(apk_path, apk_dst)
    
    print("==> Magisk installation completed!")


if __name__ == '__main__':
    main()