"""TOTP MFA service layer (B-2.5).

Encrypts the base32 TOTP secret with Fernet (SECRET_KEY-derived key) before
DB persistence. NEVER store the raw base32 secret in the database.

Recovery codes: 10 single-use base32 codes (10 chars each, hyphenated for
readability) generated at TOTP enrollment, displayed ONCE to the user, stored
as bcrypt hashes in User.totp_recovery_codes (JSON list). Each successful
recovery-code login removes that hash from the list.
"""
import base64
import hashlib
import secrets
from typing import List, Optional, Tuple

import pyotp
import qrcode
from cryptography.fernet import Fernet, InvalidToken
from flask import current_app

# 10 single-use recovery codes per enrollment
RECOVERY_CODES_COUNT = 10
RECOVERY_CODE_LENGTH = 10  # base32 chars per code, formatted as XXXXX-XXXXX


def _fernet() -> Fernet:
    """Derive a Fernet key from app SECRET_KEY (deterministic, single-key).

    Uses SHA-256 of SECRET_KEY → urlsafe-base64 (32 bytes) so the same
    SECRET_KEY always produces the same Fernet key. Single-key design (no
    rotation): rotating SECRET_KEY invalidates ALL stored TOTP secrets,
    forcing every user to re-enroll. Acceptable for MVP; revisit when we
    have key-rotation infra.
    """
    secret_key = current_app.config.get('SECRET_KEY')
    if not secret_key:
        raise RuntimeError('SECRET_KEY must be configured to use TOTP encryption')
    if isinstance(secret_key, str):
        secret_key = secret_key.encode('utf-8')
    derived = hashlib.sha256(secret_key).digest()
    return Fernet(base64.urlsafe_b64encode(derived))


def encrypt_totp_secret(plaintext_base32: str) -> str:
    """Encrypt a base32 TOTP secret. Returns the Fernet token as a string."""
    if not plaintext_base32 or not isinstance(plaintext_base32, str):
        raise ValueError('encrypt_totp_secret requires a non-empty base32 string')
    return _fernet().encrypt(plaintext_base32.encode('ascii')).decode('ascii')


def decrypt_totp_secret(ciphertext: str) -> str:
    """Decrypt a Fernet-encrypted TOTP secret. Returns the base32 string.

    Raises ValueError on bad token (key mismatch, tampered, malformed).
    """
    if not ciphertext:
        raise ValueError('decrypt_totp_secret requires a non-empty ciphertext')
    try:
        return _fernet().decrypt(ciphertext.encode('ascii')).decode('ascii')
    except InvalidToken as e:
        raise ValueError(f'Invalid TOTP ciphertext (key mismatch?): {e}') from e


def generate_totp_secret() -> str:
    """Generate a fresh base32 TOTP secret (160-bit, RFC 6238 recommended)."""
    return pyotp.random_base32()


def build_provisioning_uri(secret_base32: str, account_email: str) -> str:
    """Return the otpauth:// URI for QR encoding (RFC 6238)."""
    return pyotp.TOTP(secret_base32).provisioning_uri(
        name=account_email, issuer_name='DictIA'
    )


def render_qr_data_url(provisioning_uri: str) -> str:
    """Render the URI as a base64 PNG data URL for inline display in HTML."""
    import io
    img = qrcode.make(provisioning_uri)
    buf = io.BytesIO()
    img.save(buf, format='PNG')
    return 'data:image/png;base64,' + base64.b64encode(buf.getvalue()).decode('ascii')


def verify_totp_code(secret_base32: str, code: str) -> bool:
    """Verify a 6-digit TOTP code with a 1-window tolerance (current ±30s).

    Rejects non-digit / wrong-length input early to avoid leaking timing.
    """
    if not code or not isinstance(code, str):
        return False
    code = code.strip()
    if len(code) != 6 or not code.isdigit():
        return False
    return pyotp.TOTP(secret_base32).verify(code, valid_window=1)


def _get_bcrypt():
    """Resolve the Flask-Bcrypt extension instance (handles missing context)."""
    bcrypt = (
        current_app.extensions.get('flask-bcrypt')
        or current_app.extensions.get('bcrypt')
    )
    if bcrypt is None:
        # Fall back to the global bcrypt initialised by init_auth_extensions()
        from src.api.auth import bcrypt as _b
        bcrypt = _b
    if bcrypt is None:
        raise RuntimeError('Flask-Bcrypt extension is not initialised')
    return bcrypt


def generate_recovery_codes() -> Tuple[List[str], List[str]]:
    """Generate 10 fresh recovery codes.

    Returns (display_codes, hashed_codes_for_storage):
    - display_codes: human-readable XXXXX-XXXXX format, shown to user ONCE
    - hashed_codes_for_storage: bcrypt-style hashes — store these in
      User.totp_recovery_codes (JSON list).
    """
    bcrypt = _get_bcrypt()
    display_codes: List[str] = []
    hashed_codes: List[str] = []
    for _ in range(RECOVERY_CODES_COUNT):
        # 5+5 base32 chars hyphenated for readability (XXXXX-XXXXX)
        raw = base64.b32encode(secrets.token_bytes(7)).decode('ascii')[:RECOVERY_CODE_LENGTH]
        display = f'{raw[:5]}-{raw[5:]}'
        display_codes.append(display)
        # Hash the hyphenated form (what the user will type back) with bcrypt
        hashed = bcrypt.generate_password_hash(display).decode('ascii')
        hashed_codes.append(hashed)
    return display_codes, hashed_codes


def consume_recovery_code(user, candidate: str) -> bool:
    """Check `candidate` against the user's stored recovery code hashes.

    On match: removes that hash from the list and commits. Returns True.
    On mismatch: returns False, no DB write. Single-use: a code that
    matched once will not match again.
    """
    from src.database import db

    bcrypt = _get_bcrypt()
    if not candidate or not user.totp_recovery_codes:
        return False
    candidate = candidate.strip().upper()
    remaining: List[str] = []
    matched = False
    for h in user.totp_recovery_codes:
        if not matched and bcrypt.check_password_hash(h, candidate):
            matched = True
            # Drop this hash from the stored list (single-use)
            continue
        remaining.append(h)
    if matched:
        user.totp_recovery_codes = remaining
        db.session.commit()
    return matched


def set_user_totp(user, secret_base32: str, recovery_code_hashes: List[str]) -> None:
    """Persist the encrypted secret + recovery codes; mark MFA enabled."""
    from src.database import db
    user.totp_secret_encrypted = encrypt_totp_secret(secret_base32)
    user.totp_recovery_codes = recovery_code_hashes
    user.totp_enabled = True
    db.session.commit()


def disable_user_totp(user) -> None:
    """Disable MFA: clear encrypted secret, recovery codes, and the flag."""
    from src.database import db
    user.totp_secret_encrypted = None
    user.totp_recovery_codes = None
    user.totp_enabled = False
    db.session.commit()


def get_user_totp_secret(user) -> Optional[str]:
    """Return the decrypted base32 secret, or None if MFA not enrolled."""
    if not user.totp_secret_encrypted:
        return None
    return decrypt_totp_secret(user.totp_secret_encrypted)
