import openai, tiktoken, os
import pandas as pd
from openai.embeddings_utils import get_embedding, cosine_similarity
# embedding model parameters
embedding_model = "text-embedding-ada-002"
embedding_encoding = "cl100k_base"  # this the encoding for text-embedding-ada-002
max_tokens = 8000  # the maximum for text-embedding-ada-002 is 8191
encoding = tiktoken.get_encoding("cl100k_base")

def set_openai_api_key(api_key):
    if api_key and api_key.startswith("sk-") and len(api_key) > 50:
        openai.api_key = api_key
    else:
        raise gr.Error("OpenAI API key incorrect.")

# Prepare prompt
def prepare_prompt(prompt, results):
    tokens_limit = 16000 # Limit for gpt-3.5-turbo-16k    
    
    user_start = (
        "請只根據下述的文本，使用繁體中文(zh-TW)回答問題。排除相似重複的語意，用字精鍊而清晰，出現過一次的人名就不需要一直顯示全名以及完整職稱。請嚴格遵守只根據下述的文本的規定，如果詢問的問題超過文本的範圍，請回答你不知道。\n\n"+
        "文本：\n"
    )
    user_end = (
        f"\n\n問題： {prompt}\n 答案："
    )
    
    system = """
    你是一個萬能文字助手，你擅長從大量的文章中，辨識出相關主題，並整理成重點摘要。
    """

    count_of_tokens_consumed = len(encoding.encode("\"role\":\"system\"" +  
                                                   "\"content\" :\"" + system + 
                                                   user_start + "\n\n---\n\n" + user_end ))
    
    count_of_tokens_for_context = tokens_limit - count_of_tokens_consumed

    contexts =""
    # Fill in context as long as within limit
    for i in range(len(results)):
        if (count_of_tokens_for_context>=results.n_tokens.iloc[i]):
            contexts += results.text.iloc[i] + "\n"
            count_of_tokens_for_context -=1
            count_of_tokens_for_context -= results.n_tokens.iloc[i]
    
            
    complete_prompt = user_start + contexts + "\n\n---\n\n" + user_end
    return complete_prompt


def answer(messages):
    response = openai.ChatCompletion.create(
              model="gpt-3.5-turbo-16k",
              messages=messages,
              temperature=0, 
              stream=True
          )
    return response