# Internal Dependencies
from seedsigner.helpers import Buttons, B, CameraProcess, CameraPoll
from seedsigner.views import View

# External Dependencies
import time
from embit import bip32, bip39, ec, psbt, script
from embit.bip39 import mnemonic_to_bytes, mnemonic_from_bytes
from embit.networks import NETWORKS
from io import BytesIO
from binascii import unhexlify, hexlify, a2b_base64, b2a_base64
import textwrap

class Wallet:

    QRLOW = 0
    QRMEDIUM = 1
    QRHIGH = 2

    def __init__(self, current_network, qr_density, policy) -> None:
        self.current_network = current_network

        if policy not in self.avaliable_wallet_policies():
            policy = "PKWSH" #override policy to PKWSH when not found in wallet 

        if policy == "PKWSH" and self.current_network == "main":
            self.hardened_derivation = "m/48h/0h/0h/2h"
        elif policy == "PKWSH" and self.current_network == "test":
            self.hardened_derivation = "m/48h/1h/0h/2h"
        elif policy == "PKWPKH" and self.current_network == "main":
            self.hardened_derivation = "m/84h/0h/0h"
        elif policy == "PKWPKH" and self.current_network == "test":
            self.hardened_derivation = "m/84h/1h/0h"
        else:
            raise Exception("Unsupported Derivation Path or Policy")

        self.qrsize = 80 # Default
        self.set_qr_density(qr_density)
        self.cur_qr_density = qr_density
        self.cur_policy = policy

    def set_seed_phrase(self, seed_phrase, passphrase):
        # requires a valid seed phrase or error will be thrown
        self.seed_phrase = seed_phrase
        self.seed = bip39.mnemonic_to_seed((" ".join(self.seed_phrase)).strip(), passphrase)
        self.root = bip32.HDKey.from_seed(self.seed, version=NETWORKS[self.current_network]["xprv"])
        self.fingerprint = self.root.child(0).fingerprint
        self.bip48_xprv = self.root.derive(self.hardened_derivation)
        self.bip48_xpub = self.bip48_xprv.to_public()

        self.tx = None
        self.inp_amount = None
        self.fee = None
        self.spend = None
        self.destinationaddress = None
        self.dest_addr_cnt = None
        self.change = None
        self.controller = None
        self.buttons = None
        self.ins = None
        self.outs = None

        self.camera_loop_timer = None
        self.camera_data = None
        self.is_camera_data = False

        self.qr_total_frames = 0
        self.qr_cur_frame_count = 0
        self.qr_data = []
        self.frame_display = []
        self.percentage_complete = 0

        self.scan_started_ind = 0
        self.scan_display_working = 0

    ###
    ### Required Methods to implement for Child Wallet Class
    ###
    ### import_qr, parse_psbt, sign_transaction, total_frames_parse, current_frame_parse, data_parse, capture_complete
    ### get_name, set_network, make_xpub_qr_codes, make_signing_qr_codes, set_qr_density

    def import_qr(self) -> str:
        if self.cur_policy == "PKWPKH":
            xpubstring = "[%s%s]%s" % (
                 hexlify(self.fingerprint).decode('utf-8'),
                 self.hardened_derivation[1:],
                 self.bip48_xpub.to_base58(NETWORKS[self.current_network]["zpub"]))
        else:
            xpubstring = "[%s%s]%s" % (
                 hexlify(self.fingerprint).decode('utf-8'),
                 self.hardened_derivation[1:],
                 self.bip48_xpub.to_base58(NETWORKS[self.current_network]["Zpub"]))

        return xpubstring

    def get_xpub_info(self) -> (str, str, str):
        if self.cur_policy == "PKWPKH":
            xpub = self.bip48_xpub.to_base58(NETWORKS[self.current_network]["zpub"])
        else:
            xpub = self.bip48_xpub.to_base58(NETWORKS[self.current_network]["Zpub"])

        derivation = self.hardened_derivation[1:].replace("h","'")

        return (hexlify(self.fingerprint).decode('utf-8'), derivation, xpub)

    def parse_psbt(self, raw_psbt) -> bool:
        # decodes and parses raw_psbt, also calculates the following instance values
        self.inp_amount = None
        self.change = None
        self.fee = None
        self.spend = None
        self.destinationaddress = None
        return False

    def sign_transaction(self) -> (bool, str):
        # signs transaction/pbst last passed to parse_psbt method
        return (False, '')

    def total_frames_parse(data) -> int:
        # parse and returns total number of frames from qr data frame
        return -1

    def current_frame_parse(data) -> int:
        # parses and returns current frame number from qr data frame
        return -1

    def data_parse(data) -> str:
        # parse qr data to string to be cancatinated together into a pbst transaction
        return "empty"

    def capture_complete() -> bool:
        # returns true if the qr data list is complete
        return False

    def scan_animated_qr_pbst(self, controller) -> str:
        self.controller = controller
        self.buttons = controller.buttons
        self.controller.menu_view.draw_modal(["Initializing Camera"]) # TODO: Move to Controller
        # initialize camera
        self.controller.to_camera_queue.put(["start"])
        # First get blocking, this way it's clear when the camera is ready for the end user
        self.controller.from_camera_queue.get()
        self.camera_loop_timer = CameraPoll(0.05, self.process_camera_data)

        input = self.buttons.wait_for([B.KEY_LEFT, B.KEY_RIGHT])
        if input in (B.KEY_LEFT, B.KEY_RIGHT):
            self.camera_loop_timer.stop()
            self.controller.to_camera_queue.put(["stop"])
            return "nodata"
        elif input == B.OVERRIDE:
            self.camera_loop_timer.stop()
            self.controller.to_camera_queue.put(["stop"])
            if self.qr_data[0] == "invalid":
                return "invalid"
            time.sleep(0.5) # give time for camera loop to complete before returning data
            return "".join(self.qr_data)

    def process_camera_data(self):
        try:
            data = self.controller.from_camera_queue.get(False)
        except:
            data = ["nodata"]

        if data[0] != "nodata":
            if self.qr_total_frames == 0:
                # get total frames if not set
                self.qr_total_frames = type(self).total_frames_parse(data[0])
                if self.qr_total_frames == -1:
                    # when invalid, trigger override to display error
                    self.qr_data = ["invalid"]
                    self.buttons.trigger_override() # something went wrong, invalid QR
                    return

                # create qr_data list with number of total frames
                self.qr_data = ["empty"] * self.qr_total_frames
                # create frame display / progress with number of total frames
                self.frame_display = ["-"] * self.qr_total_frames

            # get current frame
            current_frame = type(self).current_frame_parse(data[0])
            if self.qr_data[current_frame - 1] == "empty":
                # if frame has never been captured, store data element in it
                self.qr_data[current_frame - 1] = type(self).data_parse(data[0])
                # increment number of frames captured
                self.qr_cur_frame_count += 1
                # show in frame display / progress of captured frame
                self.frame_display[current_frame - 1] = "*"
                # calculate percentage complete of captured frames
                self.percentage_complete = int((self.qr_cur_frame_count / self.qr_total_frames) * 100)

            # checking if all frames has been captured, exit camera processing
            if self.capture_complete():
                self.buttons.trigger_override()

            # if all frames have not all been captured, display progress to screen/display
            if not self.capture_complete() and self.scan_display_working == 0:
                self.scan_display_working = 1
                View.draw.rectangle((0, 0, View.canvas_width, View.canvas_height), outline=0, fill=0)
                tw, th = View.draw.textsize("Collecting QR Codes:", font=View.IMPACT22)
                View.draw.text(((240 - tw) / 2, 15), "Collecting QR Codes:", fill=View.color, font=View.IMPACT22)
                lines = textwrap.wrap("".join(self.frame_display), width=11)
                yheight = 60
                for line in lines:
                    tw, th = View.draw.textsize(line, font=View.COURIERNEW30)
                    View.draw.text(((240 - tw) / 2, yheight), line, fill=View.color, font=View.COURIERNEW30)
                    yheight += 30
                tw, th = View.draw.textsize("Right to Exit", font=View.IMPACT18)
                View.draw.text(((240 - tw) / 2, 215), "Right to Exit", fill=View.color, font=View.IMPACT18)
                View.DispShowImage()
                self.scan_display_working = 0

        elif self.scan_started_ind == 0:
            self.scan_started_ind = 1
            self.controller.menu_view.draw_modal(["Scan Animated QR"], "", "Right to Exit")

    def make_xpub_qr_codes(self, data, callback = None) -> []:
        return []

    def make_signing_qr_codes(self, data, callback = None) -> []:
        return []

    def qr_sleep(self):
        time.sleep(0.2)

    def set_qr_density(self, density):
        self.cur_qr_density = density
        if density == Wallet.QRLOW:
            self.qrsize = 60
        elif density == Wallet.QRMEDIUM:
            self.qrsize = 80
        elif density == Wallet.QRHIGH:
            self.qrsize = 100

    def get_qr_density(self):
        return self.cur_qr_density

    def get_qr_density_name(self) -> str :
        if self.cur_qr_density == Wallet.QRLOW:
            return "Low"
        elif self.cur_qr_density == Wallet.QRMEDIUM:
            return "Medium"
        elif self.cur_qr_density == Wallet.QRHIGH:
            return "High"
        else:
            return "Unknown"

    def get_wallet_policy_name(self) -> str :
        if self.cur_policy == "PKWSH":
            return "Multi Sig"
        elif self.cur_policy == "PKWPKH":
            return "Single Sig"
        else:
            return self.cur_policy

    def get_wallet_policy(self) -> str:
        return self.cur_policy

    def avaliable_wallet_policies(self) -> []:
        return ["PKWSH", "PKWPKH"]

    ###
    ### Network Related Methods
    ###

    def get_network(self) -> str:
        return self.current_network

    def get_hardened_derivation(self) -> str:
        return self.hardened_derivation

    def set_network(self, network) -> bool:
        return False

    ###
    ### Internal Wallet Transactions
    ###

    def input_amount(self, tx) -> (float, str, float):
        # Check inputs of the transaction and check that they use the same script type
        # For multisig parsed policy will look like this:
        # { script_type: p2wsh, cosigners: [xpubs strings], m: 2, n: 3}
        policy = None
        inp_amount = 0.0
        ins = 0
        for inp in tx.inputs:
            ins += 1
            inp_amount += inp.witness_utxo.value
            # get policy of the input
            inp_policy = self.get_policy(inp, inp.witness_utxo.script_pubkey, tx.xpubs)
            # if policy is None - assign current
            if policy is None:
                policy = inp_policy
            # otherwise check that everything in the policy is the same
            else:
                # check policy is the same
                if policy != inp_policy:
                    raise RuntimeError("Mixed inputs in the transaction")

        return (inp_amount, policy, ins)

    def change_fee_spend_amounts(self, tx, inp_amount, policy, currentnetwork) -> (float, float, float, str, float, float):
        spend = 0
        change = 0
        destinationaddress = ""
        dest_addr_cnt = 0
        outs = 0
        for i, out in enumerate(tx.outputs):
            outs += 1
            out_policy = self.get_policy(out, tx.tx.vout[i].script_pubkey, tx.xpubs)
            is_change = False
            # if policy is the same - probably change
            if out_policy == policy:
                # double-check that it's change
                # we already checked in get_cosigners and parse_multisig
                # that pubkeys are generated from cosigners,
                # and witness script is corresponding multisig
                # so we only need to check that scriptpubkey is generated from
                # witness script

                # empty script by default
                sc = script.Script(b"")
                # multisig, we know witness script
                if policy["type"] == "p2wsh":
                    sc = script.p2wsh(out.witness_script)
                elif policy["type"] == "p2sh-p2wsh":
                    sc = script.p2sh(script.p2wsh(out.witness_script))
                # single-sig native segwit
                elif "pkh" in policy["type"]:
                    for pub in out.bip32_derivations:
                        # check if it is our key
                        if out.bip32_derivations[pub].fingerprint == self.fingerprint:
                            hdkey = self.root.derive(out.bip32_derivations[pub].derivation)
                            mypub = hdkey.key.get_public_key()
                            if mypub != pub:
                                raise ValueError("Derivation path doesn't look right")
                            # now check if provided scriptpubkey matches
                            sc = script.p2wpkh(mypub)
                            if sc == tx.tx.vout[i].script_pubkey:
                                is_change = True
                if sc.data == tx.tx.vout[i].script_pubkey.data:
                    is_change = True
            if is_change:
                change += tx.tx.vout[i].value
                print("Change %d sats" % tx.tx.vout[i].value)
            else:
                spend += tx.tx.vout[i].value
                print("Spending %d sats to %s" % (tx.tx.vout[i].value, tx.tx.vout[i].script_pubkey.address(NETWORKS[currentnetwork])))
                destinationaddress = tx.tx.vout[i].script_pubkey.address(NETWORKS[currentnetwork])
                dest_addr_cnt += 1

        fee = inp_amount - change - spend

        return (change, fee, spend, destinationaddress, outs, dest_addr_cnt)

    def _parse_psbt(self):
        (self.inp_amount, policy, self.ins) = self.input_amount(self.tx)
        (self.change, self.fee, self.spend, self.destinationaddress, self.outs, self.dest_addr_cnt) = self.change_fee_spend_amounts(self.tx, self.inp_amount, policy, self.current_network)

    def parse_multisig(self, sc):
        """Takes a script and extracts m,n and pubkeys from it"""
        # OP_m <len:pubkey> ... <len:pubkey> OP_n OP_CHECKMULTISIG
        # check min size
        if len(sc.data) < 37 or sc.data[-1] != 0xae:
            raise ValueError("Not a multisig script")
        m = sc.data[0] - 0x50
        if m < 1 or m > 16:
            raise ValueError("Invalid multisig script")
        n = sc.data[-2] - 0x50
        if n < m or n > 16:
            raise ValueError("Invalid multisig script")
        s = BytesIO(sc.data)
        # drop first byte
        s.read(1)
        # read pubkeys
        pubkeys = []
        for i in range(n):
            char = s.read(1)
            if char != b"\x21":
                raise ValueError("Invlid pubkey")
            pubkeys.append(ec.PublicKey.parse(s.read(33)))
        # check that nothing left
        if s.read() != sc.data[-2:]:
            raise ValueError("Invalid multisig script")
        return m, n, pubkeys

    def get_cosigners(self, pubkeys, derivations, xpubs):
        """Returns xpubs used to derive pubkeys using global xpub field from psbt"""
        cosigners = []
        for i, pubkey in enumerate(pubkeys):
            if pubkey not in derivations:
                raise ValueError("Missing derivation")
            der = derivations[pubkey]
            for xpub in xpubs:
                origin_der = xpubs[xpub]
                # check fingerprint
                if origin_der.fingerprint == der.fingerprint:
                    # check derivation - last two indexes give pub from xpub
                    if origin_der.derivation == der.derivation[:-2]:
                        # check that it derives to pubkey actually
                        if xpub.derive(der.derivation[-2:]).key == pubkey:
                            # append strings so they can be sorted and compared
                            cosigners.append(xpub.to_base58())
                            break
        if len(cosigners) != len(pubkeys):
            raise RuntimeError("Can't get all cosigners")
        return sorted(cosigners)

    def get_policy(self, scope, scriptpubkey, xpubs):
        """Parse scope and get policy"""
        # we don't know the policy yet, let's parse it
        script_type = scriptpubkey.script_type()
        # p2sh can be either legacy multisig, or nested segwit multisig
        # or nested segwit singlesig
        if script_type == "p2sh":
            if scope.witness_script is not None:
                script_type = "p2sh-p2wsh"
            elif scope.redeem_script is not None and scope.redeem_script.script_type() == "p2wpkh":
                script_type = "p2sh-p2wpkh"
        policy = { "type": script_type }
        # expected multisig
        if "p2wsh" in script_type and scope.witness_script is not None:
            m, n, pubkeys = self.parse_multisig(scope.witness_script)

            # check pubkeys are derived from cosigners
            cosigners = self.get_cosigners(pubkeys, scope.bip32_derivations, xpubs)
            policy.update({
                "m": m, "n": n, "cosigners": cosigners
            })
        return policy