import logging
import struct
from typing import Union

from construct import (
    Bytes,
    Const,
    GreedyRange,
    Hex,
    Int8ub,
    Int32ub,
    PascalString,
    Rebuild,
    Struct,
    len_,
    this,
)
from intelhex import IntelHex

from ledgerwallet.crypto.ecc import PrivateKey
from ledgerwallet.crypto.scp import SCP, FakeSCP
from ledgerwallet.hsmscript import HsmScript
from ledgerwallet.hsmserver import HsmServer
from ledgerwallet.ledgerserver import LedgerServer
from ledgerwallet.manifest import AppManifest
from ledgerwallet.proto.listApps_pb2 import AppList
from ledgerwallet.simpleserver import SimpleServer
from ledgerwallet.transport import FileDevice, enumerate_devices
from ledgerwallet.utils import LedgerIns, LedgerSecureIns, VersionInfo, serialize

LOAD_SEGMENT_CHUNK_HEADER_LENGTH = 3
MIN_PADDING_LENGTH = 1
SCP_MAC_LENGTH = 0xE

MAX_CHUNK_SIZE = (
    0x10000  # Maximum size of a chunk that can be loaded at once during app install
)

LEDGER_HSM_URL = "https://hsmprod.hardwarewallet.com/hsm/process"
LEDGER_HSM_KEY = "perso_11"


ApduListAppsResponse = Struct(
    Const(b"\x01"),  # Version
    apps=GreedyRange(
        Struct(
            # Application
            # Prefixed by the size of the structure, size included.
            _size=Rebuild(Int8ub, 1 + 4 + 32 + 32 + len_(this.name)),
            flags=Hex(Int32ub),
            code_data_hash=Bytes(32),
            full_hash=Bytes(32),
            name=PascalString(Int8ub, "utf-8"),
        )
    ),
)


class AppInfo(object):
    def __init__(
        self,
        name: str,
        flags: int,
        code_data_hash: bytes,
        full_hash: bytes,
        data: bytes = None,
    ):
        self.name = name
        self.flags = flags
        self.code_data_hash = code_data_hash
        self.full_hash = full_hash
        self.data = data


class MemoryInfo:
    def __init__(
        self,
        system_size: int,
        applications_size: int,
        free_size: int,
        used_app_slots: int,
        num_app_slots: int,
    ):
        self.system_size = system_size
        self.applications_size = applications_size
        self.free_size = free_size
        self.used_app_slots = used_app_slots
        self.num_app_slots = num_app_slots


class CommException(Exception):
    def __init__(self, message, sw=0x6F00, data=None):
        self.message = message
        self.sw = sw
        self.data = data

    def __str__(self):
        buf = "Exception : " + self.message
        return buf


class NoLedgerDeviceException(Exception):
    pass


LOG = logging.getLogger("ledgerwallet")


class LedgerClient(object):
    def __init__(self, device=None, cla=0xE0, private_key=None):
        self.scp = None
        if device is None:
            devices = enumerate_devices()
            if len(devices) == 0:
                raise NoLedgerDeviceException("No Ledger device has been found.")
            device = devices[0]
        elif type(device) == FileDevice:
            self.scp = FakeSCP()

        self.device = device
        self.cla = cla
        self._target_id = None
        if private_key is None:
            self.private_key = PrivateKey()
        else:
            self.private_key = PrivateKey(private_key)
        self.device.open()

    def close(self):
        self.device.close()

    def raw_exchange(self, data: bytes) -> bytes:
        LOG.debug("=> " + data.hex())
        output_data = bytes(self.device.exchange(data))
        if len(output_data) > 0:
            LOG.debug("<= " + output_data.hex())
        return output_data

    def apdu_exchange(self, ins, data=b"", p1=0, p2=0):
        apdu = bytes([self.cla, ins, p1, p2])
        apdu += serialize(data)
        response = self.raw_exchange(apdu)

        status_word = int.from_bytes(response[-2:], "big")
        if status_word != 0x9000 and ((status_word >> 8) != 0x61):
            possible_cause = "Unknown reason"
            if status_word == 0x6982:
                possible_cause = (
                    "Have you uninstalled the existing CA with resetCustomCA first?"
                )
            elif status_word == 0x6985:
                possible_cause = "Condition of use not satisfied (denied by the user?)"
            elif status_word == 0x6A84 or status_word == 0x6A85:
                possible_cause = "Not enough space?"
            elif status_word == 0x6A83:
                possible_cause = (
                    "Maybe this app requires a library to be installed first?"
                )
            elif status_word == 0x6484:
                possible_cause = "Are you using the correct targetId?"
            elif status_word == 0x69D5:
                possible_cause = "Cannot create custom secure channels on this device"
            raise CommException(
                "Invalid status %04x (%s)" % (status_word, possible_cause),
                status_word,
                data,
            )

        return response[:-2]

    def apdu_secure_exchange(self, ins, data=b"", p1=0, p2=0):
        if self.scp is None:
            server = SimpleServer(self.private_key)
            secret = self.authenticate(server)
            self.scp = SCP(secret)

        data = self.apdu_exchange(
            LedgerIns.SECUINS, self.scp.wrap(bytes([ins]) + data), p1, p2
        )
        return self.scp.unwrap(data)

    def authenticate(self, server: LedgerServer):
        self.reset()
        if self.target_id & 0xF < 2:
            raise Exception("Target ID does not support SCP V2")

        # Exchange nonce
        server_nonce = server.get_nonce()
        data = self.apdu_exchange(LedgerIns.INITIALIZE_AUTHENTICATION, server_nonce)
        device_nonce = data[4:12]
        server.send_nonce(device_nonce)

        # Get server certificate chain
        server_chain = server.receive_certificate_chain()
        for i in range(len(server_chain)):
            if i == len(server_chain) - 1:
                self.apdu_exchange(
                    LedgerIns.VALIDATE_CERTIFICATE, server_chain[i], p1=0x80
                )
            else:
                self.apdu_exchange(LedgerIns.VALIDATE_CERTIFICATE, server_chain[i])

        # Walk the client chain
        client_chain = []
        for i in range(2):
            if i == 0:
                certificate = self.apdu_exchange(LedgerIns.GET_CERTIFICATE)
            else:
                certificate = self.apdu_exchange(LedgerIns.GET_CERTIFICATE, p1=0x80)
            if len(certificate) == 0:
                break
            client_chain.append(certificate)
        server.send_certificate_chain(client_chain)

        # Mutual authentication done, retrieve shared secret
        self.apdu_exchange(LedgerIns.MUTUAL_AUTHENTICATE)
        return server.get_shared_secret()

    def _load_chunk(self, hex_file: IntelHex, segment, offset: int):
        start_addr, end_addr = segment
        segment_load_address = start_addr - hex_file.minaddr()

        self.apdu_secure_exchange(
            LedgerSecureIns.SET_LOAD_OFFSET,
            struct.pack(">I", segment_load_address + offset),
        )

        load_size = min(end_addr - start_addr - offset, MAX_CHUNK_SIZE)
        # max_load_size = 0xf0 - LOAD_SEGMENT_CHUNK_HEADER_LENGTH - MIN_PADDING_LENGTH - SCP_MAC_LENGTH # noqa
        max_load_size = 0x80

        load_address = start_addr + offset
        chunk_offset = start_addr
        while load_size > 0:
            chunk_size = min(load_size, max_load_size)
            data = hex_file.gets(load_address, chunk_size)
            data = struct.pack(">H", chunk_offset - start_addr) + data

            self.apdu_secure_exchange(LedgerSecureIns.LOAD, data)
            load_address += chunk_size
            chunk_offset += chunk_size
            load_size -= chunk_size

    def _load_segment(self, hex_file: IntelHex, segment):
        start_addr, end_addr = segment
        # Load each segment by chunks of 64kB
        for offset in range(0, end_addr - start_addr, MAX_CHUNK_SIZE):
            self._load_chunk(hex_file, segment, offset)

    def install_app(self, app_manifest: AppManifest):
        version_info = self.get_version_info()
        device = str(version_info.target_id)

        app_manifest.assert_compatible_device(version_info.target_id)

        hex_file = IntelHex(app_manifest.get_binary(device))
        code_length = hex_file.maxaddr() - hex_file.minaddr() + 1
        data_length = app_manifest.data_size(device)

        code_length -= data_length
        assert code_length % 64 == 0  # code length must be aligned

        flags = app_manifest.get_application_flags(device)  # not handled yet

        params = app_manifest.serialize_parameters(device)
        main_address = hex_file.start_addr["EIP"] - hex_file.minaddr()

        level = app_manifest.get_api_level(device)
        if level:
            data = struct.pack(
                ">BIIIII",
                level,
                code_length,
                data_length,
                len(params),
                flags,
                main_address,
            )
        else:
            data = struct.pack(
                ">IIIII", code_length, data_length, len(params), flags, main_address
            )
        self.apdu_secure_exchange(LedgerSecureIns.CREATE_APP, data)

        hex_file.puts(hex_file.maxaddr() + 1, params)
        for segment in hex_file.segments():
            self._load_segment(hex_file, segment)
        self.apdu_secure_exchange(LedgerSecureIns.COMMIT)

    def delete_app(self, app: Union[str, bytes]):
        if isinstance(app, str):
            self.apdu_secure_exchange(
                LedgerSecureIns.DELETE_APP, serialize(app.encode())
            )
        elif isinstance(app, bytes) and len(app) == 32:
            self.apdu_secure_exchange(LedgerSecureIns.DELETE_APP_BY_HASH, app)
        else:
            raise TypeError("app parameter must be string or digest")

    def install_remote_app(
        self, app_path, key_path, url=LEDGER_HSM_URL, key=LEDGER_HSM_KEY
    ):
        script = HsmScript("distributeFirmware11", {"persoKey": key, "scpv2": "dummy"})
        server = HsmServer(script, url)
        self.authenticate(server)

        application_data = server.query(
            params={"firmware": app_path, "firmwareKey": key_path, "scpv2": "dummy"}
        )
        offset = 0
        while offset < len(application_data):
            apdu_len = application_data[offset + 4]
            self.raw_exchange(application_data[offset : offset + 5 + apdu_len])
            offset += 5 + apdu_len

    def delete_remote_app(
        self, app_path, key_path, url=LEDGER_HSM_URL, key=LEDGER_HSM_KEY
    ):
        return self.install_remote_app(app_path, key_path, url, key)

    def upgrade_firmware(
        self, firmware_name, firmware_key, url=LEDGER_HSM_URL, key=LEDGER_HSM_KEY
    ):
        script = HsmScript(
            "distributeFirmware11_scan", {"persoKey": key, "scpv2": "dummy"}
        )
        server = HsmServer(script, url)
        self.authenticate(server)

        response = b""
        while True:
            app_data = server.query(
                data=None if len(response) < 2 else b"\xff\xff" + response[:-2]
            )
            if len(app_data) == 0:
                break
            response = self.raw_exchange(app_data)

        application_data = server.query(
            params={"firmware": firmware_name, "firmwareKey": firmware_key}
        )
        offset = 0
        while offset < len(application_data):
            apdu_len = application_data[offset + 4]
            self.raw_exchange(application_data[offset : offset + 5 + apdu_len])
            offset += 5 + apdu_len

    def genuine_check(self, url=LEDGER_HSM_URL, key=LEDGER_HSM_KEY):
        script = HsmScript("checkGenuine", {"persoKey": key, "scpv2": "dummy"})
        server = HsmServer(script, url)
        self.authenticate(server)

        client_data = b""
        while True:
            application_data = server.query(client_data[:-2])  # No not send status word
            if len(application_data) < 5:
                break
            client_data = self.raw_exchange(application_data)
        # custom_ui = client_data[0]
        # custom_ca = client_data[1]
        return True

    def endorse(self, key_id: int, url=LEDGER_HSM_URL, key=LEDGER_HSM_KEY):
        script = HsmScript("signEndorsement", {"persoKey": key})
        server = HsmServer(script, url)
        self.authenticate(server)
        server.query()  # Commit agreement

        data = self.apdu_exchange(LedgerIns.ENDORSE_SET_START, p1=key_id)
        certificate = server.query(data, params={"endorsementKey": "attest_1"})

        # Commit endorsement certificate
        self.apdu_exchange(LedgerIns.ENDORSE_SET_COMMIT, certificate)
        return True

    def install_ca(self, name: str, public_key: bytes):
        data = serialize(name.encode()) + serialize(public_key)
        self.apdu_secure_exchange(LedgerSecureIns.SETUP_CUSTOM_CERTIFICATE, data)

    def delete_ca(self):
        return self.apdu_secure_exchange(LedgerSecureIns.RESET_CUSTOM_CERTIFICATE)

    def get_version_info(self):
        data = self.apdu_exchange(LedgerIns.GET_VERSION)
        version_info = VersionInfo.parse(data)
        self._target_id = version_info.target_id
        return version_info

    def get_version_info_secure(self):
        data = self.apdu_secure_exchange(LedgerSecureIns.GET_VERSION)
        version_info = VersionInfo.parse(data)
        self._target_id = version_info.target_id
        return version_info

    def validate_target_id(self, target_id: int):
        self.apdu_exchange(LedgerIns.VALIDATE_TARGET_ID, struct.pack(">I", target_id))

    def reset(self):
        return self.validate_target_id(self.target_id)

    def get_memory_info(self) -> MemoryInfo:
        response = self.apdu_secure_exchange(LedgerSecureIns.GET_MEMORY_INFORMATION)
        assert len(response) == 20

        return MemoryInfo(*struct.unpack(">IIIII", response))

    @property
    def target_id(self):
        if self._target_id is None:
            self.get_version_info()
        return self._target_id

    def list_apps_remote(self, url=LEDGER_HSM_URL, key=LEDGER_HSM_KEY):
        script = HsmScript("listApps", {"persoKey": key, "scpv2": "dummy"})
        server = HsmServer(script, url)
        self.authenticate(server)

        application_data = server.query(params={"scpv2": "dummy"})
        while True:
            if len(application_data) == 0:
                break
            client_data = self.raw_exchange(application_data)
            application_data = server.query(client_data[:-2], params={"scpv2": "dummy"})
        application_data = server.query(params={"scpv2": "dummy"})

        apps = AppList()
        apps.ParseFromString(application_data)
        for app in apps.list:
            yield AppInfo(app.name, app.flags & 0xFFFF, app.hashCodeData, app.hash)

    @property
    def apps(self):
        data = self.apdu_secure_exchange(LedgerSecureIns.LIST_APPS)
        while len(data) != 0:
            response = ApduListAppsResponse.parse(data)
            for app in response.apps:
                yield AppInfo(
                    app.name, app.flags & 0xFFFF, app.code_data_hash, app.full_hash
                )
            data = self.apdu_secure_exchange(LedgerSecureIns.LIST_APPS_CONTINUE)

    def run_app(self, app_name: str):
        return self.apdu_exchange(LedgerIns.RUN_APP, app_name.encode())
