from flask import Flask, request, jsonify
import os
from hiou import hugging
import json
import io

def get_answer(image: io.BytesIO, question: str):
    filename = image.name
    spid = os.getenv('SPID')
    hug = hugging(spid)
    filelink = hug.upload(image)
    image.seek(0, 2)
    file_size = image.tell()
    image.seek(0)
    data = [
            {
                "meta": {
                    "_type": "gradio.FileData"
                },
                "mime_type": "image/jpeg",
                "orig_name": filename,
                "path": filelink.split("=", 1)[1],
                "size": file_size,
                "url": filelink
            }
        ]
    datas = os.getenv('DATA')
    jsnd = json.loads(datas)
    data.append(jsnd)
    data.append(
        question
    )
    data.append(
        {
            "tab_index": 0
        }
    )
    hug.filnal_setup(data, 2, 15)
    hug.start()
    return hug.output.get("data")[0]

app = Flask(__name__)

@app.route('/getans', methods=['POST'])
def async_task():
    if 'file' not in request.files:
        return jsonify({"status" : False, "msg" : "No file found"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400
    file_content = io.BytesIO(file.read())
    file_content.name = file.filename
    headers = dict(request.headers)
    if not headers.get("KEY") != os.getenv("KEY"):
        return jsonify({"status" : False, "msg" : "Invalid API Key"}), 404
    print(f"QUESTION ASKED : {headers.get('QU', '')}")
    answer = get_answer(file_content, headers.get('QU', ''))
    return jsonify({"status" : True, "ANS" : answer})

@app.route('/')
def home():
    return jsonify({"message": "Welcome to my Flask API!"})