import time
import argparse
import os
import onnxruntime_genai as og
from transformers import AutoTokenizer

class ONNXInferencer:

    def __init__(self, model_path: str, **kwargs):
        self.model_path = model_path

        self.model = og.Model(self.model_path)
        
        # Tokenizers
        self.safetensors_tokenizer = AutoTokenizer.from_pretrained(self.model_path)
        self.onnx_tokenizer = og.Tokenizer(self.model)

        # Collect generation/search options
        allowed = [
            "do_sample",
            "min_length",
            "top_p",
            "top_k",
            "temperature",
            "repetition_penalty",
        ]
        search_options = {name: kwargs[name] for name in allowed if name in kwargs}
        search_options["max_length"] = kwargs.get("max_length", 1024)

        self._onnx_params = og.GeneratorParams(self.model)
        self._onnx_params.set_search_options(**search_options)

    def inference(
        self, prompt: list[dict], return_latency: bool = False
    ) -> str | tuple:
        """
        Perform inference using the ONNX model.
        Args:
            prompt (str | list[dict]): The input prompt for the model. It can be a string or a list of dictionaries.
            return_latency (bool, optional): If True, return the latency and tokens per second (TPS) along with the response. Defaults to False.

        Returns:
            str | tuple: The generated response as a string. If return_latency is True, returns a tuple containing the response,
                 the time taken to generate the first token, and the tokens per second (TPS).
        """

        query = self.safetensors_tokenizer.apply_chat_template(
            prompt, tokenize=False, add_generation_prompt=True
        )
        generator = og.Generator(self.model, self._onnx_params)
        input_tokens = self.onnx_tokenizer.encode(query)
        generator.append_tokens(input_tokens)
        tokenizer_stream = self.onnx_tokenizer.create_stream()
        response = ""
        new_tokens = []
        initial_time = time.time()

        while not generator.is_done():
            generator.generate_next_token()
            new_token = generator.get_next_tokens()[0]
            new_tokens.append(new_token)
            new_text = tokenizer_stream.decode(new_token)
            if response == "":
                first_token_timestamp = time.time()
                time_first_token = first_token_timestamp - initial_time
            response += new_text

        run_time = time.time() - first_token_timestamp
        tps = len(new_tokens) / run_time
        if return_latency:
            return (response, time_first_token, tps)

        return response


text = "<insert-text>"

def main():
    parser = argparse.ArgumentParser(description="ONNX OOM test inferencer")
    parser.add_argument(
        "--model-name",
        type=str,
        default="qwen3_1-7b_onnx_cuda_int4",
        help="Model directory name (last part of the path under --models-base-dir)",
    )
    parser.add_argument(
        "--models-base-dir",
        type=str,
        default="./",
        help="Base directory containing ONNX model folders",
    )
    parser.add_argument(
        "--max-length",
        type=int,
        default=40960,
        help="Maximum generation length passed to ONNXInferencer",
    )
    args = parser.parse_args()

    model_path = os.path.join(args.models_base_dir, args.model_name)
    if not os.path.isdir(model_path):
        raise FileNotFoundError(f"Model path does not exist: {model_path}")

    llm = ONNXInferencer(model_path, max_length=args.max_length)
    safetensors_tokenizer = llm.safetensors_tokenizer

    system_prompt = "You are a helpful, respectful and honest assistant"
    prompt = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Summarize the following text:\n\nText: {text}"},
    ]
    print(f"\033[94mModel name: {args.model_name}\033[0m")
    print(
        f"\033[94mNumber of tokens: {len(safetensors_tokenizer.encode(prompt[1]['content']))}\033[0m"
    )
    output = llm.inference(prompt=prompt)
    print(f"\n\nResponse:\n{output}")


if __name__ == "__main__":
    main()