import os
import time
from typing import Any

import numpy as np
import triton_python_backend_utils as pb_utils  # type: ignore
from transformers import AutoTokenizer


class TritonPythonModel:
    @staticmethod
    def _decode_text(value: Any) -> str:
        if isinstance(value, np.bytes_):
            return value.tobytes().decode("utf-8")
        if isinstance(value, bytes):
            return value.decode("utf-8")
        return str(value)

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

        self.logger = pb_utils.Logger
        self.max_length = int(os.environ.get("BGE_M3_MAX_LENGTH", "128"))

        model_repo = args["model_repository"]
        tokenizer_path = os.path.join(model_repo, "1", "tokenizer")
        if not os.path.isdir(tokenizer_path):
            raise pb_utils.TritonModelException(f"tokenizer directory not found: {tokenizer_path}")

        self.tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_path,
            local_files_only=True,
        )
        self.logger.log_info(f"bge_m3_preprocess initialized. tokenizer_path={tokenizer_path}")

    def execute(self, requests):
        responses = []
        st = time.time()

        for request in requests:
            input_tensor = pb_utils.get_input_tensor_by_name(request, "TEXTS")
            texts_np = input_tensor.as_numpy().reshape(-1)
            texts = [self._decode_text(v) for v in texts_np]

            if len(texts) == 0:
                input_ids = np.empty((0, 1), dtype=np.int64)
                attention_mask = np.empty((0, 1), dtype=np.int64)
            else:
                encoded = self.tokenizer(
                    texts,
                    return_tensors="np",
                    padding=True,
                    truncation=True,
                    max_length=self.max_length,
                )
                input_ids = np.asarray(encoded["input_ids"], dtype=np.int64)
                attention_mask = np.asarray(encoded["attention_mask"], dtype=np.int64)

            responses.append(
                pb_utils.InferenceResponse(
                    output_tensors=[
                        pb_utils.Tensor("input_ids", input_ids),
                        pb_utils.Tensor("attention_mask", attention_mask),
                    ]
                )
            )

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