import numpy as np
import cv2

from starlette.requests import Request
from ray import serve
from ray.serve.handle import DeploymentHandle

from texteller.api import load_model, load_tokenizer, img2latex
from texteller.utils import get_device
from texteller.globals import Globals
from typing import Literal


@serve.deployment(
    num_replicas=Globals().num_replicas,
    ray_actor_options={
        "num_cpus": Globals().ncpu_per_replica,
        "num_gpus": Globals().ngpu_per_replica * 1.0 / 2,
    },
)
class TexTellerServer:
    def __init__(
        self,
        checkpoint_dir: str,
        tokenizer_dir: str,
        use_onnx: bool = False,
        out_format: Literal["latex", "katex"] = "katex",
        keep_style: bool = False,
        num_beams: int = 1,
    ) -> None:
        self.model = load_model(
            model_dir=checkpoint_dir,
            use_onnx=use_onnx,
        )
        self.tokenizer = load_tokenizer(tokenizer_dir=tokenizer_dir)
        self.num_beams = num_beams
        self.out_format = out_format
        self.keep_style = keep_style

        if not use_onnx:
            self.model = self.model.to(get_device())

    def predict(self, image_nparray: np.ndarray) -> str:
        return img2latex(
            model=self.model,
            tokenizer=self.tokenizer,
            images=[image_nparray],
            device=get_device(),
            out_format=self.out_format,
            keep_style=self.keep_style,
            num_beams=self.num_beams,
        )[0]


@serve.deployment()
class Ingress:
    def __init__(self, rec_server: DeploymentHandle) -> None:
        self.texteller_server = rec_server

    async def __call__(self, request: Request) -> str:
        form = await request.form()
        img_rb = await form["img"].read()

        img_nparray = np.frombuffer(img_rb, np.uint8)
        img_nparray = cv2.imdecode(img_nparray, cv2.IMREAD_COLOR)
        img_nparray = cv2.cvtColor(img_nparray, cv2.COLOR_BGR2RGB)

        pred = await self.texteller_server.predict.remote(img_nparray)
        return pred