from functools import wraps
from flask import session, jsonify, redirect, url_for, request
from models import Admin
import pyotp

def admin_required(f):
    """管理员权限装饰器"""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'admin_id' not in session:
            # 如果是API请求，返回JSON
            if request.path.startswith('/admin/api/'):
                return jsonify({'success': False, 'message': '请先登录'}), 401
            # 如果是页面请求，重定向到登录页
            return redirect(url_for('admin.login'))
        return f(*args, **kwargs)
    return decorated_function

def verify_2fa(admin: Admin, code: str) -> bool:
    """验证2FA代码"""
    if not admin.is_2fa_enabled or not admin.totp_secret:
        return True

    totp = pyotp.TOTP(admin.totp_secret)
    return totp.verify(code, valid_window=1)

def generate_2fa_secret() -> str:
    """生成2FA密钥"""
    return pyotp.random_base32()

def get_2fa_qrcode_url(admin: Admin, secret: str) -> str:
    """获取2FA二维码URL"""
    totp = pyotp.TOTP(secret)
    return totp.provisioning_uri(
        name=admin.username,
        issuer_name='短视频解析平台'
    )