from Perceptrix.engine import perceptrix, robotix, identify_objects_from_text, search_keyword
from CircumSpect import answer_question, find_object_description, locate_object
from flask import Flask, request, jsonify
import numpy as np
import threading
import whisper
import cv2
import os

model = whisper.load_model("base")

def transcribe(audio):
    result = model.transcribe(audio)
    transcription = result['text']
    print(transcription)
    return transcription

app = Flask(__name__)

@app.route('/', methods=['POST', 'GET'])
def home():
    return jsonify({'message': 'WORKING'})

def handle_request(func, *args):
    try:
        result = func(*args)
        return jsonify({'message': result})
    except Exception as e:
        print(e)
        return jsonify({'error': str(e)})

@app.route('/locate_object', methods=['POST', 'GET'])
def _locate_object():
    image_data = request.json['image']
    prompt = request.json['prompt']
    image_data = np.array(image_data, dtype=np.uint8)
    image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
    cv2.imwrite('API.jpg', image)
    return handle_request(locate_object, prompt, "API.jpg")

@app.route('/vqa', methods=['POST', 'GET'])
def _vqa():
    image_data = request.json['image']
    prompt = request.json['prompt']
    image_data = np.array(image_data, dtype=np.uint8)
    image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
    cv2.imwrite('API.jpg', image)
    return handle_request(answer_question, prompt, "API.jpg")

@app.route('/object_description', methods=['POST', 'GET'])
def _object_description():
    image_data = request.json['image']
    image_data = np.array(image_data, dtype=np.uint8)
    image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
    cv2.imwrite('API.jpg', image)
    return handle_request(find_object_description, "API.jpg")

@app.route('/perceptrix', methods=['POST', 'GET'])
def _perceptrix():
    prompt = request.json['prompt']
    return handle_request(perceptrix, prompt)

@app.route('/robotix', methods=['POST', 'GET'])
def _robotix():
    prompt = request.json['prompt']
    return handle_request(robotix, prompt)

@app.route('/search_keyword', methods=['POST', 'GET'])
def _search_keyword():
    prompt = request.json['prompt']
    return handle_request(search_keyword, prompt)

@app.route('/identify_objects_from_text', methods=['POST', 'GET'])
def _identify_objects_from_text():
    prompt = request.json['prompt']
    return handle_request(identify_objects_from_text, prompt)

@app.route('/transcribe', methods=['POST', 'GET'])
def _upload_audio():
    try:
        audio_file = request.files['audio']
        filename = os.path.join("./", audio_file.filename)
        audio_file.save(filename)
        print("RECEIVED")
        return jsonify({'message': transcribe(filename)})
    except Exception as e:
        print(e)
        return jsonify({'message': "Error"})

def run_app():
    app.run(port=7777)

if __name__ == "__main__":
    runner = threading.Thread(target=run_app)
    runner.start()