#!/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, "magisk")
    overlay_magisk_path = os.path.join(overlay_path, "system", "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 ...")
    
    # 架构映射，与 magisk.py 保持一致
    arch_map = {
        "x86": "x86",
        "x86_64": "x86_64", 
        "arm": "armeabi-v7a",
        "arm64": "arm64-v8a"
    }
    
    # 默认使用 arm64 架构
    machine_arch = "arm64"
    lib_dir = os.path.join(unzip_path, "lib", arch_map[machine_arch])
    
    # 复制主要架构的所有库文件
    if os.path.exists(lib_dir):
        for parent, dirnames, filenames in os.walk(lib_dir):
            for filename in filenames:
                o_path = os.path.join(lib_dir, filename)  
                so_name = re.search(r'lib(.*)\.so', filename)
                if so_name:
                    n_path = os.path.join(overlay_magisk_path, so_name.group(1))
                    shutil.copyfile(o_path, n_path)
                    subprocess.check_call(["chmod", "+x", n_path])

    # 同时复制 arm32 的 magisk 二进制文件（如果存在）
    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()