"""
user - 用户相关功能接口

Author: Hao
Date: 2021/12/29
"""
import hashlib
import io

import openpyxl
import pymysql.cursors
from flask import session, request, make_response, Blueprint

import captcha
from utils import random_captcha_code, get_mysql_connection, check_login

user_bp = Blueprint('user', __name__, url_prefix='/user')


@user_bp.route('/captcha')
def get_captcha_image():
    cap = captcha.Captcha.instance()  # type: captcha.Captcha
    captcha_code = random_captcha_code()
    session['captcha_code'] = captcha_code.lower()
    cap_image_data = cap.generate(captcha_code)
    resp = make_response(cap_image_data)
    resp.headers['content-type'] = 'image/png'
    return resp


@user_bp.route('/login', methods=['POST', ])
def login():
    params = request.json
    captcha_from_user = params.get('captcha').lower()
    captcha_from_sess = session.get('captcha_code')
    if captcha_from_sess != captcha_from_user:
        return {'code': 10001, 'message': '验证码错误'}
    username = params.get('username')
    password = params.get('password')
    password = hashlib.md5(password.encode()).hexdigest()
    conn = get_mysql_connection()
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cursor:
            cursor.execute(
                'select user_id, nickname, avatar from tb_user where username=%s and userpass=%s',
                (username, password)
            )
            user_dict = cursor.fetchone()
    finally:
        conn.close()
    if user_dict is None:
        return {'code': 10002, 'message': '用户名或密码错误'}
    session['user_id'] = user_dict['user_id']
    session.permanent = True
    nickname, avatar = user_dict['nickname'], user_dict['avatar']
    return {'code': 10000, 'message': '登录成功', 'nickname': nickname, 'avatar': avatar}


@user_bp.route('/logout')
@check_login
def logout():
    session.clear()
    return {'code': 10003, 'message': '退出登录'}


@user_bp.route('/export')
@check_login
def export_excel():
    wb = openpyxl.Workbook()
    sheet = wb.active
    sheet.append(('交易日', '开盘价', '收盘价', '最低价', '最高价', '交易量'))
    conn = get_mysql_connection(database='stock')
    try:
        with conn.cursor() as cursor:
            cursor.execute(
                'select trade_date, open_price, close_price, '
                'low_price, high_price, trade_volume '
                'from tb_baba_stock order by stock_id'
            )
            for row in cursor.fetchall():
                sheet.append(row)
    finally:
        conn.close()
    buffer = io.BytesIO()
    wb.save(buffer)
    resp = make_response(buffer.getvalue())
    resp.headers['content-type'] = 'application/vnd.ms-excel'
    resp.headers['content-disposition'] = f'attachment; filename="stocks.xlsx"'
    return resp