import base64
import io
import json
import os
import time

import numpy as np
import torch
import triton_python_backend_utils as pb_utils  # type: ignore
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor


class TritonPythonModel:
    def initialize(self, args):
        os.environ["HF_HUB_OFFLINE"] = "1"
        os.environ["TRANSFORMERS_OFFLINE"] = "1"

        base_path = args['model_repository']
        base_path = os.path.join(base_path, "1")
        model_path = os.path.join(base_path, "icon_handler")

        instance_device_id = str(args.get("model_instance_device_id", "0")).strip()
        default_device = f"cuda:{instance_device_id}" if len(instance_device_id) > 0 else "cuda:0"
        requested_device = os.environ.get(
            "FLORENCE2_DEVICE",
            os.environ.get("ICON_DEVICE", default_device),
        ).strip().lower()
        if requested_device == "auto":
            requested_device = default_device
        if not requested_device.startswith("cuda"):
            raise RuntimeError(
                f"florence2 requires GPU-only execution; invalid device '{requested_device}'. "
                "Use FLORENCE2_DEVICE=cuda:<id>."
            )
        if not torch.cuda.is_available():
            raise RuntimeError("florence2 requires CUDA, but torch.cuda.is_available() is False.")

        self.device = requested_device
        self.dtype = torch.float16
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            trust_remote_code=True,
            local_files_only=True,
            torch_dtype=self.dtype
        ).to(self.device)
        self.processor = AutoProcessor.from_pretrained(
            model_path,
            trust_remote_code=True,
            local_files_only=True
        )
        self.model.eval()
        pb_utils.Logger.log_info(f"florence2 model initialized on {self.device} ({self.dtype})")

        self.model_config = json.loads(args["model_config"])
        output_config = pb_utils.get_output_config_by_name(self.model_config, "CAPTIONS") # type: ignore
        self.output_type = pb_utils.triton_string_to_numpy(output_config["data_type"]) # type: ignore

        self.batch_size = 8

    def execute(self, requests):
        logger = pb_utils.Logger

        responses = []
        st = time.time()
        for request in requests:
            # Force English captioning regardless of client input.
            prompt = "<CAPTION_ENG>"

            image_bytes_list = pb_utils.get_input_tensor_by_name(request, "IMAGE_BYTES_LIST").as_numpy()
            image_bytes_list = np.array(image_bytes_list).reshape(-1)
            pil_images = []
            for image_bytes in image_bytes_list:
                image_bytes = base64.b64decode(image_bytes.decode('utf-8'))
                pil_image = Image.open(io.BytesIO(image_bytes))
                pil_images.append(pil_image)

            batch_list = []
            for i in range(0, len(pil_images), self.batch_size):
                batch_list.append(pil_images[i:i + self.batch_size])

            generated_text = []
            for batch in batch_list:
                with torch.no_grad():
                    inputs = self.processor(
                        text=[prompt] * len(batch),
                        images=batch,
                        return_tensors="pt",
                        padding=True,
                        do_resize=True
                    )
                    generated_ids = self.model.generate(
                        input_ids=inputs["input_ids"].to(self.device),
                        pixel_values=inputs["pixel_values"].to(self.device, dtype=self.dtype),
                        max_new_tokens=8,
                        num_beams=1,
                        do_sample=False
                    )
                    generated_text.extend(self.processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True))

            out_tensor_0 = pb_utils.Tensor("CAPTIONS", np.array(generated_text).astype(self.output_type))

            responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0]))

        logger.log_info(f"florence2 execute duration : {int((time.time() - st)*1000)} ms")
        return responses