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

from cryptography.exceptions import InvalidTag
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

    @staticmethod
    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 = {}

    @staticmethod
    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 = {}

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

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

    @staticmethod
    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)
        if state.DHr is None or header.dh.public_bytes(serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint) != state.DHr.public_bytes(serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint):
            MessengerClient.DHRatchet(state, header)
        if header.n != state.Nr:
            return None # out of order
        state.CKr, mk = KDF_CK(state.CKr)
        state.Nr += 1
        return mk

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

    @staticmethod
    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):
        state = self.conns.get(name)
        if state is None:
            state = SimpleNamespace()
            self.conns[name] = state
            initial_secret = DH(self.dh, self.certs[name])
            self.RatchetInitAlice(state, initial_secret, self.certs[name])
        plaintext = message.encode("utf-8")
        return self.RatchetEncrypt(state, plaintext, b"")

    def receiveMessage(self, name, header, ciphertext):
        state = self.conns.get(name)
        if state is None:
            state = SimpleNamespace()
            initial_secret = DH(self.dh, self.certs[name])
            self.conns[name] = state
            self.RatchetInitBob(state, initial_secret, self.dh)
        try:
            plaintext = self.RatchetDecrypt(state, header, ciphertext, b"")
        except InvalidTag:
            return None
        return plaintext.decode("utf-8")


    # 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)