import fastapi
import uvicorn
import gradio
import os
import io
import numpy
import base64
import traceback
from PIL import Image


def extract(data):
    def check(data):
        try:
            image = base64.b64decode(data)
            image = Image.open(io.BytesIO(image))
            if image.format == "JPEG":
                image = numpy.array(image)
                return image
        except Exception as e:
            return False

    images = []

    def check_json(data):
        if isinstance(data, dict):
            for value in data.values():
                check_json(value)
        elif isinstance(data, list):
            for item in data:
                check_json(item)
        elif isinstance(data, str):
            image = check(data)
            if image is not False:
                images.append(image)
        else:
            pass

    check_json(data)
    return images


class Server:
    def __init__(self):
        self.images = [Image.fromarray(numpy.random.randint(0, 256, (360, 640, 3), dtype=numpy.uint8), 'RGB') for _ in range(8)]

    def route(self, app):
        async def error_handler(request: fastapi.Request, exc):
            return fastapi.responses.PlainTextResponse(status_code=500, content=str(exc))

        app.add_exception_handler(Exception, handler=error_handler)
            
        @app.post("/update")
        async def update(data: dict):
            for image in extract(data=data):
                self.images.insert(0, image)
                self.images = self.images[:8]
            return fastapi.responses.JSONResponse(status_code=200, content={})

    def webui(self, app, path):
        with gradio.Blocks() as webui:
            with gradio.Row():
                image_0 = gradio.Image(type="pil", streaming=True)
                image_1 = gradio.Image(type="pil", streaming=True)
                image_2 = gradio.Image(type="pil", streaming=True)
                image_3 = gradio.Image(type="pil", streaming=True)
            with gradio.Row():
                image_4 = gradio.Image(type="pil", streaming=True)
                image_5 = gradio.Image(type="pil", streaming=True)
                image_6 = gradio.Image(type="pil", streaming=True)
                image_7 = gradio.Image(type="pil", streaming=True)
            update = gradio.Button(value="update")

            @update.click(inputs=[], outputs=[image_0, image_1, image_2, image_3, image_4, image_5, image_6, image_7])
            def update():
                while True:
                    yield (gradio.update(value=self.images[0]), 
                            gradio.update(value=self.images[1]), 
                            gradio.update(value=self.images[2]), 
                            gradio.update(value=self.images[3]), 
                            gradio.update(value=self.images[4]), 
                            gradio.update(value=self.images[5]), 
                            gradio.update(value=self.images[6]), 
                            gradio.update(value=self.images[7]))

        gradio.mount_gradio_app(app, webui, path=path)
        
    def run(self, SERVER_HOST, SERVER_PORT):
        app = fastapi.FastAPI()
        self.route(app=app)
        self.webui(app=app, path="/")
        uvicorn.run(app, host=SERVER_HOST, port=int(SERVER_PORT))


if __name__ == "__main__":
    server = Server()
    server.run(
        SERVER_HOST=os.getenv("SERVER_HOST", "0.0.0.0"), 
        SERVER_PORT=os.getenv("SERVER_PORT", 60000),
    )