"""
Primitives de chiffrement : ElGamal, chiffrement homomorphe, AES.
"""

from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
import os
from typing import Tuple
from dataclasses import dataclass
import json


@dataclass
class PublicKey:
    """Clé publique ElGamal"""
    p: int  # Nombre premier
    g: int  # Générateur du groupe
    h: int  # Clé publique = g^x mod p


@dataclass
class PrivateKey:
    """Clé privée ElGamal"""
    x: int  # Clé privée


@dataclass
class Ciphertext:
    """Texte chiffré ElGamal"""
    c1: int  # c1 = g^r mod p
    c2: int  # c2 = m * h^r mod p


class ElGamalEncryption:
    """
    Chiffrement ElGamal - Fondamental pour le vote électronique.
    Propriétés:
    - Sémantiquement sûr (IND-CPA)
    - Chiffrement probabiliste
    - Support pour preuves ZK
    """

    def __init__(self, p: int = None, g: int = None):
        """
        Initialiser ElGamal avec paramètres de groupe.
        Pour le prototype, utilise des paramètres de test.
        """
        if p is None:
            # Nombres premiers de test (petits, pour prototype)
            # En production: nombres premiers cryptographiques forts (2048+ bits)
            self.p = 23  # Nombre premier
            self.g = 5   # Générateur
        else:
            self.p = p
            self.g = g

    def generate_keypair(self) -> Tuple[PublicKey, PrivateKey]:
        """Générer une paire de clés ElGamal"""
        import random
        x = random.randint(2, self.p - 2)  # Clé privée
        h = pow(self.g, x, self.p)  # Clé publique: g^x mod p
        
        public = PublicKey(p=self.p, g=self.g, h=h)
        private = PrivateKey(x=x)
        
        return public, private

    def encrypt(self, public_key: PublicKey, message: int) -> Ciphertext:
        """
        Chiffrer un message avec ElGamal.
        message: nombre entre 0 et p-1
        """
        import random
        r = random.randint(2, self.p - 2)  # Aléa
        
        c1 = pow(self.g, r, self.p)  # c1 = g^r mod p
        c2 = (message * pow(public_key.h, r, self.p)) % self.p  # c2 = m * h^r mod p
        
        return Ciphertext(c1=c1, c2=c2)

    def decrypt(self, private_key: PrivateKey, ciphertext: Ciphertext, p: int) -> int:
        """Déchiffrer un message ElGamal"""
        # m = c2 / c1^x mod p = c2 * (c1^x)^(-1) mod p
        shared_secret = pow(ciphertext.c1, private_key.x, p)
        shared_secret_inv = pow(shared_secret, -1, p)
        message = (ciphertext.c2 * shared_secret_inv) % p
        
        return message

    def add_ciphertexts(self, ct1: Ciphertext, ct2: Ciphertext, p: int) -> Ciphertext:
        """
        Addition homomorphe : E(m1) * E(m2) = E(m1 + m2)
        Propriété clé pour les dépouillement sécurisé
        """
        c1_sum = (ct1.c1 * ct2.c1) % p
        c2_sum = (ct1.c2 * ct2.c2) % p
        return Ciphertext(c1=c1_sum, c2=c2_sum)


class HomomorphicEncryption:
    """
    Chiffrement homomorphe - Paillier-like pour vote.
    Support l'addition sans déchiffrement.
    """

    def __init__(self):
        self.elgamal = ElGamalEncryption()

    def sum_encrypted_votes(self, ciphertexts: list[Ciphertext], p: int) -> Ciphertext:
        """
        Additionner les votes chiffrés sans les déchiffrer.
        C'est la base du dépouillement sécurisé.
        """
        result = ciphertexts[0]
        for ct in ciphertexts[1:]:
            result = self.elgamal.add_ciphertexts(result, ct, p)
        return result


class SymmetricEncryption:
    """
    Chiffrement symétrique AES-256 pour données sensibles.
    """

    @staticmethod
    def generate_key() -> bytes:
        """Générer une clé AES-256 aléatoire"""
        return os.urandom(32)

    @staticmethod
    def encrypt(key: bytes, plaintext: bytes) -> bytes:
        """Chiffrer avec AES-256-GCM"""
        iv = os.urandom(16)
        cipher = Cipher(
            algorithms.AES(key),
            modes.GCM(iv),
            backend=default_backend()
        )
        encryptor = cipher.encryptor()
        ciphertext = encryptor.update(plaintext) + encryptor.finalize()
        
        # Retourner IV || tag || ciphertext
        return iv + encryptor.tag + ciphertext

    @staticmethod
    def decrypt(key: bytes, encrypted_data: bytes) -> bytes:
        """Déchiffrer AES-256-GCM"""
        iv = encrypted_data[:16]
        tag = encrypted_data[16:32]
        ciphertext = encrypted_data[32:]
        
        cipher = Cipher(
            algorithms.AES(key),
            modes.GCM(iv, tag),
            backend=default_backend()
        )
        decryptor = cipher.decryptor()
        plaintext = decryptor.update(ciphertext) + decryptor.finalize()
        
        return plaintext