import argparse

from utils import get_embeddings, search_document_annoy, \
    answer_with_gpt3_with_function_calls, transform_user_question, debug_print

def main(args):
    """
    - Get & embed user's question
    - Find top choices of documents based on embedding search
    - Generate response
    """
    top_k = args.top_k
    annoy_metric = args.annoy_metric

    # Get & embed user's question
    while True:
        print("Hi! What question do you have for ANGR? press 0 to exit")
        user_input = input()
        if user_input == '0':
            break

        assert top_k > 0, 'k must be a integer greater than 0'

        if args.user_query_preprocess:
            chatgpt_question = transform_user_question(user_input, args.model)
        else:
            chatgpt_question = user_input
        debug_print("chatgpt_question: ", chatgpt_question)

        try:
            user_q_embedding = get_embeddings(chatgpt_question)
            document = search_document_annoy(user_q_embedding, top_k=top_k, metric=annoy_metric)
            reply = answer_with_gpt3_with_function_calls(document, user_input, args.model)
            print(reply)
        except Exception as e:
            print(e)
            print("Error when trying to get embedding for the user query. Please try with a shorter question.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--top_k', default=3, type=int, help="Number of documents to retrieve")
    parser.add_argument('--annoy_metric', default='angular',
                        choices=['angular', 'euclidean', 'manhattan', 'hamming', 'dot'],
                        help="metric for annoy algorithm")
    parser.add_argument('--model', default="gpt-3.5-turbo")
    parser.add_argument('--user_query_preprocess', default=False, help="Whether to ask ChatGPT to generate an additional query for itself based on user's question")
    parser.add_argument('--debug', default=False, help="define whether to print the debug statement")
    args = parser.parse_args()

    main(args)

"""
python main.py --top_k 3 --annoy_metric dot --user_query_preprocess True

Questions:
- list math-related publications
- list all publications
- Give me a publication written by J Coleman
    -  Could you please provide me with a publication authored by J Coleman?
    -  Can you find a publication authored by J Coleman? 
- Give me the link to Jared Coleman's homepage

"""