from flask import *
from gradio_client import Client
from hashlib import sha3_512



class QwenClient:
    def __init__(self):
        self.url = "https://qwen-qwen2-72b-instruct.hf.space/"
        self.client = Client(self.url)
        self.history = []
        self.system_prompt = ""

    def predict(self, prompt: str, system_prompt: str = ""):
        if system_prompt != "":
            self.set_system_prompt(system_prompt)
        _, self.history, system = self.client.predict(
            prompt,
            self.history,
            self.system_prompt,
            api_name="/model_chat"
        )
        return self.history[-1][-1]

    def clear(self):
        _, self.history = self.client.predict(
            api_name="/clear_session"
        )

    def set_system_prompt(self, system_prompt: str):
        _, self.system_prompt, self.history = self.client.predict(
            system_prompt,
            api_name="/modify_system_session"
        )






app = Flask(__name__)
app.config["qwen_client"] = QwenClient()

@app.route("/")
def index():
    return "Hello, World!"

@app.route("/predict", methods=["GET", "POST"])
def predict():
    if request.method == "POST":
        token = request.headers.get("Token")
        if token != None:
            if sha3_512(token.encode()).hexdigest() != "9b3976e38d43c561789976e377a11810a9131221b863fc54c812738230a94f354dba516a0003ef88c8efd5a1ccb12a2b72fb90d5e91aa8cbc6fe3fe2f3af5e4f":
                return "Invalid token"
            else:
                prompt = request.data.decode()
                try:
                    qwen_client = app.config["qwen_client"]
                    result = qwen_client.predict(prompt)
                    return result
                except Exception as e:
                    print(f"Error: {str(e)}")
                    app.config["qwen_client"] = QwenClient()
                    qwen_client = app.config["qwen_client"]
                    result = qwen_client.predict(prompt)
                    return result
        else:
            return "Token not provided"
    else:
        return "API use POST method only"
    


app.run(host="0.0.0.0", port=7860)