import re
import openai
from openai.embeddings_utils import distances_from_embeddings
import json
import os
import codecs
from dotenv import load_dotenv
load_dotenv()
openai.api_key = os.getenv('OPENAI_API_KEY_ML')


def remove_comments(string):
    """
    This WILL remove:

    - /* multi-line comments */
    - // single-line comments
    
    Will NOT remove:
    - String var1 = "this is /* not a comment. */";
    - char *var2 = "this is // not a comment, either.";
    - url = 'http://not.comment.com';
    
    Note: This will also work for Javascript source."""
    
    pattern = r"(\".*?\"|\'.*?\')|(/\*.*?\*/|//[^\r\n]*$)"
    # first group captures quoted strings (double or single)
    # second group captures comments (//single-line or /* multi-line */)
    regex = re.compile(pattern, re.MULTILINE|re.DOTALL)
    def _replacer(match):
        # if the 2nd group (capturing comments) is not None,
        # it means we have captured a non-quoted (real) comment string.
        if match.group(2) is not None:
            return "" # so we will return empty to remove the comment
        else: # otherwise, we will return the 1st group
            return match.group(1) # captured quoted-string
    return regex.sub(_replacer, string)



# function to create context (the content from which the chatbot should find the answer)
def create_context(question, df, max_len=1800, size="ada"):
    
    """
    Create a context for a question by finding the most similar context from the dataframe
    """

    # Get the embeddings for the question
    q_embeddings = openai.Embedding.create(input=question, engine='text-embedding-ada-002')['data'][0]['embedding']

    # Get the distances from the embeddings
    df['distances'] = distances_from_embeddings(q_embeddings, df['embeddings'].values, distance_metric='cosine')


    returns = []
    cur_len = 0

    # Sort by distance and add the text to the context until the context is too long
    for i, row in df.sort_values('distances', ascending=True).iterrows():
        
        # Add the length of the text to the current length
        cur_len += row['n_tokens'] + 4
        
        # If the context is too long, break
        if cur_len > max_len:
            break
        
        # Else add it to the text that is being returned
        returns.append(row["text"])

    # Return the context
    return "\n\n###\n\n".join(returns)



# Function to answer a question
def answer_question(
    df,
    model="gpt-3.5-turbo",
    question="Am I allowed to publish model outputs to Twitter, without a human review?",
    instructions= "Responde la siguiente pregunta basándote en el contexto dado, y si la pregunta no puede ser contestada basándote en el contexto, dí \"Lo siento, sólo puedo contestar preguntas que refieran a las soluciones de Mercado Ads.\"",
    max_len=1800,
    size="ada",
    temperature=0,
    max_tokens=300,
    stop_sequence=None
):
    """
    Answer a question based on the most similar context from the dataframe texts
    """
    context = create_context(question, df, max_len=max_len, size=size,)

    if(model=="gpt-3.5-turbo" or model=="gpt-4"):
        try:
            # Create a completions using the question and context
            response = openai.ChatCompletion.create(
                messages=[
                {"role": "user", "content": f"{instructions}\n\nContexto: {context}\n\n---\n\nPregunta: {question}\nRespuesta:",}
                ],
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=1,
                frequency_penalty=0,
                presence_penalty=0,
                stop=stop_sequence,
                model=model,
            )

            output = {
                "answer": response['choices'][0]['message']['content'],
                "context": context,
                "complete instructions": instructions + "Pregunta: " + question,
                "total_tokens": response['usage']['total_tokens'],
                "completion_tokens": response['usage']['completion_tokens'],
                "prompt_tokens": response['usage']['prompt_tokens'],            
                }
            return output
            
        except Exception as e:
            print(e)
            return {
                "error": True,
                "message": str(e)
            }
        
    elif(model=="text-davinci-003"):
        try:
            # Create a completions using the question and context
            response = openai.Completion.create(
                prompt=f"{instructions}\n\nContexto: {context}\n\n---\n\nPregunta: {question}\nRespuesta:",
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=1,
                frequency_penalty=0,
                presence_penalty=0,
                stop=stop_sequence,
                model=model,
            )

            output = {
                "answer": response["choices"][0]["text"].strip(),
                "context": context,
                "complete_instructions": instructions + "Pregunta: " + question,
                "total_tokens": response['usage']['total_tokens'],
                "completion_tokens": response['usage']['completion_tokens'],
                "prompt_tokens": response['usage']['prompt_tokens'],            
                }
            return output
            
        except Exception as e:
            print(e)
            return {
                "error": True,
                "message": str(e)
            }
    else:
        print("Model not found")
        return {
            "error": True,
            "message": "Model not found"
        }