from typing import Union
from fastapi import FastAPI, Header, BackgroundTasks,Request,Body
from urllib.parse import unquote
import re
import pickle
import os
from fastapi.responses import JSONResponse
import asyncio
from pydantic import BaseModel
import time
from transformers import AutoTokenizer, AutoModel
import uvicorn, json, datetime
import torch
import json
import base64
import requests
# from tentcentSMS import sendSms

#使用int8的量化模型
# tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
# #model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).half().quantize(8).cuda()
# model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4",trust_remote_code=True).cuda()


# ChatGLM3
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True, device='cuda')

model.eval()

#显存满配全开
#tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
#model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
#model.eval()

def submit_post(url: str, data: dict):
    """
    Submit a POST request to the given URL with the given data.
    """
    return requests.post(url, data=json.dumps(data))


def save_encoded_image(b64_image: str, output_path: str):
    """
    Save the given image to the given output path.
    """
    with open(output_path, "wb") as image_file:
        image_file.write(base64.b64decode(b64_image))





DEVICE = "cuda"
DEVICE_ID = "0"
CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE


def torch_gc():
    if torch.cuda.is_available():
        with torch.cuda.device(CUDA_DEVICE):
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()



app = FastAPI()


class Item(BaseModel):
    query: Union[str, None] = None


@app.get("/")
def index():
    return "hello"





# AI生成图片
@app.post("/img")
async def img(
    query=Body(None),
    appid: Union[str, None] = Header(default=None),
):
    # print("console :,", "appid", appid, "bid", bid, "requestid", requestid, "uid", uid)
    print('query',query)
    print('query[query]',query["query"])
    # 获取请求参数
    global model, tokenizer
    # time.sleep(10)
    # print("item", item)
    history = []  #默认空数组
    prompt=''
    max_length=None
    top_p=None
    temperature=None
    response=None
    try:
        query = unquote(query["query"], "utf-8")
        # 判断有没有'/img'标识符
        if ('/img' in query) is True:
            query=query.replace('/img', '')
            prompt=query
            txt2img_url = 'http://127.0.0.1:7860/sdapi/v1/txt2img'
            data = {'prompt': prompt}
            res = submit_post(txt2img_url, data)
            # 以时间戳命名
            now =str(time.time())
            save_encoded_image(res.json()['images'][0], f'img/{now}.png')
            # 网络图片
            response=f'#web_img http://aiimg.hackrobot.cn/{now}.png'
            print('response',response)
            # return responseimg
        else:
            prompt = query #问题
            # 请求ai模型
            response, history = model.chat(tokenizer,
                                    prompt,
                                    history=history,
                                    max_length=max_length if max_length else 2048,
                                    top_p=top_p if top_p else 0.7,
                                    temperature=temperature if temperature else 0.95)

    except(ValueError, ArithmeticError):
        print('ValueError',ValueError)
        print('ArithmeticError',ArithmeticError)
        prompt = ""

    content = {
        "answer": response,
        "answer_type": "text",
    }
    headers = {"Content-Type": "application/json"}
    return JSONResponse(content=content, headers=headers)