"""Security helpers: session timeout, CSRF protection, login throttling.

This module is intentionally self-contained so it can be wired into
the existing app without breaking the current auth flow.

Features:
- Session lifetime: configurable idle + absolute timeout
- CSRF token: per-session token, validated on all state-changing requests
- Login throttle: rate-limit failed login attempts by client IP+username
- Audit helpers
"""
from __future__ import annotations

import hashlib
import hmac
import os
import secrets
import time
from collections import defaultdict
from functools import wraps
from threading import Lock

from flask import abort, current_app, flash, redirect, request, session, url_for

# ---- Session lifetime config -------------------------------------------------

DEFAULTS = {
    "session_idle_timeout": 1800,       # 30 min - logout after this much inactivity
    "session_absolute_timeout": 28800,  # 8 hours - hard logout no matter what
    "login_throttle_window": 300,       # 5 min - window for counting failed attempts
    "login_throttle_max": 5,            # 5 failures within window triggers throttle
    "login_throttle_lockout": 900,      # 15 min - lockout duration after threshold
    "csrf_enabled": True,
}

# Module-level state for throttle (per-process; for multi-worker use Redis or DB)
_throttle_lock = Lock()
_failed_logins: dict[tuple[str, str], list[float]] = defaultdict(list)
_locked_until: dict[tuple[str, str], float] = {}


def get_security_config(key: str):
    """Read security config from Flask app config with fallback to DEFAULTS."""
    try:
        return current_app.config.get(f"SECURITY_{key.upper()}", DEFAULTS[key])
    except Exception:
        return DEFAULTS.get(key, "")


# ---- CSRF --------------------------------------------------------------------

def generate_csrf_token() -> str:
    """Return the current session's CSRF token, creating one if missing."""
    token = session.get("_csrf")
    if not token:
        token = secrets.token_urlsafe(32)
        session["_csrf"] = token
    return token


def csrf_token_input() -> str:
    """Jinja helper: render <input type=hidden name=_csrf value=...>."""
    return f'<input type="hidden" name="_csrf" value="{generate_csrf_token()}">'


def csrf_protect():
    """Flask before_request handler that validates CSRF on unsafe methods."""
    if not get_security_config("csrf_enabled"):
        return None
    # skip safe methods
    if request.method in ("GET", "HEAD", "OPTIONS"):
        return None
    # skip login (no session yet)
    if request.endpoint in ("login", "static"):
        return None
    # require logged in
    if not session.get("user_id"):
        return None
    sent = (request.form.get("_csrf")
            or (request.headers.get("X-CSRF-Token", ""))
            or (request.headers.get("X-CSRFToken", "")))
    expected = session.get("_csrf", "")
    if not expected or not sent or not hmac.compare_digest(sent, expected):
        abort(400, "CSRF token missing or invalid")


# ---- Session lifetime --------------------------------------------------------

def enforce_session_lifetime():
    """Flask before_request handler that enforces idle + absolute timeout."""
    if not session.get("user_id"):
        return None
    now = time.time()
    last_seen = session.get("_last_seen", now)
    logged_in_at = session.get("_logged_in_at", now)
    idle_to = get_security_config("session_idle_timeout")
    abs_to = get_security_config("session_absolute_timeout")
    if (now - last_seen) > idle_to:
        # idle timeout
        session.clear()
        flash("会话因长时间无操作已过期,请重新登录", "error")
        return redirect(url_for("login", next=request.path))
    if (now - logged_in_at) > abs_to:
        session.clear()
        flash("会话已达最大时长,请重新登录", "error")
        return redirect(url_for("login", next=request.path))
    session["_last_seen"] = now
    return None


# ---- Login throttle ----------------------------------------------------------

def _throttle_key() -> tuple[str, str]:
    ip = request.headers.get("X-Forwarded-For", request.remote_addr or "0.0.0.0")
    if "," in ip:
        ip = ip.split(",")[0].strip()
    username = (request.form.get("username") or "").strip().lower()
    return (ip, username)


def is_throttled() -> tuple[bool, int]:
    """Return (is_locked, seconds_remaining)."""
    key = _throttle_key()
    with _throttle_lock:
        locked = _locked_until.get(key, 0)
        now = time.time()
        if locked > now:
            return True, int(locked - now)
        # expired lockout - clean up
        if locked and locked <= now:
            _locked_until.pop(key, None)
            _failed_logins.pop(key, None)
        return False, 0


def record_failed_login():
    key = _throttle_key()
    now = time.time()
    with _throttle_lock:
        window = get_security_config("login_throttle_window")
        max_fail = get_security_config("login_throttle_max")
        lockout = get_security_config("login_throttle_lockout")
        # trim old
        _failed_logins[key] = [t for t in _failed_logins[key] if now - t < window]
        _failed_logins[key].append(now)
        if len(_failed_logins[key]) >= max_fail:
            _locked_until[key] = now + lockout
            return True, lockout
        return False, 0


def record_successful_login():
    """Clear any throttle state on success."""
    key = _throttle_key()
    with _throttle_lock:
        _failed_logins.pop(key, None)
        _locked_until.pop(key, None)


# ---- Decorator that requires login + CSRF (already done globally) ------------

def login_required_v2(func):
    """Same as login_required but also ensures CSRF token is generated for forms."""
    @wraps(func)
    def wrapper(*a, **kw):
        if not session.get("user_id"):
            return redirect(url_for("login", next=request.path))
        return func(*a, **kw)
    return wrapper


# ---- Password strength scoring (used by P3-2 user mgmt) ----------------------

def password_strength(pw: str) -> dict:
    """Score password 0-4 (very weak -> very strong)."""
    if not pw:
        return {"score": 0, "label": "空", "hint": "密码不能为空"}
    score = 0
    if len(pw) >= 8: score += 1
    if len(pw) >= 12: score += 1
    if any(c.isupper() for c in pw) and any(c.islower() for c in pw): score += 1
    if any(c.isdigit() for c in pw): score += 1
    if any(not c.isalnum() for c in pw): score += 1
    # special-case: simple/common penalties
    if pw.lower() in ("admin", "password", "123456", "squid", "changeme"):
        score = max(0, score - 2)
    score = min(4, max(0, score))
    labels = {0: "极弱", 1: "弱", 2: "中", 3: "强", 4: "很强"}
    hints = {
        0: "请使用 8+ 字符,含大小写+数字+符号",
        1: "建议加长到 12+ 字符",
        2: "可加入特殊字符进一步提升",
        3: "密码强度良好",
        4: "密码强度优秀",
    }
    return {"score": score, "label": labels[score], "hint": hints[score]}


# ---- Audit log helpers --------------------------------------------------------

def client_ip() -> str:
    """Best-effort client IP, honoring X-Forwarded-For."""
    ip = request.headers.get("X-Forwarded-For", request.remote_addr or "0.0.0.0")
    if "," in ip:
        ip = ip.split(",")[0].strip()
    return ip


def hash_password_for_audit(pw: str) -> str:
    """Stable short hash to record password changes without leaking the password."""
    return hashlib.sha256(pw.encode("utf-8")).hexdigest()[:12]