"""MFA-based authentication: TOTP, Passkey (WebAuthn), recovery codes, temp password.

Multi-user support:
- Sessions carry the authenticated `user_id` (Root = 0).
- The unified login (6-digit code) searches every active user's TOTP secret,
  then temp accounts' temp-TOTP secrets, then Root's temp password.
- Session version is per-user (stored in the users table), so deleting a user
  or kicking all sessions immediately invalidates their cookies.
"""
import functools
import secrets
import string
import time
from flask import session, redirect, url_for, request, jsonify, current_app, g
from backend.models import (
    get_auth_config, update_auth_config, get_temp_password, set_temp_password,
    clear_temp_password, get_recovery_codes, set_recovery_codes, has_mfa_configured,
    get_user_session_version, user_exists, is_admin_user, get_db,
)

# ─── Session keys ────────────────────────────
SKEY_AUTH = "mfa_authenticated"      # bool: fully authenticated
SKEY_TEMP = "mfa_temp_pending"        # bool: logged in with temp pw / temp account, config required
SKEY_VIEW = "mfa_view"                # str: "control" or "display"
SKEY_REAUTH = "mfa_reauth_at"         # float: timestamp of last re-authentication
SKEY_VERSION = "mfa_session_ver"      # int: session version (for kick_all invalidation)
SKEY_USER = "user_id"                 # int: id of the authenticated user


def _db_path(preferred: str = None) -> str:
    if preferred:
        return preferred
    return current_app.config.get("DB_PATH", "data.db")


def _validate_version():
    """If the session version is stale (kick_all / user deleted), clear the session."""
    uid = session.get(SKEY_USER)
    if uid is None:
        return
    # Cache the validation result for the duration of this request.
    if getattr(g, "_sv_ok", False) and getattr(g, "_sv_uid", None) == uid:
        return
    ver = session.get(SKEY_VERSION)
    if ver is None:
        return
    stored = get_user_session_version(uid)
    if stored == -1 or ver != stored:
        clear_session()
        return
    g._sv_ok = True
    g._sv_uid = uid


def is_authenticated() -> bool:
    """Returns True if the user has completed MFA authentication."""
    _validate_version()
    return session.get(SKEY_AUTH) is True


def is_temp_pending() -> bool:
    """True if logged in with a temp password or temp account and must set up a login method."""
    _validate_version()
    return session.get(SKEY_TEMP) is True


def current_user_id() -> int | None:
    """Return the id of the current user (session user, or scheduler thread user)."""
    try:
        from flask import has_request_context
        if has_request_context():
            _validate_version()
            uid = session.get(SKEY_USER)
            if uid is not None:
                return int(uid)
    except Exception:
        pass
    try:
        from backend import time_util
        return time_util.get_thread_user()
    except Exception:
        return None


def current_remark() -> str:
    """Return the current user's remark for display (never the raw id unless it IS the remark)."""
    uid = current_user_id()
    if uid is None:
        return ""
    try:
        from backend.models import get_user
        u = get_user(uid)
        return (u or {}).get("remark", "") or ""
    except Exception:
        return ""


def is_admin() -> bool:
    """Only Root (id 0) — and any future admin role — can manage users."""
    uid = current_user_id()
    if uid is None:
        return False
    return is_admin_user(uid)


def set_authenticated(view: str = "control", user_id: int = None):
    if user_id is None:
        user_id = session.get(SKEY_USER) or 0
    session[SKEY_AUTH] = True
    session[SKEY_TEMP] = False
    session[SKEY_VIEW] = view
    session[SKEY_USER] = user_id
    session[SKEY_VERSION] = get_user_session_version(user_id)
    # A fresh authentication must not inherit the temp-setup re-auth marker.
    session.pop(SKEY_REAUTH, None)
    session.permanent = True


def set_temp_pending(view: str = "control", user_id: int = None):
    """Mark the session as logged-in-but-must-configure (temp password / temp account)."""
    if user_id is None:
        user_id = session.get(SKEY_USER) or 0
    session[SKEY_AUTH] = False
    session[SKEY_TEMP] = True
    session[SKEY_VIEW] = view
    session[SKEY_USER] = user_id
    session[SKEY_VERSION] = get_user_session_version(user_id)
    session.permanent = True


def clear_session():
    session.pop(SKEY_AUTH, None)
    session.pop(SKEY_TEMP, None)
    session.pop(SKEY_VIEW, None)
    session.pop(SKEY_REAUTH, None)
    session.pop(SKEY_USER, None)
    session.pop(SKEY_VERSION, None)


# ─── Re-authentication guard ─────────────────
# Sensitive operations (user management, MFA changes) require re-confirming
# identity within the last N seconds. TOTP-first with a Passkey fallback
# (the fallback ceremony is handled by the routes layer).

_REAUTH_TTL = 300  # 5 minutes


def require_reauth():
    """Check if recent re-authentication is still valid.
    Skips the check for temp_pending users (first-time MFA setup flow).
    """
    if is_temp_pending():
        mark_reauth()
        return True
    last = session.get(SKEY_REAUTH)
    if last and (time.time() - last) < _REAUTH_TTL:
        return True
    return False


def mark_reauth():
    """Mark re-authentication as done now."""
    session[SKEY_REAUTH] = time.time()


# ─── Decorators ──────────────────────────────

def _unauthorized(request):
    if request.is_json or request.path.startswith("/api/"):
        return jsonify({"error": "Unauthorized", "login_required": True}), 401
    view_param = session.get(SKEY_VIEW, "control")
    return redirect(url_for("api.login_page", next=request.path, view=view_param))


def login_required(view):
    """Require full MFA authentication. Redirects to login if not authenticated."""
    @functools.wraps(view)
    def wrapped(*args, **kwargs):
        if is_authenticated():
            return view(*args, **kwargs)
        # If logged in with temp password / temp account, redirect to MFA setup
        if is_temp_pending():
            if request.is_json or request.path.startswith("/api/"):
                return jsonify({"error": "MFA configuration required", "mfa_required": True}), 403
            return redirect(url_for("api.mfa_setup_page"))
        return _unauthorized(request)
    return wrapped


def login_or_temp_required(view):
    """Allow access if fully authenticated OR in temp-pending state (MFA setup)."""
    @functools.wraps(view)
    def wrapped(*args, **kwargs):
        if is_authenticated() or is_temp_pending():
            return view(*args, **kwargs)
        return _unauthorized(request)
    return wrapped


def display_login_required(view):
    """Lighter auth check for display-end API calls."""
    @functools.wraps(view)
    def wrapped(*args, **kwargs):
        if is_authenticated():
            return view(*args, **kwargs)
        return _unauthorized(request)
    return wrapped


def admin_required(view):
    """Require full authentication AND administrator (Root) role."""
    @functools.wraps(view)
    def wrapped(*args, **kwargs):
        if not is_authenticated():
            return _unauthorized(request)
        if not is_admin():
            return jsonify({"error": "Forbidden: 仅管理员可执行此操作"}), 403
        return view(*args, **kwargs)
    return wrapped


# ─── Unified Login ───────────────────────────

def verify_login(raw: str, db_path: str = None) -> dict | None:
    """Unified login: tries every user's TOTP, then temp accounts' temp-TOTP,
    then Root's temp password.

    Returns:
        {'method': 'totp'|'temp_password', 'user_id': int, 'temp_account': bool}
    or None on failure.
    """
    if not raw or not isinstance(raw, str) or not raw.isdigit() or len(raw) != 6:
        return None
    db_path = db_path or current_app.config.get("DB_PATH", "data.db")
    conn = get_db(db_path)
    try:
        import pyotp
        # Load all users + their auth rows in ONE query (avoids per-user connections).
        rows = conn.execute(
            "SELECT u.id AS user_id, u.is_temp, u.temp_totp_secret, "
            "a.totp_secret, a.totp_enabled, a.temp_password "
            "FROM users u LEFT JOIN auth_config a ON a.user_id = u.id "
            "WHERE u.status='active'"
        ).fetchall()

        # 1. Formal users' TOTP secrets
        for r in rows:
            if r["totp_enabled"] and r["totp_secret"]:
                if pyotp.TOTP(r["totp_secret"]).verify(raw, valid_window=1):
                    return {"method": "totp", "user_id": r["user_id"], "temp_account": False}
        # 2. Temp accounts' temp-TOTP (server-only secret; admin sees only the live code)
        for r in rows:
            if r["is_temp"] and r["temp_totp_secret"]:
                if pyotp.TOTP(r["temp_totp_secret"]).verify(raw, valid_window=1):
                    return {"method": "totp", "user_id": r["user_id"], "temp_account": True}
        # 3. Root temp password (initial bootstrap)
        root_row = next((r for r in rows if r["user_id"] == 0), None)
        stored_pw = (root_row["temp_password"] or "") if root_row else ""
        if stored_pw and secrets.compare_digest(raw, stored_pw):
            return {"method": "temp_password", "user_id": 0, "temp_account": False}
        return None
    finally:
        conn.close()


# ─── Temp Password ───────────────────────────

def generate_temp_password(length: int = 6) -> str:
    """Generate a random numeric temp password."""
    return "".join(secrets.choice(string.digits) for _ in range(length))


def verify_temp_password(raw: str, db_path: str = None, user_id: int = 0) -> bool:
    """Verify a temporary password for a given user (only Root uses one)."""
    if not raw or not isinstance(raw, str) or not raw.isdigit() or len(raw) != 6:
        return False
    db_path = db_path or current_app.config.get("DB_PATH", "data.db")
    stored = get_temp_password(db_path, user_id)
    if not stored:
        return False
    return secrets.compare_digest(raw, stored)


# ─── Rate Limiter ─────────────────────────────
# Progressive per-IP ban: 5err→5min, 10err→15min, 15+err→24hr

_RATE_LIMIT_STORE: dict[str, dict] = {}  # ip -> {"failures": int, "banned_until": float}


def _get_client_ip() -> str:
    """Extract client IP for rate limiting.
    Relies on ProxyFix + Nginx config to set remote_addr correctly.
    """
    if not current_app:
        return "unknown"
    return request.remote_addr or "unknown"


def _ban_duration_sec(failures: int) -> int:
    """Progressive ban duration based on consecutive failure count."""
    if failures < 5:
        return 0
    if failures < 10:
        return 5 * 60       # 5 minutes
    if failures < 15:
        return 15 * 60      # 15 minutes
    return 24 * 60 * 60     # 24 hours


def check_rate_limit(ip: str = None) -> dict:
    """Check if this IP is currently rate-limited."""
    if ip is None:
        ip = _get_client_ip()
    entry = _RATE_LIMIT_STORE.get(ip)
    if not entry:
        return {"ok": True}

    now = time.time()
    banned_until = entry.get("banned_until", 0)
    if now < banned_until:
        remaining = int(banned_until - now)
        if remaining < 60:
            msg = f"登录尝试过多，请等待 {remaining} 秒"
        elif remaining < 3600:
            msg = f"登录尝试过多，请等待 {remaining // 60} 分钟"
        else:
            msg = f"登录尝试过多，请等待 {remaining // 3600} 小时"
        return {"ok": False, "message": msg, "retry_after": remaining}

    # Ban expired — reset to current tier (keep failure count so next ban is higher tier)
    entry["banned_until"] = 0
    return {"ok": True}


def record_failure(ip: str = None):
    """Record a failed login attempt. Applies progressive ban if threshold reached."""
    if ip is None:
        ip = _get_client_ip()
    now = time.time()
    entry = _RATE_LIMIT_STORE.get(ip, {"failures": 0, "banned_until": 0})
    entry["failures"] += 1

    duration = _ban_duration_sec(entry["failures"])
    if duration > 0:
        entry["banned_until"] = now + duration
        if current_app:
            current_app.logger.warning(
                f"Rate limit: IP {ip} banned for {duration // 60}min "
                f"after {entry['failures']} failures"
            )

    _RATE_LIMIT_STORE[ip] = entry


def record_success(ip: str = None):
    """Reset failure counter on successful login."""
    if ip is None:
        ip = _get_client_ip()
    _RATE_LIMIT_STORE.pop(ip, None)


# ─── TOTP ────────────────────────────────────

def verify_totp(raw: str, db_path: str = None, user_id: int = 0) -> bool:
    """Verify a TOTP 6-digit code against a specific user's stored secret."""
    if not raw or not isinstance(raw, str) or not raw.isdigit() or len(raw) != 6:
        return False
    db_path = db_path or current_app.config.get("DB_PATH", "data.db")
    cfg = get_auth_config(db_path, user_id)
    secret = cfg.get("totp_secret", "")
    if not secret:
        return False
    import pyotp
    totp = pyotp.TOTP(secret)
    return totp.verify(raw, valid_window=1)


def current_totp_code(user_id: int, db_path: str = None) -> str:
    """Compute the current live TOTP code for a user's secret (admin helper)."""
    db_path = db_path or current_app.config.get("DB_PATH", "data.db")
    cfg = get_auth_config(db_path, user_id)
    secret = cfg.get("totp_secret", "")
    if not secret:
        return ""
    import pyotp
    return pyotp.TOTP(secret).now()


def current_temp_totp_code(user_id: int, db_path: str = None) -> str:
    """Compute the current live temp-TOTP code for a temp account (admin helper)."""
    db_path = db_path or current_app.config.get("DB_PATH", "data.db")
    from backend.models import get_temp_totp_secret
    secret = get_temp_totp_secret(user_id, db_path)
    if not secret:
        return ""
    import pyotp
    return pyotp.TOTP(secret).now()


# ─── Recovery Codes ──────────────────────────

def generate_recovery_codes(count: int = 8) -> list:
    """Generate N random alphanumeric recovery codes, each 10 chars."""
    alphabet = string.ascii_uppercase + string.digits
    codes = []
    for _ in range(count):
        code = "".join(secrets.choice(alphabet) for _ in range(10))
        formatted = f"{code[:4]}-{code[4:]}"
        codes.append(formatted)
    return codes


def verify_recovery_code(raw: str, db_path: str = None) -> int | None:
    """Verify and consume a recovery code. Returns the user id on success, else None."""
    if not raw or not isinstance(raw, str):
        return None
    normalized = raw.replace("-", "").upper()
    if len(normalized) != 10:
        return None

    db_path = db_path or current_app.config.get("DB_PATH", "data.db")
    conn = get_db(db_path)
    try:
        rows = conn.execute(
            "SELECT u.id AS user_id, a.recovery_codes FROM users u "
            "LEFT JOIN auth_config a ON a.user_id = u.id "
            "WHERE u.status='active'"
        ).fetchall()
        for r in rows:
            raw_codes = r["recovery_codes"] or ""
            if not raw_codes:
                continue
            codes = raw_codes.split(",")
            for i, stored in enumerate(codes):
                stored_norm = stored.replace("-", "").upper()
                if secrets.compare_digest(normalized, stored_norm):
                    codes.pop(i)
                    conn.execute(
                        "UPDATE auth_config SET recovery_codes=? WHERE user_id=?",
                        (",".join(codes), r["user_id"])
                    )
                    conn.commit()
                    return r["user_id"]
        return None
    finally:
        conn.close()