import base64
import io
import sys
import torch
import onnx
import onnxruntime as rt
from torchvision import transforms as T
from tokenizer_base import Tokenizer
from PIL import Image
from huggingface_hub import hf_hub_download, try_to_load_from_cache


class DocumentParserModel:
    def __init__(self):
        charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
        img_size = (32, 128)

        self.tokenizer_base = Tokenizer(charset)
        self.transform = self.create_transform_pipeline(img_size)
        self.ort_session = self.initialize_onnx_model()

    def create_transform_pipeline(self, img_size):
        transforms = [
            T.Resize(img_size, T.InterpolationMode.BICUBIC),
            T.ToTensor(),
            T.Normalize(0.5, 0.5),
        ]
        return T.Compose(transforms)

    def initialize_onnx_model(self):
        repo_id = "stevenchang/captcha"
        filename = "captcha.onnx"

        filepath = try_to_load_from_cache(repo_id, filename)

        if isinstance(filepath, str):
            model_file = filepath
        else:
            model_file = result = hf_hub_download(repo_id, filename)

        onnx_model = onnx.load(model_file)
        onnx.checker.check_model(onnx_model)
        return rt.InferenceSession(model_file)

    def load_image_from_base64(self, base64_string):
        img_data = base64.b64decode(base64_string)
        image_buffer = io.BytesIO(img_data)

        try:
            image = Image.open(image_buffer)
            return image
        except IOError:
            print(f"Error: Cannot open image {image_blob}")
            return None

    def predict_text(self, image_blob):
        with self.load_image_from_base64(image_blob) as img_org:
            x = self.transform(img_org.convert("RGB")).unsqueeze(0)
            ort_inputs = {self.ort_session.get_inputs()[0].name: x.cpu().numpy()}
            logits = self.ort_session.run(None, ort_inputs)[0]
            probs = torch.tensor(logits).softmax(-1)
            preds, _ = self.tokenizer_base.decode(probs)
            return preds[0]


if __name__ == "__main__":
    import sys

    doc_parser = DocumentParserModel()

    if len(sys.argv) > 1:
        image_blob = sys.argv[1]
        result = doc_parser.predict_text(image_blob)
        print(result)
    else:
        print("Please provide an image blob.")