from fastapi import Request, Response, status
from fastapi.responses import JSONResponse
from .blackbox import Blackbox
from typing import Optional

import requests
import base64


def is_base64(value) -> bool:
    try:
        base64.b64decode(base64.b64decode(value)) == value.encode()
        return True
    except Exception:
        return False


class VLMS(Blackbox):

    def __call__(self, *args, **kwargs):
        return self.processing(*args, **kwargs)

    def valid(self, *args, **kwargs) -> bool:
        data = args[0]
        return isinstance(data, list)

    def processing(self, prompt, images, model_name: Optional[str] = None) -> str:

        if model_name == "Qwen-VL-Chat":
            model_name = "infer-qwen-vl"
        elif model_name == "llava-llama-3-8b-v1_1-transformers":
            model_name = "infer-lav-lam-v1-1"
        else:
            model_name = "infer-qwen-vl"

        url = 'http://120.196.116.194:48894/' + model_name + '/'

        if is_base64(images):
            images_data = images
        else:
            with open(images, "rb") as img_file:
                images_data = str(base64.b64encode(img_file.read()), 'utf-8')

        data_input = {'model': model_name, 'prompt': prompt, 'img_data': images_data}

        data = requests.post(url, json=data_input)

        return data.text

    async def fast_api_handler(self, request: Request) -> Response:
        try:
            data = await request.json()
        except:
            return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)

        model_name = data.get("model_name")
        prompt = data.get("prompt")
        img_data = data.get("img_data")

        if prompt is None:
            return JSONResponse(content={'error': "Question is required"}, status_code=status.HTTP_400_BAD_REQUEST)

        if model_name is None or model_name.isspace():
            model_name = "Qwen-VL-Chat"

        jsonresp = str(JSONResponse(content={"response": self.processing(prompt, img_data, model_name)}).body, "utf-8")
        return JSONResponse(content={"response": jsonresp}, status_code=status.HTTP_200_OK)