import argparse
import os

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import dh
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
import cryptography.hazmat.primitives.serialization as serialization
from cryptography.hazmat.primitives.serialization import (
    Encoding,
    NoEncryption,
    PrivateFormat,
    PublicFormat,
)

def setup(user):
    parameters = dh.generate_parameters(generator=2, key_size=512)
    private_key = parameters.generate_private_key()
    public_key = private_key.public_key()

    skey = open(f"{user}.sk", 'wb')
    pkey = open(f"{user}.pk", 'wb')
    skey.write(private_key.private_bytes(encoding=Encoding.PEM, format=PrivateFormat.PKCS8, encryption_algorithm=NoEncryption()))
    pkey.write(public_key.public_bytes(encoding=Encoding.PEM, format=PublicFormat.SubjectPublicKeyInfo))
    skey.close()
    pkey.close()


def encrypt(user,filename):
    peer_pkey_file = open(f"{user}.pk", 'rb')
    if peer_pkey_file is None:
        print(f"No public key found for {user}")
        return

    file_r = open(filename, 'rb')
    if file_r is None:
        print(f"File {filename} not found")
        return

    # FIX: Add try
    parameters = dh.generate_parameters(generator=2, key_size=512)
    skey = parameters.generate_private_key()
    peer_pkey= serialization.load_pem_public_key(peer_pkey_file.read())
    derived_key = None
    if isinstance(peer_pkey, dh.DHPublicKey):
        shared_key = skey.exchange(peer_pkey)
        derived_key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=b'handshake data').derive(shared_key)
    else:
        print("Error: Invalid key type")
        return

    message = file_r.read()
    chacha = ChaCha20Poly1305(derived_key)
    nonce = os.urandom(12)
    ciphertext = chacha.encrypt(nonce, message, None)
    file_w = open(f"{filename}.enc", 'wb')
    nonce_ciphertext = mkpair(nonce, ciphertext)
    to_send = mkpair(peer_pkey.public_bytes(Encoding.PEM, PublicFormat.SubjectPublicKeyInfo), nonce_ciphertext)
    file_w.write(to_send)

    file_r.close()
    file_w.close()
    skey_file.close()
    peer_pkey_file.close()


def decrypt(user,filename):
    skey_file = open(f"{user}.sk", 'rb')
    if skey_file is None:
        print(f"No private key found for {user}")
        return

    file_r = open(filename, 'rb')
    if file_r is None:
        print(f"File {filename} not found")
        return

    message_bytes = file_r.read()
    peer_pkey_bytes, ciphertext = unpair(message_bytes)
    # FIX: Add try
    skey  = serialization.load_pem_private_key(skey_file.read(), None)
    peer_pkey = serialization.load_pem_public_key(peer_pkey_bytes)
    derived_key = None
    if isinstance(skey, dh.DHPrivateKey) and isinstance(peer_pkey, dh.DHPublicKey):
        shared_key = skey.exchange(peer_pkey)
        derived_key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=b'handshake data').derive(shared_key)
    else:
        print("Error: Invalid key type")
        return

    nonce, encrypted_message = unpair(ciphertext)
    chacha = ChaCha20Poly1305(derived_key)
    message = chacha.decrypt(nonce, encrypted_message, None)
    file_w = open(f"{filename}.decrypt", 'wb')
    file_w.write(message)
    file_r.close()
    file_w.close()
    skey_file.close()




def mkpair(x, y):
    """ produz uma byte-string contendo o tuplo '(x,y)' ('x' e 'y' são byte-strings) """
    len_x = len(x)
    len_x_bytes = len_x.to_bytes(2, 'little')
    return len_x_bytes + x + y

def unpair(xy):
    """ extrai componentes de um par codificado com 'mkpair' """
    len_x = int.from_bytes(xy[:2], 'little')
    x = xy[2:len_x+2]
    y = xy[len_x+2:]
    return x, y

def main():
    parser = argparse.ArgumentParser(
        description="Program to perform operations using ChaCha20 cipher on files",
    )

    subparsers = parser.add_subparsers(dest="operation", help="Operation to perform")

    # Setup subcommand
    setup_parser = subparsers.add_parser("setup", help="Setup diffie helman key pair for user")
    setup_parser.add_argument("user", help="User for which to setup the key pair")

    # Encrypt subcommand
    enc_parser = subparsers.add_parser("enc", help="Encrypt a file")
    enc_parser.add_argument("user", help="User for which to encrypt the file")
    enc_parser.add_argument("file", help="File to encrypt")

    # Decrypt subcommand
    decrypt_parser = subparsers.add_parser("dec", help="Decrypt a file")
    decrypt_parser.add_argument("user", help="Target user for the decryption")
    decrypt_parser.add_argument("file", help="File to be decrypted")

    args = parser.parse_args()
    match args.operation:
        case "setup":
            user = args.user
            setup(user)
        case "enc":
            user = args.user
            file = args.file
            encrypt(user,file)
        case "dec":
            user = args.user
            file = args.file
            decrypt(user,file)

if __name__ == "__main__":
    main()