"""FIPS-aware cryptographic primitives for GreyMatter.

Centralizes all cryptographic operations behind a FIPS-compliant interface.
When ``GM_FIPS_MODE=1`` is set, routes operations through the ``cryptography``
library's OpenSSL backend (which supports FIPS 140-3 when the underlying
OpenSSL is compiled with the FIPS provider enabled).

When FIPS mode is disabled or the ``cryptography`` library is unavailable,
falls back to Python stdlib (``os.urandom``, ``hashlib``, ``hmac``) with a
logged warning.

Environment variables:
    GM_FIPS_MODE: Set to ``"1"`` to enable FIPS mode. Default: disabled.

Usage::

    from greymatter_core.security.fips import (
        fips_random_bytes,
        fips_sha256,
        fips_hmac_sha256,
        fips_pbkdf2,
        fips_token_urlsafe,
        fips_token_bytes,
        fips_choice,
        is_fips_mode,
    )

    # Generate 32 random bytes (FIPS-approved RNG when GM_FIPS_MODE=1)
    key = fips_random_bytes(32)

    # SHA-256 hash (returns raw 32-byte digest)
    digest = fips_sha256(b"data")

    # HMAC-SHA256 (returns raw 32-byte MAC)
    mac = fips_hmac_sha256(key, b"message")

    # PBKDF2-HMAC-SHA256 key derivation
    derived = fips_pbkdf2(b"password", b"salt", iterations=100_000)
"""

from __future__ import annotations

import base64
import hashlib
import hmac as _hmac
import logging
import os
import secrets
from typing import Sequence, TypeVar

logger = logging.getLogger(__name__)

T = TypeVar("T")

# ---------------------------------------------------------------------------
# FIPS mode detection
# ---------------------------------------------------------------------------

_FIPS_MODE: bool = os.environ.get("GM_FIPS_MODE", "0") == "1"
_HAS_CRYPTO: bool = False
_FIPS_PROVIDER_AVAILABLE: bool = False

# Attempt to import the cryptography library for FIPS-grade primitives.
# These are only used when _FIPS_MODE is True, but we detect availability
# at import time so we can warn early.
_ossl_backend = None
_crypto_hashes = None
_CryptoHMAC = None
_PBKDF2HMAC = None

try:
    from cryptography.hazmat.primitives import hashes as _crypto_hashes_mod
    from cryptography.hazmat.primitives.hmac import HMAC as _CryptoHMAC_cls
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC as _PBKDF2HMAC_cls
    from cryptography.hazmat.backends.openssl import backend as _ossl_backend_inst

    _HAS_CRYPTO = True
    _crypto_hashes = _crypto_hashes_mod
    _CryptoHMAC = _CryptoHMAC_cls
    _PBKDF2HMAC = _PBKDF2HMAC_cls
    _ossl_backend = _ossl_backend_inst

    # Check if the OpenSSL FIPS provider is actually enabled at the C level.
    _FIPS_PROVIDER_AVAILABLE = getattr(_ossl_backend, "_fips_enabled", False)
except ImportError:
    pass

# Log FIPS mode status at import time for operational visibility.
if _FIPS_MODE:
    if not _HAS_CRYPTO:
        logger.warning(
            "GM_FIPS_MODE=1 but cryptography library is not installed. "
            "Falling back to stdlib -- FIPS compliance is NOT guaranteed. "
            "Install: pip install cryptography"
        )
    elif not _FIPS_PROVIDER_AVAILABLE:
        logger.warning(
            "GM_FIPS_MODE=1 but OpenSSL FIPS provider is not enabled. "
            "Routing through cryptography library's OpenSSL backend, but "
            "full FIPS 140-3 compliance requires an OpenSSL build with "
            "FIPS provider. OpenSSL version: %s",
            _ossl_backend.openssl_version_text(),
        )
    else:
        logger.info(
            "FIPS mode active -- all crypto routed through OpenSSL FIPS provider. "
            "OpenSSL version: %s",
            _ossl_backend.openssl_version_text(),
        )


# ---------------------------------------------------------------------------
# Public helpers
# ---------------------------------------------------------------------------

def is_fips_mode() -> bool:
    """Return True if FIPS mode is requested *and* the cryptography library is available."""
    return _FIPS_MODE and _HAS_CRYPTO


# ---------------------------------------------------------------------------
# Random bytes generation
# ---------------------------------------------------------------------------

def fips_random_bytes(n: int) -> bytes:
    """Generate *n* cryptographically secure random bytes.

    In FIPS mode with the ``cryptography`` library available, generates
    bytes through OpenSSL's ``RAND_bytes`` (which uses the FIPS-approved
    DRBG when the FIPS provider is active).

    Falls back to ``os.urandom()`` otherwise -- still a CSPRNG on all
    modern operating systems, but not a FIPS-validated code path.
    """
    if _FIPS_MODE and _HAS_CRYPTO:
        try:
            # Route through OpenSSL's RAND_bytes via the backend's CFFI
            # bindings.  This ensures the FIPS provider's DRBG is used
            # when the provider is enabled.
            buf = _ossl_backend._ffi.new("unsigned char[]", n)
            rc = _ossl_backend._lib.RAND_bytes(buf, n)
            if rc != 1:
                raise RuntimeError("OpenSSL RAND_bytes returned error")
            return bytes(_ossl_backend._ffi.buffer(buf, n))
        except Exception as exc:
            logger.warning(
                "FIPS RAND_bytes failed, falling back to os.urandom: %s", exc
            )
    return os.urandom(n)


# ---------------------------------------------------------------------------
# SHA-256
# ---------------------------------------------------------------------------

def fips_sha256(data: bytes) -> bytes:
    """Compute SHA-256 digest of *data*.  Returns the raw 32-byte hash.

    In FIPS mode, uses the ``cryptography`` library (OpenSSL EVP_SHA256).
    Falls back to ``hashlib.sha256`` otherwise.

    Call ``.hex()`` on the result if you need a hex string.
    """
    if _FIPS_MODE and _HAS_CRYPTO:
        try:
            h = _crypto_hashes.Hash(_crypto_hashes.SHA256())
            h.update(data)
            return h.finalize()
        except Exception as exc:
            logger.warning(
                "FIPS SHA-256 failed, falling back to hashlib: %s", exc
            )
    return hashlib.sha256(data).digest()


# ---------------------------------------------------------------------------
# HMAC-SHA256
# ---------------------------------------------------------------------------

def fips_hmac_sha256(key: bytes, data: bytes) -> bytes:
    """Compute HMAC-SHA256 of *data* with *key*.  Returns the raw 32-byte MAC.

    In FIPS mode, uses the ``cryptography`` library (OpenSSL HMAC).
    Falls back to ``hmac.new(key, data, hashlib.sha256)`` otherwise.

    Call ``.hex()`` on the result if you need a hex string.
    """
    if _FIPS_MODE and _HAS_CRYPTO:
        try:
            h = _CryptoHMAC(key, _crypto_hashes.SHA256())
            h.update(data)
            return h.finalize()
        except Exception as exc:
            logger.warning(
                "FIPS HMAC-SHA256 failed, falling back to stdlib hmac: %s", exc
            )
    return _hmac.new(key, data, hashlib.sha256).digest()


# ---------------------------------------------------------------------------
# PBKDF2-HMAC-SHA256
# ---------------------------------------------------------------------------

def fips_pbkdf2(
    password: bytes,
    salt: bytes,
    iterations: int,
    dklen: int = 32,
) -> bytes:
    """Derive a key using PBKDF2-HMAC-SHA256.

    In FIPS mode, uses the ``cryptography`` library (OpenSSL PKCS5_PBKDF2).
    Falls back to ``hashlib.pbkdf2_hmac`` otherwise.

    Args:
        password: Key material (raw bytes).
        salt: Unique salt per derivation.
        iterations: PBKDF2 iteration count (OWASP recommends >= 600,000 for
            password hashing; lower counts are acceptable for non-password KDF).
        dklen: Desired derived key length in bytes (default 32 = AES-256).

    Returns:
        Derived key bytes of length *dklen*.
    """
    if _FIPS_MODE and _HAS_CRYPTO:
        try:
            kdf = _PBKDF2HMAC(
                algorithm=_crypto_hashes.SHA256(),
                length=dklen,
                salt=salt,
                iterations=iterations,
            )
            return kdf.derive(password)
        except Exception as exc:
            logger.warning(
                "FIPS PBKDF2 failed, falling back to hashlib: %s", exc
            )
    return hashlib.pbkdf2_hmac("sha256", password, salt, iterations, dklen=dklen)


# ---------------------------------------------------------------------------
# Token generation helpers (drop-in replacements for secrets.* functions)
# ---------------------------------------------------------------------------

def fips_token_urlsafe(nbytes: int = 32) -> str:
    """Generate a URL-safe random token string.

    FIPS-aware replacement for ``secrets.token_urlsafe()``.
    Returns a base64url-encoded string (without padding) from *nbytes*
    random bytes.
    """
    return base64.urlsafe_b64encode(fips_random_bytes(nbytes)).rstrip(b"=").decode("ascii")


def fips_token_bytes(nbytes: int = 32) -> bytes:
    """Generate *nbytes* random bytes for use as a token.

    FIPS-aware replacement for ``secrets.token_bytes()``.
    """
    return fips_random_bytes(nbytes)


def fips_choice(seq: Sequence[T]) -> T:
    """Choose a random element from a non-empty sequence.

    FIPS-aware replacement for ``secrets.choice()``.  In FIPS mode, uses
    ``fips_random_bytes()`` for index generation; otherwise delegates to
    ``secrets.choice()``.

    Raises:
        IndexError: If *seq* is empty.
    """
    if not seq:
        raise IndexError("Cannot choose from an empty sequence")
    if _FIPS_MODE and _HAS_CRYPTO:
        # 4 bytes of randomness -> 2^32 possible values.  For typical
        # alphabet sizes (< 100), the modulo bias is < 2^-24, which is
        # negligible for ID generation.
        index = int.from_bytes(fips_random_bytes(4), "big") % len(seq)
        return seq[index]
    return secrets.choice(seq)
