from flask import Flask, request
from gradio_client import Client
from uuid import uuid4
import os, json, requests, hashlib

app = Flask(__name__)

secret = os.environ.get("auth")
headers = {'Authorization': os.environ.get("auth")}

res_dict = {}

@app.route('/queue', methods=['POST'])
def queue():
    if request.json.get('secret') != secret:
        return {}, 400
        
    uuid = str(uuid4())
    res_dict[uuid] = request.json
    
    return { 'uuid': uuid }, 200


@app.route('/6asr', methods=['POST'])
def sixasr():
    data = res_dict.get(request.json.get('uuid'))
    if not data:
        return {}, 400

    base64 = request.json.get('base64')
    dialectCode = request.json.get('dialectCode')
    if not base64 or not dialectCode:
        return {}, 400
    
    if hashlib.sha256(base64.encode('utf-8')).hexdigest() != data.get("hash"):
        return {}, 400

    URL = 'https://nuu-lab-gradio-6asr-whisper.hf.space/run/predict'
    body = {'data': ["whisper-large-v3", f"htia_{dialectCode}", None, base64, '']}

    res = requests.post(URL, json=body, headers=headers)
    if res.status_code == 200:
        res_json = res.json()

        callbackURL = data.get('callbackURL')
        callback_body = data.get('callbackBody')

        if res_json.get('data'):
            callback_body['length'] = len(res_json['data'][0])
        if res_json.get('srt'):
            callback_body['srt'] = res_json['srt'][0]

        requests.post(callbackURL, json=callback_body)

        print(callback_body)
        print(res_json)

        return res_json

    return {'error': '請求失敗'}, 500


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)