import hashlib
import os
from typing import Dict, Optional

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ed25519
from cryptography.hazmat.primitives.asymmetric.x25519 import (
    X25519PrivateKey,
    X25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from bitchat_python._logger import logger


class EncryptionError(Exception):
    """Encryption-related errors"""

    pass


class EncryptionService:
    def __init__(self):
        # Generate ephemeral keys
        self.private_key = X25519PrivateKey.generate()
        self.public_key = self.private_key.public_key()

        self.signing_key = ed25519.Ed25519PrivateKey.generate()
        self.verifying_key = self.signing_key.public_key()

        # Generate persistent identity key
        self._identity_key = ed25519.Ed25519PrivateKey.generate()
        self.identity_public = self._identity_key.public_key()

        # Storage
        self.peer_public_keys: Dict[str, X25519PublicKey] = {}
        self.peer_signing_keys: Dict[str, ed25519.Ed25519PublicKey] = {}
        self.peer_identity_keys: Dict[str, ed25519.Ed25519PublicKey] = {}
        self.shared_secrets: Dict[str, bytes] = {}

    def get_combined_public_key_data(self) -> bytes:
        """Get 96-byte combined public key data"""
        data = bytearray()
        # X25519 public key (32 bytes)
        data.extend(
            self.public_key.public_bytes(
                encoding=serialization.Encoding.Raw,
                format=serialization.PublicFormat.Raw,
            )
        )
        # Ed25519 signing key (32 bytes)
        data.extend(
            self.verifying_key.public_bytes(
                encoding=serialization.Encoding.Raw,
                format=serialization.PublicFormat.Raw,
            )
        )
        # Identity key (32 bytes)
        data.extend(
            self.identity_public.public_bytes(
                encoding=serialization.Encoding.Raw,
                format=serialization.PublicFormat.Raw,
            )
        )
        return bytes(data)

    def store_peer_keys(self, peer_id: str, key_data: bytes):
        """Store peer's public keys from key exchange"""
        if len(key_data) < 96:
            raise EncryptionError(
                f"Invalid key data size: {len(key_data)} (expected 96)"
            )

        # Parse X25519 public key (first 32 bytes)
        public_key_bytes = key_data[0:32]
        signing_key_bytes = key_data[32:64]
        identity_key_bytes = key_data[64:96]

        # Parse X25519 key
        try:
            public_key = X25519PublicKey.from_public_bytes(public_key_bytes)
        except Exception:
            raise EncryptionError("Invalid X25519 public key")

        # Parse Ed25519 signing key
        try:
            signing_key = ed25519.Ed25519PublicKey.from_public_bytes(signing_key_bytes)
        except Exception:
            raise EncryptionError("Invalid signing key")

        # Parse identity key with Android compatibility
        try:
            identity_key = ed25519.Ed25519PublicKey.from_public_bytes(
                identity_key_bytes
            )
        except Exception as e:
            # Android bug compatibility
            logger.error(f"[ERROR] Identity key parse error: {e}")
            print(
                f"[CRYPTO] Note: Peer {peer_id} appears to be Android (invalid identity key format)"
            )
            identity_key = signing_key

        # Store keys
        self.peer_public_keys[peer_id] = public_key
        self.peer_signing_keys[peer_id] = signing_key
        self.peer_identity_keys[peer_id] = identity_key

        # Generate shared secret
        shared_secret = self.private_key.exchange(public_key)

        # Derive symmetric key using HKDF
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=b"bitchat-v1",
            info=b"",
            backend=default_backend(),
        )
        symmetric_key = hkdf.derive(shared_secret)

        self.shared_secrets[peer_id] = symmetric_key
        print(f"[CRYPTO] Successfully established shared secret with {peer_id}")

    def get_peer_identity_key(self, peer_id: str) -> Optional[bytes]:
        """Get peer's identity key bytes"""
        if peer_id in self.peer_identity_keys:
            return self.peer_identity_keys[peer_id].public_bytes(
                encoding=serialization.Encoding.Raw,
                format=serialization.PublicFormat.Raw,
            )
        return None

    def get_peer_fingerprint(self, peer_id: str) -> Optional[str]:
        """Get SHA256 fingerprint of peer's identity key"""
        key_bytes = self.get_peer_identity_key(peer_id)
        if key_bytes:
            digest = hashlib.sha256(key_bytes).digest()
            return digest[:16].hex()
        return None

    def encrypt(self, data: bytes, peer_id: str) -> bytes:
        """Encrypt data for a peer"""
        if peer_id not in self.shared_secrets:
            raise EncryptionError("No shared secret")

        aesgcm = AESGCM(self.shared_secrets[peer_id])
        nonce = os.urandom(12)
        ciphertext = aesgcm.encrypt(nonce, data, None)

        return nonce + ciphertext

    def decrypt(self, data: bytes, peer_id: str) -> bytes:
        """Decrypt data from a peer"""
        if len(data) < 12:
            raise EncryptionError("Data too short")

        if peer_id not in self.shared_secrets:
            raise EncryptionError("No shared secret")

        aesgcm = AESGCM(self.shared_secrets[peer_id])
        nonce = data[:12]
        ciphertext = data[12:]

        try:
            plaintext = aesgcm.decrypt(nonce, ciphertext, None)
            return plaintext
        except Exception:
            raise EncryptionError("Decryption failed")

    def sign(self, data: bytes) -> bytes:
        """Sign data"""
        signature = self.signing_key.sign(data)
        return signature

    def verify(self, signature: bytes, data: bytes, peer_id: str) -> bool:
        """Verify signature"""
        if peer_id not in self.peer_signing_keys:
            raise EncryptionError("No signing key")

        try:
            self.peer_signing_keys[peer_id].verify(signature, data)
            return True
        except Exception as e:
            logger.debug(e)
            return False

    @staticmethod
    def derive_channel_key(password: str, channel_name: str) -> bytes:
        """Derive channel key from password"""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=channel_name.encode(),
            iterations=100_000,
            backend=default_backend(),
        )
        return kdf.derive(password.encode())

    @staticmethod
    def encrypt_with_key(data: bytes, key: bytes) -> bytes:
        """Encrypt with a specific key"""
        aesgcm = AESGCM(key)
        nonce = os.urandom(12)
        ciphertext = aesgcm.encrypt(nonce, data, None)
        return nonce + ciphertext

    @staticmethod
    def decrypt_with_key(data: bytes, key: bytes) -> bytes:
        """Decrypt with a specific key"""
        if len(data) < 12:
            raise EncryptionError("Data too short")

        aesgcm = AESGCM(key)
        nonce = data[:12]
        ciphertext = data[12:]

        try:
            plaintext = aesgcm.decrypt(nonce, ciphertext, None)
            return plaintext
        except Exception:
            raise EncryptionError("Decryption failed")

    def has_peer_key(self, peer_id: str) -> bool:
        """Check if we have a peer's key"""
        return peer_id in self.shared_secrets

    def encrypt_for_peer(self, peer_id: str, data: bytes) -> bytes:
        """Encrypt specifically for a peer"""
        return self.encrypt(data, peer_id)


# Export classes and errors
__all__ = ("EncryptionError", "EncryptionService")
