# score.py {{SCORE_VERSION}} generated by ADS {{ADS_VERSION}} on {{time_created}}
import os
import sys
import json
import subprocess
from functools import lru_cache
import onnxruntime as ort
import jsonschema
from jsonschema import validate, ValidationError
from transformers import AutoTokenizer
import logging

model_name = '{{model_file_name}}'
openapi_schema = 'openapi.json'


"""
   Inference script. This script is used for prediction by scoring server when schema is known.
"""


@lru_cache(maxsize=10)
def load_model(model_file_name=model_name):
    """
    Loads model from the serialized format

    Returns
    -------
    model:  a model instance on which predict API can be invoked
    """
    model_dir = os.path.dirname(os.path.realpath(__file__))
    if model_dir not in sys.path:
        sys.path.insert(0, model_dir)
    contents = os.listdir(model_dir)
    if model_file_name in contents:
        print(f'Start loading {model_file_name} from model directory {model_dir} ...')
        providers= ['CPUExecutionProvider']
        if is_gpu_available():
            providers=['CUDAExecutionProvider','CPUExecutionProvider']
        model = ort.InferenceSession(os.path.join(model_dir, model_file_name), providers=providers)
        print("Model is successfully loaded.")
        return model
    else:
        raise Exception(f'{model_file_name} is not found in model directory {model_dir}')


def is_gpu_available():
    """Check if gpu is available on the infrastructure."""
    try:
        result = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.returncode == 0:
           return True
    except FileNotFoundError:
        return False


@lru_cache(maxsize=1)
def load_tokenizer(model_full_name):

    model_dir = os.path.dirname(os.path.realpath(__file__))
    # initialize tokenizer
    return AutoTokenizer.from_pretrained(model_dir, clean_up_tokenization_spaces=True)

@lru_cache(maxsize=1)
def load_openapi_schema():
    """
    Loads the input schema for the incoming request

    Returns
    -------
    schema:  openapi schema as json
    """
    model_dir = os.path.dirname(os.path.realpath(__file__))
    if model_dir not in sys.path:
        sys.path.insert(0, model_dir)
    contents = os.listdir(model_dir)

    try:
        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), openapi_schema), 'r') as file:
            return json.load(file)
    except:
        raise Exception(f'{openapi_schema} is not found in model directory {model_dir}')


def validate_inputs(data):

    api_schema = load_openapi_schema()

    # use a reference resolver for internal $refs
    resolver = jsonschema.RefResolver.from_schema(api_schema)

    # get the actual schema part to validate against
    request_schema = api_schema["components"]["schemas"]["OpenAICompatRequest"]

    try:
        # validate the input JSON
        validate(instance=data, schema=request_schema, resolver=resolver)
    except ValidationError as e:
        example_value = {
            "input": ["What are activation functions?"],
            "encoding_format": "float",
            "model": "sentence-transformers/all-MiniLM-L6-v2",
            "user": "user"
        }
        message = f"JSON is invalid. Error: {e.message}\n An example of the expected format for 'OpenAICompatRequest' looks like: \n {json.dumps(example_value, indent=2)}"
        raise ValueError(message) from e


def pre_inference(data):
    """
    Preprocess data

    Parameters
    ----------
    data: Data format as expected by the predict API.

    Returns
    -------
    onnx_inputs: Data format after any processing
    total_tokens: total tokens that will be processed by the model

    """
    validate_inputs(data)

    tokenizer = load_tokenizer(data['model'])
    inputs = tokenizer(data['input'], return_tensors="np", padding=True)

    padding_token_id = tokenizer.pad_token_id
    total_tokens = (inputs["input_ids"] != padding_token_id).sum().item()
    onnx_inputs = {key: [l.tolist()for l in inputs[key] ] for key in inputs}

    return onnx_inputs, total_tokens

def convert_embeddings_to_openapi_format(embeddings, model_name, total_tokens):

    formatted_data = []
    openai_compat_response = {}
    for idx, embedding in enumerate(embeddings):

        formatted_embedding = {
            "object": "embedding",
            "embedding": embedding,
            "index": idx
        }
        formatted_data.append(formatted_embedding)

    # create the final OpenAICompatResponse format
    openai_compat_response = {
        "object": "list",
        "data": formatted_data,
        "model": model_name,  # Use the provided model name
        "usage": {
            "prompt_tokens": total_tokens,  # represents the token count for just the text input
            "total_tokens": total_tokens     # total number of tokens involved in the request, same in case of embeddings
        }
    }

    return openai_compat_response


def post_inference(outputs, model_name, total_tokens):
    """
    Post-process the model results

    Parameters
    ----------
    outputs: Data format after calling model.run
    model_name: name of model
    total_tokens: total tokens that will be processed by the model

    Returns
    -------
    outputs: Data format after any processing.

    """
    results = [embed.tolist() for embed in outputs]
    response = convert_embeddings_to_openapi_format(results, model_name, total_tokens)
    return response

def predict(data, model=load_model()):
    """
    Returns prediction given the model and data to predict

    Parameters
    ----------
    model: Model instance returned by load_model API.
    data: Data format as expected by the predict API of the core estimator. For eg. in case of sckit models it could be numpy array/List of list/Pandas DataFrame.

    Returns
    -------
    predictions: Output from scoring server
        Format: {'prediction': output from model.predict method}

    """
    # inputs contains 'input_ids', 'token_type_ids', 'attention_mask' but 'token_type_ids' is optional
    inputs, total_tokens = pre_inference(data)

    onnx_inputs = [inp.name for inp in model.get_inputs()]
    embeddings = model.run(None, {key: inputs[key] if key in inputs else None for key in onnx_inputs})[0]

    response = post_inference(embeddings, data['model'], total_tokens)
    return response
