import json
import os
import sys
import shutil
import pickle
import maya
import traceback
from timeit import default_timer as timer

from nucypher.characters.lawful import Bob, Ursula
from nucypher.crypto.kits import UmbralMessageKit
from nucypher.crypto.powers import EncryptingPower, SigningPower
from nucypher.data_sources import DataSource
from nucypher.keystore.keypairs import EncryptingKeypair, SigningKeypair
from nucypher.network.middleware import RestMiddleware

from umbral.keys import UmbralPublicKey


######################
# Boring setup stuff #
######################

SEEDNODE_URL = sys.argv[1]

# TODO: path joins?
TEMP_DOCTOR_DIR = "{}/doctor-files".format(os.path.dirname(os.path.abspath(__file__)))

TEMP_URSULA_CERTIFICATE_DIR = "{}/ursula-certs".format(TEMP_DOCTOR_DIR)
TEMP_DOCTOR_CERTIFICATE_DIR = "{}/doctor-certs".format(TEMP_DOCTOR_DIR)

# Remove previous demo files and create new ones
shutil.rmtree(TEMP_DOCTOR_DIR, ignore_errors=True)
os.mkdir(TEMP_DOCTOR_DIR)
os.mkdir(TEMP_URSULA_CERTIFICATE_DIR)
os.mkdir(TEMP_DOCTOR_CERTIFICATE_DIR)

ursula = Ursula.from_seed_and_stake_info(host=SEEDNODE_URL,
                                         certificates_directory=TEMP_URSULA_CERTIFICATE_DIR,
                                         federated_only=True,
                                         minimum_stake=0)

# To create a Bob, we need the doctor's private keys previously generated.
from doctor_keys import get_doctor_privkeys
doctor_keys = get_doctor_privkeys()

bob_enc_keypair = EncryptingKeypair(private_key=doctor_keys["enc"])
bob_sig_keypair = SigningKeypair(private_key=doctor_keys["sig"])
enc_power = EncryptingPower(keypair=bob_enc_keypair)
sig_power = SigningPower(keypair=bob_sig_keypair)
power_ups = [enc_power, sig_power]

print("Creating the Doctor ...")

doctor = Bob(
    is_me=True,
    federated_only=True,
    crypto_power_ups=power_ups,
    start_learning_now=True,
    known_certificates_dir=TEMP_DOCTOR_CERTIFICATE_DIR,
    abort_on_learning_error=True,
    known_nodes=[ursula],
    save_metadata=False,
    network_middleware=RestMiddleware(),
)

print("Doctor = ", doctor)

# Let's join the policy generated by Alicia. We just need some info about it.
with open("policy-metadata.json", 'r') as f:
    policy_data = json.load(f)

policy_pubkey = UmbralPublicKey.from_bytes(bytes.fromhex(policy_data["policy_pubkey"]))
alices_sig_pubkey = UmbralPublicKey.from_bytes(bytes.fromhex(policy_data["alice_sig_pubkey"]))
label = policy_data["label"].encode()

print("The Doctor joins policy for label '{}'".format(label.decode("utf-8")))
doctor.join_policy(label, alices_sig_pubkey)

# Now that the Doctor joined the policy in the NuCypher network,
# he can retrieve encrypted data which he can decrypt with his private key.
# But first we need some encrypted data!
# Let's read the file produced by the heart monitor and unpack the MessageKits,
# which are the individual ciphertexts.
data = pickle.load(open("heart_data.pickle", "rb"))
message_kits = (UmbralMessageKit.from_bytes(k) for k in data['kits'])

# The doctor also needs to create a view of the Data Source from its public keys
data_source = DataSource.from_public_keys(
        policy_public_key=policy_pubkey,
        datasource_public_key=data['data_source'],
        label=label
)

# Now he can ask the NuCypher network to get a re-encrypted version of each MessageKit.
for message_kit in message_kits:
    try:
        start = timer()
        retrieved_plaintexts = doctor.retrieve(
            message_kit=message_kit,
            data_source=data_source,
            alice_verifying_key=alices_sig_pubkey
        )
        end = timer()

        plaintext = pickle.loads(retrieved_plaintexts[0])

        # Now we can get the heart rate and the associated timestamp,
        # generated by the heart rate monitor.
        heart_rate = plaintext['heart_rate']
        timestamp = maya.MayaDT(plaintext['timestamp'])

        # This code block simply pretty prints the heart rate info
        terminal_size = shutil.get_terminal_size().columns
        max_width = min(terminal_size, 120)
        columns = max_width - 12 - 27
        scale = columns/40
        scaled_heart_rate = int(scale * (heart_rate - 60))
        retrieval_time = "Retrieval time: {:8.2f} ms".format(1000*(end - start))
        line = ("-" * scaled_heart_rate) + "❤︎ ({} BPM)".format(heart_rate)
        line = line.ljust(max_width - 27, " ") + retrieval_time
        print(line)
    except Exception as e:
        # We just want to know what went wrong and continue the demo
        traceback.print_exc()