import os
import pickle
import string
import hmac
from dataclasses import dataclass

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

CURVE = ec.SECP256R1()

@dataclass
class MessageHeader:
    dh: object
    pn: int
    n: int

def GENERATE_DH():
    return ec.generate_private_key(CURVE)

def DH(dh_pair, dh_pub):
    return dh_pair.exchange(ec.ECDH(), dh_pub)

def KDF_RK(rk, dh_out):
    output = HKDF(
        algorithm=hashes.SHA256(),
        length=64,
        salt=rk,
        info=b"",
    ).derive(dh_out)
    return output[:32], output[32:]

def KDF_CK(ck):
    return (
        hmac.new(ck, b"\x01", "sha256").digest(),
        hmac.new(ck, b"\x02", "sha256").digest(),
    )

def ENCRYPT(mk, plaintext, associated_data):
    nonce = os.urandom(12)
    return nonce + AESGCM(mk).encrypt(nonce, plaintext, associated_data)

def DECRYPT(mk, ciphertext, associated_data):
    return AESGCM(mk).decrypt(ciphertext[:12], ciphertext[12:], associated_data)

def HEADER(dh_pair, pn, n):
    public_key = dh_pair.public_key()
    return MessageHeader(public_key, int(pn), int(n))

def CONCAT(ad, header):
    header_tup = pickle.dumps(
        (header.dh.public_bytes(serialization.Encoding.X962,serialization.PublicFormat.UncompressedPoint), header.pn, header.n),
    )
    return len(ad).to_bytes(8) + ad + header_tup

class MessengerServer:
    def __init__(self, server_signing_key, server_decryption_key):
        self.server_signing_key = server_signing_key
        self.server_decryption_key = server_decryption_key

    def decryptReport(self, ct):
        data = pickle.loads(ct)
        key = ec.EllipticCurvePublicKey.from_encoded_point(
            CURVE, data["bytes"]
        )
        shared = DH(self.server_decryption_key, key)
        derived_key = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b""
        ).derive(shared)
        aesgcm = AESGCM(derived_key)
        plaintext = aesgcm.decrypt(data["nonce"], data["ciphertext"], data["bytes"])
        return plaintext.decode("latin1")

    def signCert(self, cert):
        return self.server_signing_key.sign(cert, ec.ECDSA(hashes.SHA256()))

class MessengerClient:

    def __init__(self, name, server_signing_pk, server_encryption_pk):
        self.name = name
        self.server_signing_pk = server_signing_pk
        self.server_encryption_pk = server_encryption_pk
        self.conns = {}
        self.certs = {}
        self.dh = GENERATE_DH()
        self.certificate = None

    def RatchetInitAlice(state, SK, bob_dh_public_key):
        state.DHs = GENERATE_DH()
        state.DHr = bob_dh_public_key
        state.RK, state.CKs = KDF_RK(SK, DH(state.DHs, state.DHr))
        state.CKr = None
        state.Ns = 0
        state.Nr = 0
        state.PN = 0
        state.MKSKIPPED = {}

    def RatchetInitBob(state, SK, bob_dh_key_pair):
        state.DHs = bob_dh_key_pair
        state.DHr = None
        state.RK = SK
        state.CKs = None
        state.CKr = None
        state.Ns = 0
        state.Nr = 0
        state.PN = 0
        state.MKSKIPPED = {}

    def RatchetSendKey(state):
        state.CKs, mk = KDF_CK(state.CKs)
        Ns = state.Ns
        state.Ns += 1
        return Ns, mk

    def RatchetEncrypt(state, plaintext, AD):
        Ns, mk = RatchetSendKey(state)
        header = HEADER(state.DHs, state.PN, Ns)
        return header, ENCRYPT(mk, plaintext, CONCAT(AD, header))

    def RatchetReceiveKey(state, header):
        mk = TrySkippedMessageKeys(state, header)
        if mk != None:
            return mk
        if header.dh != state.DHr:
            SkipMessageKeys(state, header.pn)
            DHRatchet(state, header)
        SkipMessageKeys(state, header.n)
        state.CKr, mk = KDF_CK(state.CKr)
        state.Nr += 1
        return mk

    def RatchetDecrypt(state, header, ciphertext, AD):
        mk = RatchetReceiveKey(state, header)
        return DECRYPT(mk, ciphertext, CONCAT(AD, header))

    def TrySkippedMessageKeys(state, header):
        if (header.dh, header.n) in state.MKSKIPPED:
            mk = state.MKSKIPPED[header.dh, header.n]
            del state.MKSKIPPED[header.dh, header.n]
            return mk
        else:
            return None

    def SkipMessageKeys(state, until):
        if state.Nr + MAX_SKIP < until:
            raise Error()
        if state.CKr != None:
            while state.Nr < until:
                state.CKr, mk = KDF_CK(state.CKr)
                state.MKSKIPPED[state.DHr, state.Nr] = mk
                state.Nr += 1

    def DHRatchet(state, header):
        state.PN = state.Ns
        state.Ns = 0
        state.Nr = 0
        state.DHr = header.dh
        state.RK, state.CKr = KDF_RK(state.RK, DH(state.DHs, state.DHr))
        state.DHs = GENERATE_DH()
        state.RK, state.CKs = KDF_RK(state.RK, DH(state.DHs, state.DHr))

    def generateCertificate(self):
        self.certificate = pickle.dumps((self.name, self.dh.public_key().public_bytes(serialization.Encoding.X962,
        serialization.PublicFormat.UncompressedPoint)))
        return self.certificate

    def receiveCertificate(self, certificate, signature):
        self.server_signing_pk.verify(signature, certificate, ec.ECDSA(hashes.SHA256()))
        name, public_bytes = pickle.loads(certificate)
        public_key = ec.EllipticCurvePublicKey.from_encoded_point(CURVE, public_bytes)
        self.certs[name] = public_key

    def sendMessage(self, name, message):
        raise Exception("not implemented!")
        return

    def receiveMessage(self, name, header, ciphertext):
        raise Exception("not implemented!")
        return

    # this is technically hybrid encryption instead of just el gamal but I blame the instructions instead of myself
    # maybe i'll change this if I want to use a late day
    def report(self, name, message):
        report = (name, message)
        plaintext = pickle.dumps(report)
        temp_key = GENERATE_DH()
        temp_bytes = temp_key.public_key().public_bytes(
            serialization.Encoding.X962,
            serialization.PublicFormat.UncompressedPoint,
        )
        shared_secret = DH(temp_key, self.server_encryption_pk)
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b""
        )
        key = hkdf.derive(shared_secret)
        nonce = os.urandom(12)
        aesgcm = AESGCM(key)
        ciphertext = aesgcm.encrypt(nonce, plaintext, temp_bytes) # not sure if passing in temp_bytes is correct, but it seems like it should be
        data = {
            "bytes": temp_bytes,
            "nonce": nonce,
            "ciphertext": ciphertext
        }
        return plaintext.decode("latin1"), pickle.dumps(data)