"""Two-factor sign-in with an authenticator app (app/core/totp.py).

Turning it on is two steps, so a typo never locks anyone out: start_setup()
keeps a new secret (encrypted) and shows it as a QR code; confirm_setup()
switches it on only once a code from the app checks out, and hands back ten
one-time recovery codes — shown once, stored hashed — for when the phone is
lost. An administrator can also reset someone's two-factor (they then set
it up again at their next sign-in, if the tenant requires it).

At sign-in the password step answers with a short-lived challenge token
instead of a session; the code (or a recovery code) and that token make
the session. The token's purpose says which step it's for, and it carries
no tenant_id/user_id claims, so it's useless as a session token.
"""

import hashlib
import hmac
import secrets
import uuid
from datetime import datetime, timedelta, timezone

from jose import JWTError, jwt

from app.core import totp
from app.core.config import settings
from app.core.crypto import decrypt_secret, encrypt_secret
from app.models.identity import Tenant, User

RECOVERY_CODES = 10
CHALLENGE_TTL = timedelta(minutes=10)
# "mfa": enter a code; "mfa_setup": the tenant requires two-factor and this
# account hasn't set it up yet — do that now.
# "password_change": the password was right but has expired, or an
# administrator set it — a new one is chosen before sign-in carries on.
PURPOSES = ("mfa", "mfa_setup", "password_change")


class MfaError(Exception):
    def __init__(self, message: str, status_code: int = 400):
        super().__init__(message)
        self.status_code = status_code


def enabled(user: User) -> bool:
    return user.mfa_enabled_at is not None


# ---- setting up --------------------------------------------------------------------------------------------


def start_setup(user: User, tenant: Tenant) -> dict:
    """A new secret for the authenticator app — not in force until confirmed."""
    if enabled(user):
        raise MfaError("Two-factor sign-in is already on — turn it off first to move it to another app.", 409)
    secret = totp.new_secret()
    user.mfa_secret_encrypted = encrypt_secret(secret)
    uri = totp.provisioning_uri(secret, user.email, f"FlowFabric ({tenant.name})")
    return {"secret": secret, "otpauth_uri": uri, "qr_svg": totp.qr_svg(uri)}


def confirm_setup(user: User, code: str) -> list[str]:
    """Switches two-factor on once a code from the app checks out. Returns
    the recovery codes — the only time they're shown."""
    if enabled(user):
        raise MfaError("Two-factor sign-in is already on.", 409)
    if not user.mfa_secret_encrypted:
        raise MfaError("Start the setup first.")
    step = totp.verify(decrypt_secret(user.mfa_secret_encrypted), code)
    if step is None:
        raise MfaError("That code doesn't match — check the app's clock and try the newest code.")
    user.mfa_last_step = step
    user.mfa_enabled_at = datetime.now(timezone.utc)
    return new_recovery_codes(user)


def disable(user: User) -> None:
    user.mfa_secret_encrypted = None
    user.mfa_enabled_at = None
    user.mfa_recovery_hashes = []
    user.mfa_last_step = None


def _hash(code: str) -> str:
    # Recovery codes are 50 random bits each, so a fast hash is enough.
    return hashlib.sha256(code.replace("-", "").lower().encode()).hexdigest()


def new_recovery_codes(user: User) -> list[str]:
    codes = []
    for _ in range(RECOVERY_CODES):
        raw = "".join(secrets.choice("abcdefghjkmnpqrstuvwxyz23456789") for _ in range(10))
        codes.append(f"{raw[:5]}-{raw[5:]}")
    user.mfa_recovery_hashes = [_hash(c) for c in codes]
    return codes


# ---- checking a code ---------------------------------------------------------------------------------------


def check(user: User, code: str) -> str:
    """Accepts a current code from the app, or an unused recovery code (used
    up by this). Returns which ("totp" | "recovery"); raises MfaError."""
    if not enabled(user) or not user.mfa_secret_encrypted:
        raise MfaError("Two-factor sign-in isn't set up for this account.")
    code = code.strip()
    if sum(ch.isdigit() for ch in code) == totp.DIGITS and len(code.replace(" ", "")) == totp.DIGITS:
        step = totp.verify(decrypt_secret(user.mfa_secret_encrypted), code, last_step=user.mfa_last_step)
        if step is None:
            raise MfaError("That code isn't right, or it was already used — wait for the next one.")
        user.mfa_last_step = step
        return "totp"
    hashed = _hash(code)
    remaining = list(user.mfa_recovery_hashes or [])
    match = next((h for h in remaining if hmac.compare_digest(h, hashed)), None)
    if match is None:
        raise MfaError("That code isn't right.")
    remaining.remove(match)
    user.mfa_recovery_hashes = remaining
    return "recovery"


# ---- the sign-in challenge ---------------------------------------------------------------------------------


def challenge_token(user: User, purpose: str) -> str:
    assert purpose in PURPOSES
    claims = {"purpose": purpose, "challenge_user_id": str(user.id), "challenge_tenant_id": str(user.tenant_id), "exp": datetime.now(timezone.utc) + CHALLENGE_TTL}
    return jwt.encode(claims, settings.jwt_secret, algorithm=settings.jwt_algorithm)


def read_challenge(token: str, tenant_id: uuid.UUID, purpose: str) -> uuid.UUID:
    try:
        claims = jwt.decode(token, settings.jwt_secret, algorithms=[settings.jwt_algorithm])
    except JWTError as exc:
        raise MfaError("This sign-in attempt has expired — sign in again.") from exc
    if claims.get("purpose") != purpose or claims.get("challenge_tenant_id") != str(tenant_id):
        raise MfaError("This sign-in attempt has expired — sign in again.")
    return uuid.UUID(claims["challenge_user_id"])