"""
TOTP (Time-based One-Time Password, RFC 6238) helpers.

Used as optional MFA for local-auth users. SSO users get MFA from their IdP
and should not enable TOTP here (UI hides the setup button when
auth_provider != local).

Storage: `users.totp_secret` holds the base32 secret encrypted with Fernet.
Key resolution order:
  1. `AUTH_PROVIDER_CRYPTO_KEY` env var (recommended for production — key
     stays outside the encrypted-data store)
  2. `auth_provider_crypto_key` row in `settings` table (auto-generated on
     first use if neither is present; survives container restarts)

We never log the secret. We never log the key.
"""
from __future__ import annotations

import logging
import os
import threading
import time
from typing import Optional

import pyotp
from cryptography.fernet import Fernet, InvalidToken

logger = logging.getLogger(__name__)

ISSUER = "VulnCheck"
TOTP_DIGITS = 6
TOTP_INTERVAL = 30
TOTP_VALID_WINDOW = 1  # accept current ± 1 step to tolerate clock skew

SETTINGS_KEY_CRYPTO = "auth_provider_crypto_key"

# Cache the resolved key for the process lifetime. Re-resolved on each restart.
_cached_key: Optional[bytes] = None

# Replay-protection cache: (user_id, code) -> expiry epoch. Single-worker
# in-memory store; for multi-worker deployments, swap this for Redis with
# the same TTL semantics. Window of (2*VALID_WINDOW+1)*INTERVAL = 90s.
_REPLAY_TTL_SEC = (2 * TOTP_VALID_WINDOW + 1) * TOTP_INTERVAL
_replay_cache: dict[tuple, float] = {}
_replay_lock = threading.Lock()


def _replay_seen(user_id, code: str) -> bool:
    """Return True if (user_id, code) was already accepted recently."""
    now = time.monotonic()
    with _replay_lock:
        # Lazy purge of expired entries (bounded — TOTP cardinality is small).
        expired = [k for k, exp in _replay_cache.items() if exp <= now]
        for k in expired:
            _replay_cache.pop(k, None)
        return (user_id, code) in _replay_cache


def _replay_record(user_id, code: str) -> None:
    with _replay_lock:
        _replay_cache[(user_id, code)] = time.monotonic() + _REPLAY_TTL_SEC


def _resolve_key() -> bytes:
    """
    Resolve the Fernet key from AUTH_PROVIDER_CRYPTO_KEY env var.

    Refuses to fall back to DB storage — co-locating the key with the
    encrypted ciphertext defeats the purpose of encrypting at rest.
    """
    global _cached_key
    if _cached_key is not None:
        return _cached_key

    env_key = os.getenv("AUTH_PROVIDER_CRYPTO_KEY")
    if not env_key:
        raise RuntimeError(
            "AUTH_PROVIDER_CRYPTO_KEY environment variable is not set. "
            "Generate one with: python -c 'from cryptography.fernet import "
            "Fernet; print(Fernet.generate_key().decode())' and set it in "
            "the environment. Required for TOTP secret encryption and "
            "LDAP bind password encryption."
        )

    key_bytes = env_key.encode() if isinstance(env_key, str) else env_key
    try:
        Fernet(key_bytes)
    except Exception as e:
        raise RuntimeError(
            f"AUTH_PROVIDER_CRYPTO_KEY is not a valid Fernet key: {e}"
        )

    _cached_key = key_bytes
    return _cached_key


def _fernet() -> Fernet:
    """Build Fernet instance from the resolved key (env > DB > auto-generated)."""
    return Fernet(_resolve_key())


def generate_secret() -> str:
    """Generate a new base32 TOTP secret (160-bit)."""
    return pyotp.random_base32()


def encrypt_secret(plain_secret: str) -> str:
    return _fernet().encrypt(plain_secret.encode()).decode()


# A base32-encoded TOTP secret (the legacy plain format) is 32 chars
# of [A-Z2-7] when generated by pyotp.random_base32. Padding is rare
# but we accept '=' for safety. A Fernet ciphertext starts with the
# version byte 'gAAAAA' after URL-safe base64 — its first char is 'g'.
import re as _re
_BASE32_RE = _re.compile(r"^[A-Z2-7=]{16,64}$")


def _looks_like_plain_base32(value: str) -> bool:
    return bool(_BASE32_RE.match(value or ""))


def decrypt_secret(stored: str) -> str:
    """
    Decrypt a stored TOTP secret, with a one-shot migration fallback
    for users enrolled BEFORE this Fernet-encryption layer landed.

    Resolution order:
    1. Try Fernet decryption (the standard, post-migration path).
    2. On InvalidToken, check if the stored value looks like a plain
       base32 secret (the legacy on-disk format). If so, return it as
       plain so login still works. Callers should call ``encrypt_secret``
       on the returned value and write it back so the next read goes
       through the modern path — see ``app/routers/auth.py`` mfa flows.
    3. If neither succeeds, log and re-raise so the user is rejected.

    Without this fallback, every operator who had MFA enabled before
    the AUTH_PROVIDER_CRYPTO_KEY requirement landed would see "Invalid
    MFA Code" on the next login and have to reset enrolment, even
    though their authenticator app + DB row still hold the correct
    secret. This is the bug the tester hit during the v0.12 upgrade.
    """
    if not stored:
        raise InvalidToken("empty stored secret")
    try:
        return _fernet().decrypt(stored.encode()).decode()
    except InvalidToken as e:
        if _looks_like_plain_base32(stored):
            logger.warning(
                "TOTP secret for user appears to be stored in legacy "
                "plain base32 format — accepting for login and flagging "
                "for re-encryption on next save."
            )
            return stored
        # Corrupted secret or wrong key. Log generic message, never the value.
        logger.error("TOTP secret decryption failed: %s", e)
        raise


def provisioning_uri(username: str, plain_secret: str) -> str:
    """Generate the otpauth:// URI for QR-code rendering in the UI."""
    return pyotp.totp.TOTP(plain_secret).provisioning_uri(
        name=username, issuer_name=ISSUER
    )


def verify(plain_secret: str, code: str) -> bool:
    """Verify a 6-digit OTP code against the secret with skew tolerance."""
    if not code or not plain_secret:
        return False
    code = code.strip().replace(" ", "")
    if len(code) != TOTP_DIGITS or not code.isdigit():
        return False
    totp = pyotp.TOTP(plain_secret, digits=TOTP_DIGITS, interval=TOTP_INTERVAL)
    return totp.verify(code, valid_window=TOTP_VALID_WINDOW)


def verify_for_user(user, code: str, db=None) -> bool:
    """Verify code for a User instance with encrypted totp_secret in DB.

    Rejects replayed codes within the 90s validation window so a sniffed
    OTP cannot be reused.

    Migration: if ``decrypt_secret`` returned a plain-base32 legacy
    secret (see its docstring), re-encrypt and persist via the optional
    ``db`` session. Skipped when no session is passed so the function
    stays usable from contexts without DB access — login flows always
    have one.
    """
    if not user.totp_enabled or not user.totp_secret:
        return False
    try:
        secret = decrypt_secret(user.totp_secret)
    except InvalidToken:
        return False

    cleaned = (code or "").strip().replace(" ", "")
    if _replay_seen(user.id, cleaned):
        logger.warning("TOTP replay rejected for user_id=%s", user.id)
        return False

    if not verify(secret, cleaned):
        return False

    _replay_record(user.id, cleaned)

    # One-shot migration for legacy plain-base32 secrets — re-encrypt
    # so the next read goes through the modern Fernet path.
    if db is not None and user.totp_secret == secret:
        try:
            user.totp_secret = encrypt_secret(secret)
            db.commit()
            logger.info(
                "Re-encrypted legacy plain-base32 TOTP secret for user_id=%s",
                user.id,
            )
        except Exception as e:
            logger.warning(
                "Could not migrate legacy TOTP secret for user_id=%s: %s",
                user.id, e,
            )
            db.rollback()

    return True
