import base64
from sqlalchemy.orm import Session
import pyotp
import os
import traceback
from qrcode import QRCode, constants
from six import BytesIO

from app.api.users.login_verification import get_user, add_google_key
from app.api.users.schemas import GoogleLogin


def get_qrcode(username: str, gtoken: str):
    # gtoken = pyotp.random_base32(64)
    data = pyotp.totp.TOTP(gtoken).provisioning_uri(username, issuer_name="IAM MFA Code")
    qr = QRCode(
        version=1,
        error_correction=constants.ERROR_CORRECT_L,
        box_size=6,
        border=4, )
    try:
        qr.add_data(data)
        qr.make(fit=True)
        img = qr.make_image()
        # 图片保存
        # dirpath = os.path.join(os.getcwd())
        # filepath = dirpath + os.sep + username + '.png'
        # img.save(filepath)  # 保存条形码图片
        # 图片以二进制形式写入
        buf = BytesIO()
        img.save(buf, 'jpeg')
        buf_str = buf.getvalue()
        img_b = b"data:image/png;base64," + base64.b64encode(buf_str)
        return img_b
    except Exception as e:
        traceback.print_exc()
        return False


def google_verify_result(db: Session, data: GoogleLogin):
    """谷歌动态码效验"""
    if data.google_key == "1":
        user_data = get_user(db, data.username)
        secret_key = user_data.google_key
    else:
        secret_key = data.google_key
    t = pyotp.TOTP(secret_key)
    result = t.verify(data.google_code)  # 对输入验证码进行校验，正确返回True
    msg = result if result is True else False
    if msg and data.google_key != "1":
        update_info = add_google_key(db=db, username=data.username, google_key=data.google_key)
    return msg