import os
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import json
import base64
from config.config_loader import load_config

def encrypt_data(data, symmetric_key):
    # Convert the data to a JSON string
    json_data = json.dumps(data).encode('utf-8')
    
    # Generate a random IV for encryption
    iv = os.urandom(16)
    
    # Create AES Cipher and encrypt the data
    cipher = Cipher(algorithms.AES(symmetric_key), modes.CFB(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    encrypted_data = encryptor.update(json_data) + encryptor.finalize()

    # Combine IV and encrypted data
    encrypted_payload = iv + encrypted_data
    
    # Encode the result in base64 to make it JSON-compatible
    return base64.b64encode(encrypted_payload).decode('utf-8')

# AES Decryption
def decrypt_data(encrypted_data, symmetric_key):
    # Decode base64
    encrypted_data = base64.b64decode(encrypted_data)
    
    iv = encrypted_data[:16]
    cipher = Cipher(algorithms.AES(symmetric_key), modes.CFB(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    decrypted_data = decryptor.update(encrypted_data[16:]) + decryptor.finalize()
    return decrypted_data.decode('utf-8')

def generate_ecc_key_pair():
    private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
    public_key = private_key.public_key()
    return private_key, public_key

def load_client_public_key(config):
    client_public_key_pem = config.get('client_keys', {}).get('public_key')
    if not client_public_key_pem:
        raise ValueError("Client public key not found")
    return serialization.load_pem_public_key(client_public_key_pem.encode(), backend=default_backend())

def decrypt_symmetric_key(encrypted_symmetric_key, private_key):
    return private_key.decrypt(
        encrypted_symmetric_key,
        ec.ECIES(hashes.SHA256())
    )

# Load server/client public and private keys
SERVER_PRIVATE_KEY, SERVER_PUBLIC_KEY = generate_ecc_key_pair()
CONFIG = load_config("config.toml")
CLIENT_PUBLIC_KEY = load_client_public_key(CONFIG)