import asyncio
import contextlib
import hashlib
import logging
from typing import Optional
from zipfile import ZipFile

from usb import USBError

from pymobiledevice3.exceptions import PyMobileDevice3Exception
from pymobiledevice3.irecv import IRecv, Mode
from pymobiledevice3.lockdown import create_using_usbmux
from pymobiledevice3.restore.base_restore import BaseRestore, Behavior
from pymobiledevice3.restore.consts import lpol_file
from pymobiledevice3.restore.device import Device
from pymobiledevice3.restore.tss import TSSRequest, TSSResponse

RESTORE_VARIANT_ERASE_INSTALL = "Erase Install (IPSW)"
RESTORE_VARIANT_UPGRADE_INSTALL = "Upgrade Install (IPSW)"
RESTORE_VARIANT_MACOS_RECOVERY_OS = "macOS Customer"


class Recovery(BaseRestore):
    def __init__(self, ipsw: ZipFile, device: Device, tss: Optional[dict] = None, behavior: Behavior = Behavior.Update):
        super().__init__(ipsw, device, tss, behavior)
        self.tss_localpolicy = None
        self.tss_recoveryos_root_ticket = None
        self.restore_boot_args = None

    async def reconnect_irecv(self, is_recovery=None):
        self.logger.debug("waiting for device to reconnect...")
        self.device.set_irecv(IRecv(ecid=await self.device.get_ecid(), is_recovery=is_recovery))
        self.logger.debug(f"connected mode: {self.device.irecv.mode}")

    def get_preboard_manifest(self):
        overrides = {
            "@APTicket": True,
            "ApProductionMode": 0,
            "ApSecurityDomain": 0,
        }

        parameters = {
            "ApProductionMode": False,
            "ApSecurityMode": False,
            "ApSupportsImg4": True,
        }

        self.populate_tss_request_from_manifest(parameters)

        tss = TSSRequest()
        tss.add_common_tags(parameters, overrides)

        parameters["_OnlyFWComponents"] = True

        tss.add_ap_tags(parameters)

        return tss.img4_create_local_manifest(build_identity=self.build_identity)

    async def get_tss_response(self) -> TSSResponse:
        # populate parameters
        parameters = {}

        parameters["ApECID"] = await self.device.get_ecid()
        ap_nonce = await self.device.get_ap_nonce()
        if ap_nonce is not None:
            parameters["ApNonce"] = ap_nonce

        sep_nonce = await self.device.get_sep_nonce()
        if sep_nonce is not None:
            parameters["ApSepNonce"] = sep_nonce

        parameters["ApProductionMode"] = True

        if await self.device.get_is_image4_supported():
            parameters["ApSecurityMode"] = True
            parameters["ApSupportsImg4"] = True
        else:
            parameters["ApSupportsImg4"] = False

        self.populate_tss_request_from_manifest(parameters)

        tss = TSSRequest()
        tss.add_common_tags(parameters)
        tss.add_ap_tags(parameters)

        # TODO: This break updating iPhone 15P Pro. Consider re-adding it once we figure out what went wrong
        # build_manifest_info = self.build_identity['Info']
        # for manifest_property in build_manifest_info.get('RequestManifestProperties', []):
        #     tss.add_tags({manifest_property: build_manifest_info[manifest_property]})

        # add personalized parameters
        if await self.device.get_is_image4_supported():
            tss.add_ap_img4_tags(parameters)
        else:
            tss.add_ap_img3_tags(parameters)

        # normal mode; request baseband ticket as well
        if self.device.lockdown is not None:
            pinfo = await self.device.get_firmware_preflight_info()
            if pinfo:
                self.logger.debug("adding firmware preflight info")

                node = pinfo.get("Nonce")
                if node is not None:
                    parameters["BbNonce"] = node

                node = pinfo.get("ChipID")
                if node is not None:
                    parameters["BbChipID"] = node

                node = pinfo.get("CertID")
                if node is not None:
                    parameters["BbGoldCertId"] = node

                node = pinfo.get("ChipSerialNo")
                if node is not None:
                    parameters["BbSNUM"] = node

                    # add baseband parameters
                    tss.add_baseband_tags(parameters)

                    euiccchipid = pinfo.get("EUICCChipID")
                    if euiccchipid:
                        self.logger.debug("adding EUICCChipID info")
                        parameters["eUICC,ChipID"] = euiccchipid

                        if euiccchipid >= 5:
                            node = pinfo.get("EUICCCSN")
                            if node is not None:
                                parameters["eUICC,EID"] = node

                            node = pinfo.get("EUICCCertIdentifier")
                            if node is not None:
                                parameters["eUICC,RootKeyIdentifier"] = node

                            node = pinfo.get("EUICCGoldNonce")
                            if node is not None:
                                parameters["EUICCGoldNonce"] = node

                            node = pinfo.get("EUICCMainNonce")
                            if node is not None:
                                parameters["EUICCMainNonce"] = node

                            tss.add_vinyl_tags(parameters)

        # send request and grab response
        return await tss.send_receive()

    async def get_local_policy_tss_response(self):
        # populate parameters
        parameters = {
            "ApECID": await self.device.get_ecid(),
            "Ap,LocalBoot": False,
            "ApProductionMode": True,
        }

        ap_nonce = await self.device.get_ap_nonce()
        if ap_nonce:
            parameters["ApNonce"] = ap_nonce

        sep_nonce = await self.device.get_sep_nonce()

        if sep_nonce:
            parameters["ApSepNonce"] = sep_nonce

        if await self.device.get_is_image4_supported():
            parameters["ApSecurityMode"] = True
            parameters["ApSupportsImg4"] = True
        else:
            parameters["ApSupportsImg4"] = False

        self.populate_tss_request_from_manifest(parameters)

        # Add Ap,LocalPolicy
        lpol = {
            "Digest": hashlib.sha384(lpol_file).digest(),
            "Trusted": True,
        }

        parameters["Ap,LocalPolicy"] = lpol

        # Add Ap,NextStageIM4MHash
        # Get previous TSS ticket
        ticket = self.tss.ap_img4_ticket
        # Hash it and add it as Ap,NextStageIM4MHash
        parameters["Ap,NextStageIM4MHash"] = hashlib.sha384(ticket).digest()

        # create basic request
        request = TSSRequest()

        # add common tags from manifest
        request.add_local_policy_tags(parameters)

        return await request.send_receive()

    async def get_recoveryos_root_ticket_tss_response(self):
        # populate parameters
        parameters = {
            "ApECID": await self.device.get_ecid(),
            "Ap,LocalBoot": False,
            "ApProductionMode": True,
        }

        ap_nonce = await self.device.get_ap_nonce()
        if ap_nonce:
            parameters["ApNonce"] = ap_nonce

        sep_nonce = await self.device.get_sep_nonce()

        if sep_nonce:
            parameters["ApSepNonce"] = sep_nonce

        if await self.device.get_is_image4_supported():
            parameters["ApSecurityMode"] = True
            parameters["ApSupportsImg4"] = True
        else:
            parameters["ApSupportsImg4"] = False

        self.populate_tss_request_from_manifest(parameters)

        # create basic request
        # Adds @HostPlatformInfo, @VersionInfo, @UUID
        request = TSSRequest()

        # add common tags from manifest
        # Adds Ap,OSLongVersion, AppNonce, @ApImg4Ticket
        request.add_ap_img4_tags(parameters)

        # add AP tags from manifest
        request.add_common_tags(parameters)

        # add AP tags from manifest
        # Fills digests & co
        request.add_ap_recovery_tags(parameters)

        return await request.send_receive()

    async def fetch_tss_record(self) -> TSSResponse:
        await self.ensure_image4_supported()

        if self.ipsw.build_manifest.build_major > 8 and await self.device.get_ap_nonce() is None:
            # the first nonce request with older firmware releases can fail, and it's OK
            self.logger.info("NOTE: Unable to get nonce from device")

        self.tss = await self.get_tss_response()

        if self.macos_variant:
            self.tss_localpolicy = await self.get_local_policy_tss_response()
            self.tss_recoveryos_root_ticket = await self.get_recoveryos_root_ticket_tss_response()
        else:
            recovery_variant = self.build_identity["Info"].get("RecoveryVariant")
            if recovery_variant is not None:
                self.tss_recoveryos_root_ticket = await self.get_tss_response()

        return self.tss

    async def send_component(self, name: str):
        # Use a specific TSS ticket for the Ap,LocalPolicy component
        data = None
        tss = self.tss
        if name == "Ap,LocalPolicy":
            tss = self.tss_localpolicy
            # If Ap,LocalPolicy => Inject an empty policy
            data = lpol_file

        data = await self.get_personalized_data(name, data=data, tss=tss)
        self.logger.info(f"Sending {name} ({len(data)} bytes)...")
        self.device.irecv.send_buffer(data)

    async def send_component_and_command(self, name, command):
        await self.send_component(name)
        self.device.irecv.send_command(command)

    async def send_ibec(self):
        component = "iBEC"
        await self.send_component(component)
        self.device.irecv.send_command("go", b_request=1)
        self.device.irecv.ctrl_transfer(0x21, 1)

    async def send_applelogo(self, allow_missing=True):
        component = "RestoreLogo"

        if not self.build_identity.has_component(component):
            if allow_missing:
                logging.warning(f"build_identity has no {component}")
                return
            else:
                raise PyMobileDevice3Exception(f"missing component: {component}")

        await self.send_component(component)
        self.device.irecv.send_command("setpicture 4")
        self.device.irecv.send_command("bgcolor 0 0 0")

    async def send_loaded_by_iboot(self):
        manifest = self.build_identity["Manifest"]
        for key, node in manifest.items():
            iboot = node["Info"].get("IsLoadedByiBoot", False)
            iboot_stg1 = node["Info"].get("IsLoadedByiBootStage1", False)

            assert isinstance(iboot, bool)
            assert isinstance(iboot_stg1, bool)

            if iboot and not iboot_stg1:
                self.logger.debug(f"{key} is loaded by iBoot")
                await self.send_component_and_command(key, "firmware")

    async def send_iboot_stage1_components(self):
        manifest = self.build_identity["Manifest"]
        for key, node in manifest.items():
            iboot = node["Info"].get("IsLoadedByiBoot", False)
            iboot_stg1 = node["Info"].get("IsLoadedByiBootStage1", False)

            assert isinstance(iboot, bool)
            assert isinstance(iboot_stg1, bool)

            if iboot and iboot_stg1:
                self.logger.debug(f"{key} is loaded by iBoot Stage 1")
                await self.send_component_and_command(key, "firmware")

    async def send_ramdisk(self):
        component = "RestoreRamDisk"
        ramdisk_size = self.device.irecv.getenv("ramdisk-size")
        self.logger.info(f"ramdisk-size: {ramdisk_size}")

        await self.send_component(component)
        ramdisk_delay = self.device.irecv.getenv("ramdisk-delay")
        self.logger.info(f"ramdisk-delay: {ramdisk_delay}")

        self.device.irecv.send_command("ramdisk")

        await asyncio.sleep(2)

    async def send_kernelcache(self):
        component = "RestoreKernelCache"

        await self.send_component(component)
        with contextlib.suppress(USBError):
            self.device.irecv.ctrl_transfer(0x21, 1)

        if self.restore_boot_args:
            self.device.irecv.send_command(f"setenv boot-args {self.restore_boot_args}")

        with contextlib.suppress(USBError):
            self.device.irecv.send_command("bootx", b_request=1)

    def set_autoboot(self, enable: bool):
        self.device.irecv.set_autoboot(enable)

    async def enter_restore(self):
        if self.macos_variant:
            self.restore_boot_args = "rd=md0 nand-enable-reformat=1 -progress -restore"
        elif self.ipsw.build_manifest.build_major >= 8:
            self.restore_boot_args = "rd=md0 nand-enable-reformat=1 -progress"

        # upload data to make device boot restore mode

        # Recovery Mode Environment:
        build_version = None
        while not build_version:
            self.logger.debug("build-version not yet supported. reconnecting...")
            await asyncio.sleep(1)

            # sometimes we manage to connect before iBEC actually started running
            build_version = self.device.irecv.getenv("build-version")
            await self.reconnect_irecv()

        self.logger.info(f"iBoot build-version={build_version}")

        build_style = self.device.irecv.getenv("build-style")
        self.logger.info(f"iBoot build-style={build_style}")

        radio_error = self.device.irecv.getenv("radio-error")
        if radio_error:
            radio_error = int(radio_error)
            self.logger.info(f"radio-error: {radio_error}")
            radio_error_string = self.device.irecv.getenv("radio-error-string")
            if radio_error_string:
                self.logger.info(f"radio-error-string: {radio_error_string}")

        self.set_autoboot(False)

        # send logo and show it
        await self.send_applelogo()

        # send components loaded by iBoot
        await self.send_loaded_by_iboot()

        # send ramdisk and run it
        await self.send_ramdisk()

        # send devicetree and load it
        await self.send_component_and_command("RestoreDeviceTree", "devicetree")

        if self.build_identity.has_component("RestoreSEP"):
            # attempt to send rsepfirmware and load it, otherwise continue
            with contextlib.suppress(USBError):
                await self.send_component_and_command("RestoreSEP", "rsepfirmware")

        await self.send_kernelcache()

    async def dfu_enter_recovery(self) -> None:
        await self.send_component("iBSS")
        await self.reconnect_irecv()

        if "SRTG" in self.device.irecv._device_info:
            raise PyMobileDevice3Exception("Device failed to enter recovery")

        if self.build_identity.build_manifest.build_major > 8:
            old_nonce = self.device.irecv.ap_nonce

            # reconnect
            await self.reconnect_irecv()
            nonce = self.device.irecv.ap_nonce

            if old_nonce != nonce:
                # Welcome iOS5. We have to re-request the TSS with our nonce.
                self.tss = await self.get_tss_response()

            self.device.irecv.set_configuration(1)

            # Now, before sending iBEC, we must send necessary firmwares on new versions.
            if self.macos_variant:
                # Without this empty policy file & its special signature, iBEC won't start.
                await self.send_component_and_command("Ap,LocalPolicy", "lpolrestore")
                await self.send_iboot_stage1_components()
                self.device.irecv.set_autoboot(False)
                self.device.irecv.send_command("setenvnp boot-args rd=md0 nand-enable-reformat=1 -progress -restore")
                await self.send_applelogo(allow_missing=False)

            mode = self.device.irecv.mode
            # send iBEC
            await self.send_component("iBEC")

            if self.device.irecv and mode.is_recovery:
                await asyncio.sleep(1)
                self.device.irecv.send_command("go", b_request=1)

                if self.build_identity.build_manifest.build_major < 20:
                    with contextlib.suppress(USBError):
                        self.device.irecv.ctrl_transfer(0x21, 1, timeout=5000)

                self.logger.debug("Waiting for device to disconnect...")
                await asyncio.sleep(10)

        self.logger.debug("Waiting for device to reconnect in recovery mode...")
        await self.reconnect_irecv(is_recovery=True)

    async def boot_ramdisk(self) -> None:
        if self.tss is None:
            self.logger.info("fetching TSS record")
            await self.fetch_tss_record()

        if self.device.lockdown:
            # normal mode
            self.logger.info("going into Recovery")

            # In case lockdown has disconnected while waiting for a ticket
            self.device.set_lockdown(await create_using_usbmux(serial=self.device.lockdown.udid, connection_type="USB"))
            await self.device.lockdown.enter_recovery()

            self.device.set_lockdown(None)
            self.device.set_irecv(IRecv(await self.device.get_ecid()))
            await self.reconnect_irecv()

        if self.device.irecv.mode == Mode.DFU_MODE:
            # device is currently in DFU mode, place it into recovery mode
            await self.dfu_enter_recovery()

        elif self.device.irecv.mode.is_recovery:
            # now we load the iBEC
            with contextlib.suppress(USBError):
                await self.send_ibec()

            await self.reconnect_irecv(is_recovery=True)

        self.logger.info("device booted into recovery")

        # now finally do the magic to put the device into restore mode
        await self.enter_restore()
