import atexit
import functools
from queue import Queue
from threading import Event, Thread

from paddleocr import PaddleOCR, draw_ocr
from PIL import Image
import gradio as gr


LANG_CONFIG = {
    "ch": {"num_workers": 1},
    "en": {"num_workers": 1},
    "fr": {"num_workers": 1},
    "german": {"num_workers": 1},
    "korean": {"num_workers": 1},
    "japan": {"num_workers": 1},
}

CONCURRENCY_LIMIT = 2


class PaddleOCRModelManager:
    def __init__(self, num_workers, model_factory):
        self._model_factory = model_factory
        self._queue = Queue()
        self._workers = []
        self._model_initialized_event = Event()

        for _ in range(num_workers):
            worker = Thread(target=self._worker, daemon=True)
            worker.start()
            self._model_initialized_event.wait()
            self._model_initialized_event.clear()
            self._workers.append(worker)

    def infer(self, *args, **kwargs):
        result_queue = Queue(maxsize=1)
        self._queue.put((args, kwargs, result_queue))
        success, payload = result_queue.get()

        if success:
            return payload
        else:
            raise payload

    def close(self):
        for _ in self._workers:
            self._queue.put(None)

    def _worker(self):
        model = self._model_factory()
        self._model_initialized_event.set()

        while True:
            item = self._queue.get()

            if item is None:
                break

            args, kwargs, result_queue = item

            try:
                result = model.ocr(*args, **kwargs)
                result_queue.put((True, result))
            except Exception as e:
                result_queue.put((False, e))


def create_model(lang):
    return PaddleOCR(lang=lang, use_angle_cls=True, use_gpu=False)


model_managers = {}

for lang, config in LANG_CONFIG.items():
    model_managers[lang] = PaddleOCRModelManager(
        config["num_workers"],
        functools.partial(create_model, lang=lang)
    )


def close_model_managers():
    for manager in model_managers.values():
        manager.close()


atexit.register(close_model_managers)


def inference(img, lang):
    ocr = model_managers[lang]

    result = ocr.infer(img, cls=True)

    if not result or not result[0]:
        return Image.open(img)

    result = result[0]

    image = Image.open(img).convert("RGB")

    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]

    im_show = draw_ocr(
        image,
        boxes,
        txts,
        scores,
        font_path="./simfang.ttf"
    )

    return im_show


title = "PaddleOCR"

description = """
Gradio demo for PaddleOCR.

Supported languages:
Chinese, English, French, German, Korean, Japanese.

Upload an image and select the language.
"""

examples = [
    ["en_example.jpg", "en"],
    ["cn_example.jpg", "ch"],
    ["jp_example.jpg", "japan"],
]

css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"


demo = gr.Interface(
    fn=inference,
    inputs=[
        gr.Image(type="filepath", label="Input"),
        gr.Dropdown(
            choices=list(LANG_CONFIG.keys()),
            value="en",
            label="Language"
        ),
    ],
    outputs=gr.Image(type="pil", label="Output"),
    title=title,
    description=description,
    examples=examples,
    cache_examples=False,
    css=css,
    concurrency_limit=CONCURRENCY_LIMIT,
)

demo.launch()