import os
import secrets
import hashlib
import time
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import dh
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import padding as sym_padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.hmac import HMAC as HMAC_PRIMITIVE
from doubleratchet import DoubleRatchet, DiffieHellmanRatchet, KDFChain, MessageKey, AEAD, Header

class FantomXProtocol:
    def __init__(self):
        self.backend = default_backend()

        # Generate long-term DH private/public key pair
        self.private_key = dh.generate_parameters(generator=2, key_size=2048, backend=self.backend).generate_private_key()
        self.public_key = self.private_key.public_key()

        self.current_session_keys = {}

        # Store previously seen replay protection tokens
        self.replay_protection_store = {}

        # Initialize Double Ratchet
        self.dh_ratchet = DiffieHellmanRatchet(self.private_key)
        self.ratchet = DoubleRatchet(self.dh_ratchet, KDFChain, MessageKey, AEAD)

    def rotate_session_keys(self, recipient_public_key):
        shared_key = self.private_key.exchange(recipient_public_key)

        derived_key = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b'session data',
            backend=self.backend
        ).derive(shared_key)

        self.current_session_keys = {
            'derived_key': derived_key
        }

        return self.public_key

    def generate_replay_protection_token(self, message, recipient_public_key):
        """Generates a token to protect against message replay, using timestamp and nonce."""
        timestamp = int(time.time())
        nonce = secrets.token_bytes(16)  # Generate a 16-byte random nonce
        message_hash = hashlib.sha256(message.encode()).hexdigest()

        token = {
            'message_hash': message_hash,
            'timestamp': timestamp,
            'nonce': nonce.hex()
        }

        recipient_key = recipient_public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        ).hex()
        
        if recipient_key not in self.replay_protection_store:
            self.replay_protection_store[recipient_key] = []

        self.replay_protection_store[recipient_key].append(token)

        return token

    def verify_replay_protection_token(self, message, token, recipient_public_key):
        """Verifies if a message has been replayed using the stored token data."""
        recipient_key = recipient_public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        ).hex()

        stored_tokens = self.replay_protection_store.get(recipient_key, [])

        message_hash = hashlib.sha256(message.encode()).hexdigest()

        for stored_token in stored_tokens:
            if stored_token['message_hash'] == message_hash and stored_token['nonce'] == token['nonce']:
                # If we find the same message hash with the same nonce, this is a replay
                return False

        # If token is not found or the nonce is different, this message is new
        return True

    def encrypt(self, content, recipient_public_key):
        """Encrypts a message or file content for the recipient with forward secrecy."""
        ephemeral_public_key = self.rotate_session_keys(recipient_public_key)

        # Generate a nonce (IV) for GCM mode
        nonce = os.urandom(12)

        # Create the cipher with the nonce
        cipher = Cipher(
            algorithms.AES(self.current_session_keys['derived_key']),
            modes.GCM(nonce),
            backend=self.backend
        )

        encryptor = cipher.encryptor()

        # If the content is a string, apply padding for AES encryption
        if isinstance(content, str):  # Treat it as text content if it's a string
            padder = sym_padding.PKCS7(128).padder()
            padded_content = padder.update(content.encode()) + padder.finalize()
        else:  # If it's binary data (like a file), we don't need padding
            padded_content = content

        # Encrypt the padded content
        ciphertext = encryptor.update(padded_content) + encryptor.finalize()

        # Generate HMAC for message integrity
        hmac = HMAC_PRIMITIVE(self.current_session_keys['derived_key'], hashes.SHA256(), backend=self.backend)
        hmac.update(ciphertext)
        hmac_value = hmac.finalize()

        # Generate the replay protection token
        token = self.generate_replay_protection_token(content, recipient_public_key)

        # Return serialized content with nonce, tag, ciphertext in hex, replay protection token, and HMAC
        encrypted_content = {
            'ephemeral_public_key': ephemeral_public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            ).hex(),  # Convert ephemeral public key to hex
            'nonce': nonce.hex(),  # Convert nonce to hex
            'ciphertext': ciphertext.hex(),  # Convert ciphertext to hex
            'tag': encryptor.tag.hex(),  # Convert tag to hex
            'replay_token': token,  # Include the replay protection token
            'hmac': hmac_value.hex()  # Include the HMAC
        }

        return encrypted_content

    def decrypt(self, encrypted_content, private_key):
        """Decrypts a message or file content and ensures message integrity and freshness."""
        ephemeral_public_key = serialization.load_pem_public_key(
            bytes.fromhex(encrypted_content['ephemeral_public_key']),
            backend=self.backend
        )
        shared_key = private_key.exchange(ephemeral_public_key)

        derived_key = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b'session data',
            backend=self.backend
        ).derive(shared_key)

        cipher = Cipher(
            algorithms.AES(derived_key),
            modes.GCM(
                bytes.fromhex(encrypted_content['nonce']),
                bytes.fromhex(encrypted_content['tag'])
            ),
            backend=self.backend
        )
        decryptor = cipher.decryptor()
        decrypted_padded_content = decryptor.update(bytes.fromhex(encrypted_content['ciphertext'])) + decryptor.finalize()

        # Unpad the content if it's text (for file, we don't need unpadding)
        try:
            unpadder = sym_padding.PKCS7(128).unpadder()
            decrypted_content = unpadder.update(decrypted_padded_content) + unpadder.finalize()

            # Verify the replay protection token
            token = encrypted_content['replay_token']
            if not self.verify_replay_protection_token(decrypted_content.decode(), token, private_key.public_key()):
                raise Exception("Replay detected. The message has been replayed.")

            # Verify the HMAC
            hmac = HMAC_PRIMITIVE(derived_key, hashes.SHA256(), backend=self.backend)
            hmac.update(bytes.fromhex(encrypted_content['ciphertext']))
            hmac.verify(bytes.fromhex(encrypted_content['hmac']))

            return decrypted_content.decode()  # Return decrypted content as text
        except ValueError:
            # If unpadding fails, treat as binary file content
            decrypted_content = decrypted_padded_content
            
            # Verify the replay protection token
            token = encrypted_content['replay_token']
            if not self.verify_replay_protection_token(decrypted_content, token, private_key.public_key()):
                raise Exception("Replay detected. The message has been replayed.")

            # Verify the HMAC
            hmac = HMAC_PRIMITIVE(derived_key, hashes.SHA256(), backend=self.backend)
            hmac.update(bytes.fromhex(encrypted_content['ciphertext']))
            hmac.verify(bytes.fromhex(encrypted_content['hmac']))

            return decrypted_content  # Return binary content

    def create_user_key_pair(self, password):
        """Generates a user-specific key pair based on a password."""
        salt = os.urandom(16)
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=self.backend
        )
        key = kdf.derive(password.encode())
        private_key = dh.generate_parameters(generator=2, key_size=2048, backend=self.backend).generate_private_key()
        public_key = private_key.public_key()
        return private_key, public_key, key, salt

    def encrypt_user_private_key(self, private_key, key, salt):
        """Encrypts a user private key using AES."""
        cipher = Cipher(
            algorithms.AES(key),
            modes.GCM(os.urandom(12)),
            backend=self.backend
        )
        encryptor = cipher.encryptor()
        private_key_bytes = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        ciphertext = encryptor.update(private_key_bytes) + encryptor.finalize()
        return {
            'ciphertext': ciphertext.hex(),
            'nonce': cipher.algorithm.mode.nonce.hex(),
            'tag': encryptor.tag.hex(),
            'salt': salt.hex()
        }

    def decrypt_user_private_key(self, encrypted_private_key, password, salt):
        """Decrypts a user private key using AES."""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=self.backend
        )
        key = kdf.derive(password.encode())
        cipher = Cipher(
            algorithms.AES(key),
            modes.GCM(
                bytes.fromhex(encrypted_private_key['nonce']),
                bytes.fromhex(encrypted_private_key['tag'])
            ),
            backend=self.backend
        )
        decryptor = cipher.decryptor()
        private_key_bytes = decryptor.update(bytes.fromhex(encrypted_private_key['ciphertext'])) + decryptor.finalize()
        return serialization.load_pem_private_key(private_key_bytes, password=None, backend=self.backend)

    def hide_ip(self, message, ip_address):
        """Adds an IP hash to a message for privacy."""
        return f"{message} [IP: {hashlib.sha256(ip_address.encode()).hexdigest()}]"

    def verify_ip_hash(self, message, ip_hash, ip_address):
        """Verifies the IP hash for the message."""
        return hashlib.sha256(ip_address.encode()).hexdigest() == ip_hash

    # Methods for socket integration
    def initiate_key_exchange(self):
        """Initiate key exchange and return the public key to be sent to the recipient."""
        return self.public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )

    def receive_key_exchange(self, recipient_public_key_bytes):
        """Receive recipient's public key and rotate session keys."""
        recipient_public_key = serialization.load_pem_public_key(
            recipient_public_key_bytes,
            backend=self.backend
        )
        self.rotate_session_keys(recipient_public_key)

    def prepare_message(self, message, recipient_public_key_bytes):
        """Prepare an encrypted message for sending via socket."""
        recipient_public_key = serialization.load_pem_public_key(
            recipient_public_key_bytes,
            backend=self.backend
        )
        return self.encrypt(message, recipient_public_key)

    def process_received_message(self, encrypted_content):
        """Process a received encrypted message."""
        return self.decrypt(encrypted_content, self.private_key)

    def decrypt_previous_messages(self, encrypted_contents, private_key):
        """Decrypts a list of encrypted messages, allowing for decryption of previous messages despite key changes."""
        decrypted_messages = []
        for encrypted_content in encrypted_contents:
            try:
                decrypted_message = self.decrypt(encrypted_content, private_key)
                decrypted_messages.append(decrypted_message)
            except Exception as e:
                print(f"Failed to decrypt message: {e}")
                continue
        return decrypted_messages
    
