# from transformers import GPT2TokenizerFast
import tiktoken
from openai import AzureOpenAI
import openai
import os
from time import sleep
import dotenv
# loads .env file with your OPENAI_API_KEY
dotenv.load_dotenv()

# tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
tokenizer = tiktoken.get_encoding("cl100k_base")
def get_embedding(text, engine=os.getenv("EMBEDDING_ENGINE")):
    count = 0
    error = ''
    while count < 5:
        try:
            if count > 0:
                print(f'retrying the {count} time for getting text embedding...')
            return openai.Embedding.create(input=text, engine=engine)['data'][0]['embedding']
        except Exception as e:
            error = str(e)
            print(error)
            count += 1
            sleep(1)


def num_tokens_from_string(string: str) -> int:
    """Returns the number of tokens in a text string."""
    num_tokens = len(tokenizer.encode(string))
    return num_tokens


def num_tokens_from_messages(messages, model="gpt-35-turbo-16k"):
    """Returns the number of tokens used by a list of messages."""
    encoding = tiktoken.get_encoding("cl100k_base")
    if model == "gpt-35-turbo-16k":
        tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
        tokens_per_name = -1  # if there's a name, the role is omitted
    elif model == "gpt-4-32k":
        tokens_per_message = 3
        tokens_per_name = 1
    else:
        tokens_per_message = 3
        tokens_per_name = 1
    num_tokens = 0
    for message in messages:
        num_tokens += tokens_per_message
        for key, value in message.items():
            num_tokens += len(encoding.encode(value))
            if key == "name":
                num_tokens += tokens_per_name
    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
    return num_tokens


def chat(prompt: str, 
         engine = os.getenv("Engine_GPT4o"), 
         azure_endpoint=os.getenv("OPENAI_API_BASE_GPT4o"),
         api_key=os.getenv("OPENAI_API_KEY_GPT4o"),
         api_version=os.getenv("OPENAI_API_VERSION_GPT4o"),
         temperature: float = 0.0):
    client = AzureOpenAI(
        azure_endpoint=azure_endpoint,
        api_key=api_key,
        api_version=api_version
    )

    count = 0
    error = ''
    max_tokens = 4000
    request_timeout = 120
    while count < 8:
        try:
            if count > 0:
                print(f'retrying the {count} time...')
            response  = client.chat.completions.create(
                    model=engine,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    top_p=0.95,
                    frequency_penalty=0,
                    presence_penalty=0,
                    timeout=request_timeout,
                    stop=None,
                    messages=[
                        {"role": "user", "content": prompt}
                    ]
                )
            return response.choices[0].message.content, False
        except Exception as e:
            error = str(e)
            print(f"error message: {error}")
            if 'maximum context length' in error:
                return error, True
            count += 1
        sleep(3)
    return error, True