import flask
import base64
import tempfile
import traceback
from flask import Flask, Response, stream_with_context
from inference_vision import OmniVisionInference


class OmniChatServer(object):
    def __init__(self, ip='0.0.0.0', port=60808, run_app=True,
                 ckpt_dir='./checkpoint', device='cuda:0') -> None:
        server = Flask(__name__)
        # CORS(server, resources=r"/*")
        # server.config["JSON_AS_ASCII"] = False

        self.client = OmniVisionInference(ckpt_dir, device)
        self.client.warm_up()

        server.route("/chat", methods=["POST"])(self.chat)

        if run_app:
            server.run(host=ip, port=port, threaded=False)
        else:
            self.server = server

    def chat(self) -> Response:

        req_data = flask.request.get_json()
        try:
            audio_data_buf = req_data["audio"].encode("utf-8")
            audio_data_buf = base64.b64decode(audio_data_buf)
            stream_stride = req_data.get("stream_stride", 4)
            max_tokens = req_data.get("max_tokens", 2048)

            image_data_buf = req_data.get("image", None)
            if image_data_buf:
                image_data_buf = image_data_buf.encode("utf-8")
                image_data_buf = base64.b64decode(image_data_buf)

            audio_path, img_path = None, None
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as audio_f, \
                 tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as img_f:
                audio_f.write(audio_data_buf)
                audio_path = audio_f.name

                if image_data_buf:
                    img_f.write(image_data_buf)
                    img_path = img_f.name
                else:
                    img_path = None

                if img_path is not None:
                    resp_generator = self.client.run_vision_AA_batch_stream(audio_f.name, img_f.name, 
                                                                             stream_stride, max_tokens,
                                                                             save_path='./vision_qa_out_cache.wav')
                else:
                    resp_generator = self.client.run_AT_batch_stream(audio_f.name, stream_stride, 
                                                                      max_tokens,
                                                                      save_path='./audio_qa_out_cache.wav')
                return Response(stream_with_context(self.generator(resp_generator)), 
                                mimetype='multipart/x-mixed-replace; boundary=frame')
        except Exception as e:
            print(traceback.format_exc())
            return Response("An error occurred", status=500)
            
    def generator(self, resp_generator):
        for audio_stream, text_stream in resp_generator:
            yield b'\r\n--frame\r\n'
            yield b'Content-Type: audio/wav\r\n\r\n'
            yield audio_stream
            yield b'\r\n--frame\r\n'
            yield b'Content-Type: text/plain\r\n\r\n'
            yield text_stream.encode()


# CUDA_VISIBLE_DEVICES=1 gunicorn -w 2 -b 0.0.0.0:60808 'server:create_app()'
def create_app():
    server = OmniChatServer(run_app=False)
    return server.server


def serve(ip='0.0.0.0', port=60808, device='cuda:0'):

    OmniChatServer(ip, port=port,run_app=True, device=device)


if __name__ == "__main__":
    import fire
    fire.Fire(serve)