"""Layer 3: Encryption at rest — SQLite database encryption.
Layer 4: Sync payload encryption — AES-256-GCM for sync deltas.

Layer 3 uses SQLCipher when available, falling back to application-level
field encryption for sensitive columns.

Layer 4 encrypts sync batches before transmission using AES-256-GCM
with HMAC integrity verification. Supports ML-KEM-768 key encapsulation
for post-quantum forward secrecy (version 2 payloads).

All cryptographic primitives (random bytes, HMAC-SHA256, PBKDF2 key
derivation) are sourced from the FIPS-aware wrapper
(:mod:`greymatter_core.security.fips`).  When ``GM_FIPS_MODE=1``, every
operation routes through the OpenSSL FIPS provider.  See
``docs/fips-crypto.md`` for the full API reference and deployment guide.
"""
from __future__ import annotations

import hmac
import json
import os
import base64
import logging
from typing import Optional

from greymatter_core.security.fips import (
    fips_hmac_sha256,
    fips_pbkdf2,
    fips_random_bytes,
)

logger = logging.getLogger(__name__)

# ML-KEM availability check
try:
    from kyber_py.ml_kem import ML_KEM_768
    MLKEM_AVAILABLE = True
except ImportError:
    MLKEM_AVAILABLE = False


def _derive_key(secret: bytes, salt: bytes, length: int = 32) -> bytes:
    """Derive an encryption key using PBKDF2-HMAC-SHA256.

    Args:
        secret: Base secret material (cluster secret or node cert fingerprint).
        salt: Unique salt per database/payload.
        length: Desired key length in bytes (default 32 = AES-256).
    """
    return fips_pbkdf2(secret, salt, iterations=100_000, dklen=length)


# =========================================================================
# Layer 3: Encryption at Rest
# =========================================================================

def get_sqlcipher_key(node_id: str, secret: Optional[bytes] = None) -> str:
    """Generate a SQLCipher encryption key for a node's local database.

    The key is derived from the cluster secret + node_id, so each node
    has a unique database encryption key. If the cluster secret is
    compromised, all local databases are at risk — but an attacker who
    only gets one machine's disk can't read the DB without the secret.

    Args:
        node_id: The node's unique identifier.
        secret: Cluster secret. Uses GM_CLUSTER_SECRET if not provided.
    """
    if secret is None:
        env_secret = os.environ.get("GM_CLUSTER_SECRET", "")
        if not env_secret:
            from greymatter_core.security.tokens import _get_cluster_secret
            secret = _get_cluster_secret()
        else:
            secret = env_secret.encode()

    salt = f"gm-sqlcipher-{node_id}".encode()
    key = _derive_key(secret, salt, length=32)
    return key.hex()


def open_encrypted_sqlite(db_path: str, key: str):
    """Open a SQLCipher-encrypted database.

    Falls back to standard sqlite3 if SQLCipher (pysqlcipher3) is
    not available. In that case, a warning is logged but the DB
    opens unencrypted.

    Args:
        db_path: Path to the SQLite database file.
        key: Hex-encoded encryption key.

    Returns:
        sqlite3.Connection (or pysqlcipher3 connection).
    """
    try:
        from pysqlcipher3 import dbapi2 as sqlcipher
        conn = sqlcipher.connect(db_path)
        conn.execute(f"PRAGMA key = \"x'{key}'\";")
        conn.execute("PRAGMA cipher_page_size = 4096;")
        conn.execute("PRAGMA kdf_iter = 256000;")
        conn.execute("PRAGMA cipher_hmac_algorithm = HMAC_SHA256;")
        # Verify the key works
        conn.execute("SELECT count(*) FROM sqlite_master;")
        return conn
    except ImportError:
        import logging
        logging.getLogger(__name__).warning(
            "SQLCipher not available — database will NOT be encrypted at rest. "
            "Install pysqlcipher3 for encryption: pip install pysqlcipher3"
        )
        import sqlite3
        return sqlite3.connect(db_path)


# =========================================================================
# ML-KEM Key Management
# =========================================================================

def generate_mlkem_keypair() -> tuple[bytes, bytes]:
    """Generate an ML-KEM-768 key pair for this node.

    Returns:
        Tuple of (encapsulation_key, decapsulation_key).
        encapsulation_key is public (1184 bytes), decapsulation_key is private (2400 bytes).
    """
    if not MLKEM_AVAILABLE:
        raise ImportError("kyber-py not installed. Install: pip install kyber-py")
    return ML_KEM_768.keygen()


def save_mlkem_keys(ek: bytes, dk: bytes, keys_dir: str | None = None) -> None:
    """Save ML-KEM keys to disk with restricted permissions."""
    from pathlib import Path
    keys_dir = Path(keys_dir or os.environ.get(
        "GM_CERTS_DIR",
        str(Path.home() / ".greymatter" / "certs"),
    ))
    keys_dir.mkdir(parents=True, exist_ok=True)

    ek_path = keys_dir / "mlkem.ek"
    dk_path = keys_dir / "mlkem.dk"

    ek_path.write_bytes(ek)
    dk_path.write_bytes(dk)
    dk_path.chmod(0o600)
    logger.info("ML-KEM-768 keys saved to %s", keys_dir)


def load_mlkem_keys(keys_dir: str | None = None) -> tuple[bytes, bytes]:
    """Load ML-KEM keys from disk.

    Returns:
        Tuple of (encapsulation_key, decapsulation_key).

    Raises:
        FileNotFoundError: If keys don't exist (call generate_mlkem_keypair first).
    """
    from pathlib import Path
    keys_dir = Path(keys_dir or os.environ.get(
        "GM_CERTS_DIR",
        str(Path.home() / ".greymatter" / "certs"),
    ))

    ek = (keys_dir / "mlkem.ek").read_bytes()
    dk = (keys_dir / "mlkem.dk").read_bytes()
    return ek, dk


# =========================================================================
# Layer 4: Sync Payload Encryption
# =========================================================================

def encrypt_sync_payload(
    payload: dict,
    secret: Optional[bytes] = None,
    peer_ek: Optional[bytes] = None,
) -> dict:
    """Encrypt a sync payload for transmission.

    Version selection:
        v2: ML-KEM-768 key encapsulation + AES-256-GCM (if peer_ek provided + kyber-py available)
        v1: AES-256-GCM with derived key (if cryptography available)
        v0: HMAC-only integrity (fallback, no confidentiality)

    Args:
        payload: Dict to encrypt (will be JSON-serialized).
        secret: Encryption key material. Uses cluster secret if not provided.
        peer_ek: Peer's ML-KEM encapsulation key (public key). If provided,
            uses ML-KEM key encapsulation for post-quantum forward secrecy.

    Returns:
        Dict with version-specific keys including ct, hmac, and v.
    """
    if secret is None:
        from greymatter_core.security.tokens import _get_cluster_secret
        secret = _get_cluster_secret()

    plaintext = json.dumps(payload, separators=(",", ":")).encode()

    # Version 2: ML-KEM-768 + AES-256-GCM (post-quantum forward secrecy)
    if peer_ek is not None and MLKEM_AVAILABLE:
        try:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM

            # Encapsulate: generates ephemeral shared secret + ciphertext
            shared_secret, kem_ct = ML_KEM_768.encaps(peer_ek)

            # Use shared secret as AES key (it's already 32 bytes)
            nonce = fips_random_bytes(12)
            aesgcm = AESGCM(shared_secret)
            ciphertext = aesgcm.encrypt(nonce, plaintext, None)

            # HMAC over everything: KEM ciphertext + nonce + AES ciphertext
            hmac_key = _derive_key(secret, b"gm-sync-mlkem-hmac", length=32)
            hmac_data = kem_ct + nonce + ciphertext
            payload_hmac = fips_hmac_sha256(hmac_key, hmac_data).hex()

            return {
                "v": 2,
                "ct": base64.b64encode(ciphertext).decode(),
                "nonce": base64.b64encode(nonce).decode(),
                "kem_ct": base64.b64encode(kem_ct).decode(),
                "hmac": payload_hmac,
            }
        except ImportError:
            logger.warning("cryptography not available, falling back to v1")

    # Version 1: AES-256-GCM with derived key
    sync_key = _derive_key(secret, b"gm-sync-payload", length=32)

    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        nonce = fips_random_bytes(12)
        aesgcm = AESGCM(sync_key)
        ciphertext = aesgcm.encrypt(nonce, plaintext, None)

        hmac_data = nonce + ciphertext
        payload_hmac = fips_hmac_sha256(sync_key, hmac_data).hex()

        return {
            "v": 1,
            "ct": base64.b64encode(ciphertext).decode(),
            "nonce": base64.b64encode(nonce).decode(),
            "hmac": payload_hmac,
        }
    except ImportError:
        # Fallback: HMAC-only integrity (no confidentiality)
        logger.warning(
            "cryptography library not available — sync payloads will NOT be "
            "encrypted. Install cryptography for encryption: pip install cryptography"
        )
        payload_bytes = base64.b64encode(plaintext).decode()
        payload_hmac = fips_hmac_sha256(sync_key, plaintext).hex()
        return {
            "v": 0,
            "ct": payload_bytes,
            "hmac": payload_hmac,
        }


def decrypt_sync_payload(
    encrypted: dict,
    secret: Optional[bytes] = None,
    dk: Optional[bytes] = None,
) -> dict:
    """Decrypt a sync payload received from another node.

    Verifies HMAC integrity before decryption.

    Args:
        encrypted: Dict with ct, nonce, hmac, v keys.
        secret: Decryption key material. Uses cluster secret if not provided.
        dk: ML-KEM decapsulation key (private key). Required for v2 payloads.

    Returns:
        Decrypted dict.

    Raises:
        ValueError: If HMAC verification fails (tampered payload).
        ValueError: If decryption fails (wrong key or corrupted data).
    """
    if secret is None:
        from greymatter_core.security.tokens import _get_cluster_secret
        secret = _get_cluster_secret()

    version = encrypted.get("v", 0)

    if version == 2:
        # ML-KEM-768 + AES-256-GCM (post-quantum forward secrecy)
        if not MLKEM_AVAILABLE:
            raise ImportError("kyber-py required for v2 decryption")
        if dk is None:
            raise ValueError("ML-KEM decapsulation key (dk) required for v2 payloads")

        from cryptography.hazmat.primitives.ciphers.aead import AESGCM

        ciphertext = base64.b64decode(encrypted["ct"])
        nonce = base64.b64decode(encrypted["nonce"])
        kem_ct = base64.b64decode(encrypted["kem_ct"])

        # Verify HMAC first
        hmac_key = _derive_key(secret, b"gm-sync-mlkem-hmac", length=32)
        hmac_data = kem_ct + nonce + ciphertext
        expected_hmac = fips_hmac_sha256(hmac_key, hmac_data).hex()
        if not hmac.compare_digest(expected_hmac, encrypted["hmac"]):
            raise ValueError("HMAC verification failed — payload may be tampered")

        # Decapsulate to recover the shared secret (AES key)
        shared_secret = ML_KEM_768.decaps(dk, kem_ct)
        aesgcm = AESGCM(shared_secret)
        plaintext = aesgcm.decrypt(nonce, ciphertext, None)
        return json.loads(plaintext)

    elif version == 1:
        # Full encryption (AES-256-GCM)
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM

        sync_key = _derive_key(secret, b"gm-sync-payload", length=32)
        ciphertext = base64.b64decode(encrypted["ct"])
        nonce = base64.b64decode(encrypted["nonce"])

        hmac_data = nonce + ciphertext
        expected_hmac = fips_hmac_sha256(sync_key, hmac_data).hex()
        if not hmac.compare_digest(expected_hmac, encrypted["hmac"]):
            raise ValueError("HMAC verification failed — payload may be tampered")

        aesgcm = AESGCM(sync_key)
        plaintext = aesgcm.decrypt(nonce, ciphertext, None)
        return json.loads(plaintext)

    elif version == 0:
        # Integrity-only (no encryption)
        sync_key = _derive_key(secret, b"gm-sync-payload", length=32)
        plaintext = base64.b64decode(encrypted["ct"])
        expected_hmac = fips_hmac_sha256(sync_key, plaintext).hex()
        if not hmac.compare_digest(expected_hmac, encrypted["hmac"]):
            raise ValueError("HMAC verification failed — payload may be tampered")
        return json.loads(plaintext)

    else:
        raise ValueError(f"Unknown encryption version: {version}")
