from PDAPI import *

TOTP_CONFIG = {
    "digits": 6,
    "interval": 30,
}


__metadata__ = {
    "methods": ["POST"],
    "blueprint": "UGM",
    "permission": 0.0,
    "attributes": [],
    "limit": "5 per minute",
    "form": [
        {"Name": "pendingKey", "Type": "str", "MaxLen": 64},
        {"Name": "token", "Type": "str", "MaxLen": 6}
    ]
}

def HandleRequest(_CleanedData):
    pendingKey = _CleanedData["form"].get("pendingKey", "").strip()
    Token = _CleanedData["form"].get("token", "").strip()

    if not pendingKey or not Token or len(Token) != 6:
        return {"code": 1, "msg": "参数错误"}, 200

    pending = GERPUserPending2FA.query.filter(
        GERPUserPending2FA.SecretKey == pendingKey,
        GERPUserPending2FA.ExpireAt > datetime.datetime.utcnow()
    ).first()
    if not pending:
        return {"code": 1, "msg": "验证已过期，请重新登录"}, 200

    user = GERPUser.query.filter(GERPUser.UserUUID == pending.UserUUID).first()
    if not user:
        return {"code": 1, "msg": "用户不存在"}, 200

    two_fa = GERPUser2FA.query.filter(GERPUser2FA.UserUUID == user.UserUUID, GERPUser2FA.IsEnabled == 1).first()
    if not two_fa:
        return {"code": 1, "msg": "2FA未启用"}, 200

    Totp = pyotp.TOTP(two_fa.SecretKey, digits=TOTP_CONFIG["digits"], interval=TOTP_CONFIG["interval"])
    if Totp.verify(Token, valid_window=1):
        two_fa.LastUsedAt = datetime.datetime.utcnow()

        db.session.delete(pending)

        ELTSession = UGM.CreateELT(user)
        if not ELTSession:
            db.session.commit()
            return {"code": 1, "msg": "凭证创建失败"}, 200

        user.UserLastLoginDate = datetime.datetime.utcnow()
        try:
            db.session.commit()
        except:
            db.session.rollback()
            return {"code": 1, "msg": "系统错误"}, 200

        Cookie.set("ELToken", UUIDObj(ELTSession.ELTUUID).UpperStr, domain=".gvsds.com")

        current_role_uuid = flask.request.cookies.get("CurrentRoleUUID")
        if not current_role_uuid:
            first_role = GERPRole.query.join(UserRole).filter(UserRole.UserUUID == user.UserUUID).first()
            if first_role:
                Cookie.set("CurrentRoleUUID", UUIDObj(first_role.RoleUUID).UpperStr, domain=".gvsds.com")

        return {"code": 0, "msg": "2FA验证成功", "redirect": flask.url_for("GERP.GERPIndex")}, 200

    return {"code": 1, "msg": "验证码错误"}, 200