import flask , requests , json , pymysql

# 设置请求的目标URL
url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"  # 替换为你的API端点URL

header = {
    "Content-Type":"application/json",
    "Authorization":"Bearer sk-69129a5d7fc6468a9f6f30d6935254c6"

}

def request_TY(userkey,prompt,history = ""):
    # 设置请求数据
    if history == "":
        data = {
            "model": "qwen-turbo",
            "input":{
                "prompt":f"{prompt}"
            }
        }
    else:
        data = {
            "model": "qwen-turbo",
            "input":{
                "prompt":f"{prompt}",
                "history":history
            }
        }
    # 发送POST请求
    response = json.loads(requests.post(url, json=data ,headers=header).text)
    if 'code' in response:
        if response['code'] == 'DataInspectionFailed':
            return response['message']
    reduce_value(userkey, response["usage"]["total_tokens"])
    return response['output']['text']

def userSurplus(userkey):
    #打开数据库连接
    db = pymysql.connect(host='www.kakuweb.top',
                     port=33306,
                     user='root',
                     password='KAku3.14159..',
                     database='AIapi')
    # 使用 cursor() 方法创建一个游标对象 cursor
    cursor = db.cursor()
 
    # 使用 execute()  方法执行 SQL 查询 
    cursor.execute(f"SELECT surplus FROM usersurplus WHERE userkey = '{userkey}';")
    # 使用 fetchone() 方法获取单条数据.
    data = cursor.fetchone()

    # 关闭连接
    db.close()

    if data != None:
        return data[0]
    return -99999

def reduce_value(userkey, value):
    #打开数据库连接
    db = pymysql.connect(host='www.kakuweb.top',
                     port=33306,
                     user='root',
                     password='KAku3.14159..',
                     database='AIapi')
    # 使用 cursor() 方法创建一个游标对象 cursor
    cursor = db.cursor()

    # 执行 SQL 查询以获取当前值
    cursor.execute(f"SELECT surplus FROM usersurplus WHERE userkey = '{userkey}';")
    current_value = cursor.fetchone()[0]

    # 如果没有找到用户，则返回错误信息
    if current_value is None:
        db.close()
        return -1

    # 计算新的值
    new_value = current_value - value

    # 更新数据库中的值
    cursor.execute(f"UPDATE usersurplus SET surplus={new_value} WHERE userkey='{userkey}'")

    # 提交事务
    db.commit()

    # 关闭连接
    db.close()

    # 返回新值
    return 0




app = flask.Flask(__name__)

@app.route('/api/user', methods=['POST'])
def post_data():
    userrequest = flask.request.json
    surplus = userSurplus(userrequest['userkey'])

    if userrequest["prompt"] == "":
        return {"code":400,"output":"Input is empty"}
    
    if userrequest["prompt"] == "":
        return {"code":400,"output":"UserKey is empty"}

    if surplus > 0:
        if userrequest["context"] == 1:
            return {"code":200,"output":f"{request_TY(userrequest['userkey'],userrequest['prompt'],userrequest['history'])}","surplus":userSurplus(userrequest['userkey'])}
        else:
            return {"code":200,"output":f"{request_TY(userrequest['userkey'],userrequest['prompt'])}","surplus":userSurplus(userrequest['userkey'])}
    elif surplus == -99999:
        return {"code":400,"output":"No User"}
    elif surplus <= 0:
        return {"code":403,"output":"No prompt Tokens","surplus":surplus}


if __name__ == '__main__':
    app.run(debug=True,host='0.0.0.0',port=5000)