"""Dual-token authentication: short-lived access token (Authorization header,
validated per-request) + long-lived refresh token (httpOnly cookie, rotated
on each refresh). TOTP / Passkey / recovery codes / temp password still drive
the login ceremony; the result is a token pair instead of a Flask session.

Multi-user support:
- Every request resolves its user from the access token (`g.user_id`).
- Access tokens snapshot the user's `session_version`; bumping it (delete user,
  kick devices) invalidates every outstanding access token immediately.
- Refresh tokens are stored hashed in `auth_tokens`; deleting them (logout of
  one device / kick) ends the corresponding device session.
- Temp-account / temp-password logins get `temp=1` tokens; the user is forced
  through MFA setup until they complete it, then a fresh non-temp pair is
  issued.
"""
import functools
import os
import secrets
import string
import time
from flask import g, redirect, url_for, request, jsonify, current_app
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,
    user_exists, is_admin_user, get_db,
    create_auth_token, find_access_token,
    delete_auth_token, REFRESH_TOKEN_TTL, ACCESS_TOKEN_TTL,
)

# ─── Token / cookie constants ────────────────
REFRESH_COOKIE = "refresh_token"   # httpOnly, SameSite=Lax, Secure under HTTPS


def _https_enabled() -> bool:
    return os.environ.get("FLASK_HTTPS", "").lower() in ("1", "true", "yes")


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


# ─── Auth context (per-request, populated by before_request) ───

def load_auth_context():
    """Parse `Authorization: Bearer <access>` and fill `g` with the auth context.
    Runs for every request (cheap: no token → no-op)."""
    token = _extract_bearer()
    if not token:
        return
    row = find_access_token(token)
    if not row:
        return
    g.user_id = row["user_id"]
    g.auth_view = row["view"]
    g.temp_pending = bool(row["temp"])
    g.auth_authenticated = True
    g.access_token_id = row["id"]
    g.reauth_at = row.get("reauth_at") or 0


def _extract_bearer() -> str | None:
    auth = request.headers.get("Authorization", "")
    if auth.startswith("Bearer "):
        return auth[7:].strip()
    return None


def _g():
    return g


def is_authenticated() -> bool:
    """True when the request carries a valid, non-temp access token."""
    try:
        return (getattr(_g(), "auth_authenticated", False)
                and not getattr(_g(), "temp_pending", False))
    except Exception:
        return False


def is_temp_pending() -> bool:
    """True when logged in with a temp password / temp account (MFA setup required)."""
    try:
        return (getattr(_g(), "auth_authenticated", False)
                and bool(getattr(_g(), "temp_pending", False)))
    except Exception:
        return False


def current_user_id() -> int | None:
    """Return the id of the current user (token user, or scheduler thread user)."""
    try:
        from flask import has_request_context
        if has_request_context():
            uid = getattr(_g(), "user_id", None)
            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)


# ─── Token issuance / refresh / revocation ───

def issue_tokens(user_id: int, view: str = "control", temp: bool = False) -> dict:
    """Create a fresh access+refresh token pair. The refresh token must be
    delivered to the client via an httpOnly cookie by the caller."""
    access = create_auth_token(user_id, "access", view=view, temp=temp)
    refresh = create_auth_token(user_id, "refresh", view=view, temp=temp)
    return {"access_token": access, "refresh_token": refresh}


def set_refresh_cookie(resp, refresh: str):
    """Attach the refresh token as an httpOnly cookie on the response."""
    resp.set_cookie(
        REFRESH_COOKIE, refresh,
        max_age=REFRESH_TOKEN_TTL,
        httponly=True,
        samesite="Lax",
        secure=_https_enabled(),
        path="/",
    )


def clear_refresh_cookie(resp):
    resp.delete_cookie(REFRESH_COOKIE, path="/")


def rotate_refresh(refresh_raw: str) -> dict | None:
    """Validate a refresh token, revoke it and issue a fresh pair (rotation).
    Atomic (single transaction): a replayed old token can never mint twice.
    Returns the new {access_token, refresh_token} or None when invalid."""
    from backend.models import rotate_refresh_token
    return rotate_refresh_token(refresh_raw)


def logout_current_device():
    """Revoke this device's refresh token (from the httpOnly cookie) and also
    the presented access token, if any (so a stolen access dies immediately)."""
    raw = request.cookies.get(REFRESH_COOKIE)
    if raw:
        delete_auth_token(raw, kind="refresh")
    acc = _extract_bearer()
    if acc:
        delete_auth_token(acc, kind="access")


def load_page_context() -> bool:
    """Authenticate a top-level page navigation via the httpOnly refresh cookie
    (browsers cannot attach Authorization headers to navigations).
    Fills g.user_id / g.auth_view / g.temp_pending on success.
    Returns False when unauthenticated."""
    raw = request.cookies.get(REFRESH_COOKIE)
    if not raw:
        return False
    from backend.models import find_refresh_token
    row = find_refresh_token(raw)
    if not row:
        return False
    g.user_id = row["user_id"]
    g.auth_view = row["view"]
    g.temp_pending = bool(row["temp"])
    g.page_authenticated = True
    return True


# ─── Re-authentication guard ─────────────────
# Sensitive operations (user management, MFA changes) require re-confirming
# identity within the last N seconds. The marker lives on the ACCESS TOKEN row,
# so it is per-device (each device's token must re-auth independently) and
# lost on token rotation (a fresh access token requires a fresh re-auth).

_REAUTH_TTL = 300  # 5 minutes


def require_reauth() -> bool:
    """Check if recent re-authentication is still valid for THIS access token.
    Skips the check for temp_pending users (first-time MFA setup flow).
    """
    if is_temp_pending():
        mark_reauth()
        return True
    try:
        last = getattr(g, "reauth_at", 0) or 0
    except Exception:
        return False
    return bool(last) and (time.time() - last) < _REAUTH_TTL


def mark_reauth():
    """Mark re-authentication as done now for the current access token."""
    try:
        from flask import g as _g
        from backend.models import mark_token_reauth
        mark_token_reauth(getattr(_g, "access_token_id", None))
    except Exception as e:
        try:
            current_app.logger.warning(f"mark_reauth failed: {e}")
        except Exception:
            pass


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

def _unauthorized(request):
    if request.is_json or request.path.startswith("/api/"):
        return jsonify({"error": "Unauthorized", "login_required": True}), 401
    view_param = getattr(_g(), "auth_view", None) or "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()