import os
import shutil
import re
import zipfile
import requests
from stuff.general import General
from tools.helper import bcolors, download_file, host, print_color, run, get_download_dir

class Magisk(General):
    stable_json = "https://raw.githubusercontent.com/os-fork/HuskyDG-magisk-files/main/stable.json"
    machine = host()

    def __init__(self):
        self.download_loc = get_download_dir()
        self.dl_file_name = os.path.join(self.download_loc, "magisk.apk")
        self.release_meta_file = os.path.join(self.download_loc, "magisk.version")
        self.extract_to = os.path.join(self.download_loc, "magisk_unpack")
        self.copy_dir = "./magisk"
        self.magisk_dir = os.path.join(self.copy_dir, "vendor", "etc", "magisk")
        self.rc_dir = os.path.join(self.copy_dir, "vendor", "etc", "init")
        self.app_package = os.environ.get("REDROID_MAGISK_APP_PACKAGE", "io.github.huskydg.magisk")
        self.custom_apk = os.environ.get("REDROID_MAGISK_APK", "").strip()
        self.custom_url = os.environ.get("REDROID_MAGISK_URL", "").strip()
        self.boot_complete_enabled = os.environ.get("REDROID_MAGISK_BOOT_COMPLETE", "0").strip().lower() in {
            "1",
            "true",
            "yes",
            "on",
        }
        self.preseed_sql = os.environ.get(
            "REDROID_MAGISK_PRESEED_SQL",
            "REPLACE INTO settings (key,value) VALUES "
            "('root_access',3),('su_biometric',0),('multiuser_mode',0),('mnt_ns',1),"
            "('denylist',0),('sulist',0),('zygisk',1); "
            "REPLACE INTO policies (uid,policy,until,logging,notification) VALUES "
            "(2000,2,0,0,0);",
        ).strip()
        self.release_version = "unknown"
        self.magisk_tmp = "/debug_ramdisk"
        self.magisk_bootstrap = "magisk64" if self.machine[1] == 64 else "magisk32"
        self.magisk_system_dir = "/vendor/etc/magisk"
        boot_complete_component = ""
        if self.boot_complete_enabled:
            boot_complete_component = """
    exec u:r:su:s0 root root -- {MAGISKTMP}/magisk --boot-complete""".format(
                MAGISKTMP=self.magisk_tmp,
            )
        preseed_component = ""
        if self.preseed_sql:
            escaped_sql = self.preseed_sql.replace("\\", "\\\\").replace('"', '\\"')
            preseed_component = """
    exec u:r:su:s0 root root -- {MAGISKTMP}/magisk --sqlite "{PRESEED_SQL}" """.format(
                MAGISKTMP=self.magisk_tmp,
                PRESEED_SQL=escaped_sql,
            )
        self.bootanim_component = """
on post-fs-data
    start logd
    mount tmpfs tmpfs {MAGISKTMP} mode=1777
    mkdir /data/adb 0700
    mkdir /data/adb/magisk 0755
    copy {MAGISKSYSTEMDIR}/busybox /data/adb/magisk/busybox
    copy {MAGISKSYSTEMDIR}/magiskboot /data/adb/magisk/magiskboot
    copy {MAGISKSYSTEMDIR}/magiskinit /data/adb/magisk/magiskinit
    copy {MAGISKSYSTEMDIR}/magiskpolicy /data/adb/magisk/magiskpolicy
    copy {MAGISKSYSTEMDIR}/util_functions.sh /data/adb/magisk/util_functions.sh
    copy {MAGISKSYSTEMDIR}/boot_patch.sh /data/adb/magisk/boot_patch.sh
    copy {MAGISKSYSTEMDIR}/addon.d.sh /data/adb/magisk/addon.d.sh
    chmod 0755 /data/adb/magisk/busybox
    chmod 0755 /data/adb/magisk/magiskboot
    chmod 0755 /data/adb/magisk/magiskinit
    chmod 0755 /data/adb/magisk/magiskpolicy
    chmod 0644 /data/adb/magisk/util_functions.sh
    chmod 0644 /data/adb/magisk/boot_patch.sh
    chmod 0644 /data/adb/magisk/addon.d.sh
    copy {MAGISKSYSTEMDIR}/magisk32 {MAGISKTMP}/magisk32
    copy {MAGISKSYSTEMDIR}/magisk64 {MAGISKTMP}/magisk64
    copy {MAGISKSYSTEMDIR}/magiskpolicy {MAGISKTMP}/magiskpolicy
    copy {MAGISKSYSTEMDIR}/stub.apk {MAGISKTMP}/stub.apk
    chmod 0755 {MAGISKTMP}/magisk32
    chmod 0755 {MAGISKTMP}/magisk64
    chmod 0755 {MAGISKTMP}/magiskpolicy
    symlink {MAGISKTMP}/{magisk_bootstrap} {MAGISKTMP}/magisk
    symlink {MAGISKTMP}/magisk {MAGISKTMP}/su
    symlink {MAGISKTMP}/magisk {MAGISKTMP}/resetprop
    symlink {MAGISKTMP}/magiskpolicy {MAGISKTMP}/supolicy
    rm /vendor/bin/su
    rm /vendor/bin/magisk
    rm /vendor/bin/resetprop
    rm /vendor/bin/magiskpolicy
    symlink {MAGISKTMP}/magisk /vendor/bin/su
    symlink {MAGISKTMP}/magisk /vendor/bin/magisk
    symlink {MAGISKTMP}/magisk /vendor/bin/resetprop
    symlink {MAGISKTMP}/magiskpolicy /vendor/bin/magiskpolicy
    mkdir {MAGISKTMP}/.magisk 0711
    mkdir {MAGISKTMP}/.magisk/device 0711
    mkdir {MAGISKTMP}/.magisk/mirror 0755
    mkdir {MAGISKTMP}/.magisk/block 0755
    mkdir {MAGISKTMP}/.magisk/worker 0755
    mkdir {MAGISKTMP}/.magisk/rootdir 0755
    copy {MAGISKSYSTEMDIR}/config {MAGISKTMP}/.magisk/config
    exec u:r:su:s0 root root -- {MAGISKTMP}/magisk --post-fs-data

on nonencrypted
    exec u:r:su:s0 root root -- {MAGISKTMP}/magisk --service

on property:vold.decrypt=trigger_restart_framework
    exec u:r:su:s0 root root -- {MAGISKTMP}/magisk --service

on property:sys.boot_completed=1
    mkdir /data/adb/magisk 755
{PRESEED}
{BOOTCOMPLETE}
    exec -- /system/bin/sh -c "if [ ! -e /data/data/{app_package} ] ; then pm install {MAGISKSYSTEMDIR}/magisk.apk ; fi"
   
on property:init.svc.zygote=restarting
    exec u:r:su:s0 root root -- {MAGISKTMP}/magisk --zygote-restart
   
on property:init.svc.zygote=stopped
    exec u:r:su:s0 root root -- {MAGISKTMP}/magisk --zygote-restart
        """.format(
            MAGISKSYSTEMDIR=self.magisk_system_dir,
            MAGISKTMP=self.magisk_tmp,
            PRESEED=preseed_component,
            BOOTCOMPLETE=boot_complete_component,
            magisk_bootstrap=self.magisk_bootstrap,
            app_package=self.app_package,
        )

    def _cached_version(self):
        if os.path.isfile(self.release_meta_file):
            with open(self.release_meta_file, "r", encoding="utf-8") as meta_file:
                return meta_file.read().strip()
        return ""

    def _write_version(self):
        with open(self.release_meta_file, "w", encoding="utf-8") as meta_file:
            meta_file.write(self.release_version)

    def _validate_archive(self):
        if not zipfile.is_zipfile(self.dl_file_name):
            raise ValueError(
                "Downloaded Magisk/Kitsune archive is not a valid APK/ZIP. "
                "Set REDROID_MAGISK_APK to a local APK or REDROID_MAGISK_URL to a working direct APK URL."
            )

    def resolve_release(self):
        print_color("Resolving Kitsune Mask stable release ...", bcolors.GREEN)
        response = requests.get(self.stable_json, timeout=30)
        response.raise_for_status()
        payload = response.json()
        magisk_info = payload["magisk"]
        self.dl_link = magisk_info["link"]
        self.release_version = magisk_info.get("version", "unknown")

    def download(self):
        if self.custom_apk:
            self.release_version = "local:" + os.path.basename(self.custom_apk)
            if (
                os.path.isfile(self.dl_file_name)
                and self._cached_version() == self.release_version
            ):
                print_color(
                    "Using cached local Magisk/Kitsune APK {} ...".format(self.release_version),
                    bcolors.GREEN,
                )
            else:
                print_color(
                    "Copying local Magisk/Kitsune APK {} ...".format(self.custom_apk),
                    bcolors.GREEN,
                )
                shutil.copyfile(self.custom_apk, self.dl_file_name)
                self._write_version()
            self._validate_archive()
            return

        if self.custom_url:
            self.dl_link = self.custom_url
            self.release_version = "url:" + os.path.basename(self.custom_url)
        else:
            self.resolve_release()

        if (
            os.path.isfile(self.dl_file_name)
            and self._cached_version() == self.release_version
        ):
            print_color(
                "Using cached Kitsune Mask {} ...".format(self.release_version),
                bcolors.GREEN,
            )
            self._validate_archive()
            return

        print_color(
            "Downloading Kitsune Mask {} now .....".format(self.release_version),
            bcolors.GREEN,
        )
        if os.path.isfile(self.dl_file_name):
            os.remove(self.dl_file_name)
        download_file(self.dl_link, self.dl_file_name)
        self._validate_archive()
        self._write_version()

    def extract(self):
        print_color("Extracting archive...", bcolors.GREEN)
        if os.path.exists(self.extract_to):
            shutil.rmtree(self.extract_to)
        os.makedirs(self.extract_to, exist_ok=True)
        with zipfile.ZipFile(self.dl_file_name) as zip_handle:
            zip_handle.extractall(self.extract_to)

    def copy(self):
        if os.path.exists(self.copy_dir):
            shutil.rmtree(self.copy_dir)
        os.makedirs(self.magisk_dir, exist_ok=True)
        os.makedirs(self.rc_dir, exist_ok=True)
        os.chmod(os.path.join(self.copy_dir, "vendor"), 0o755)
        os.chmod(os.path.join(self.copy_dir, "vendor", "etc"), 0o755)
        os.chmod(self.magisk_dir, 0o755)
        os.chmod(self.rc_dir, 0o755)

        print_color("Copying magisk libs now ...", bcolors.GREEN)
        
        arch_map = {
            "x86": "x86",
            "x86_64": "x86_64",
            "arm": "armeabi-v7a",
            "arm64": "arm64-v8a"
        }
        lib_dir = os.path.join(self.extract_to, "lib", arch_map[self.machine[0]])
        for parent, dirnames, filenames in os.walk(lib_dir):
            for filename in filenames:
                o_path = os.path.join(lib_dir, filename)  
                filename = re.search(r'lib(.*)\.so', filename)
                n_path = os.path.join(self.magisk_dir, filename.group(1))
                shutil.copyfile(o_path, n_path)
                os.chmod(n_path, 0o755)

        if self.machine[1] == 64:
            secondary_arch_map = {
                "x86_64": "x86",
                "arm64": "armeabi-v7a",
            }
            secondary_dir = os.path.join(
                self.extract_to,
                "lib",
                secondary_arch_map[self.machine[0]],
            )
            magisk32_src = os.path.join(secondary_dir, "libmagisk32.so")
            if os.path.exists(magisk32_src):
                magisk32_dst = os.path.join(self.magisk_dir, "magisk32")
                shutil.copyfile(magisk32_src, magisk32_dst)
                os.chmod(magisk32_dst, 0o755)

        assets_dir = os.path.join(self.extract_to, "assets")
        for asset_name in ["util_functions.sh", "addon.d.sh", "boot_patch.sh", "stub.apk"]:
            asset_src = os.path.join(assets_dir, asset_name)
            if os.path.exists(asset_src):
                asset_dst = os.path.join(self.magisk_dir, asset_name)
                shutil.copyfile(asset_src, asset_dst)
                os.chmod(asset_dst, 0o644)

        config_path = os.path.join(self.magisk_dir, "config")
        with open(config_path, "w", encoding="utf-8") as config_file:
            config_file.write("SYSTEMMODE=true\nRECOVERYMODE=false\n")
        os.chmod(config_path, 0o600)

        magisk_apk_path = os.path.join(self.magisk_dir, "magisk.apk")
        shutil.copyfile(self.dl_file_name, magisk_apk_path)
        os.chmod(magisk_apk_path, 0o644)

        magisk_rc_path = os.path.join(self.rc_dir, "magisk.rc")
        with open(magisk_rc_path, "w", encoding="utf-8") as initfile:
            initfile.write(self.bootanim_component)
        os.chmod(magisk_rc_path, 0o644)