import logging
import os
from time import asctime

import gradio as gr
from llama_index.core import Document, VectorStoreIndex

from generate_response import generate_chat_response_with_history, set_llm, is_search_query, \
    generate_chat_response_with_history_rag_return_response, google_question
from read_write_index import read_write_index
from web_search import search

logger = logging.getLogger("agent_logger")
sourced = False
query = False
rag_similarity = False


def google_search_chat(message, history):
    gquestion = google_question(message, history)
    if is_search_query(gquestion):
        search_results = search(message, gquestion)
        print(f'Search results returned: {len(search_results)}')

        relevant_content = ""
        for result in search_results:
            relevant_content = relevant_content + "\n" + ''.join(result['text'])

        print(f'Relevant content char length: {len(relevant_content)}')
        if relevant_content != "":
            documents = [Document(text=relevant_content)]
            index = VectorStoreIndex.from_documents(documents)
            print('Search results vectorized...')
            response = generate_chat_response_with_history_rag_return_response(index, message, history)
        else:
            print(f'Assistant Response: Sorry, no search results found, trying with offline resources.')
            index = read_write_index(path='storage_search/')
            response = generate_chat_response_with_history_rag_return_response(index, message, history)
        response_text = []
        string_output = ""

        for text in response.response_gen:
            response_text.append(text)
            string_output = ''.join(response_text)
            yield string_output

        print(f'Assistant Response: {string_output}')

    else:
        yield from generate_chat_response_with_history(message, history)


if __name__ == '__main__':
    logging.root.setLevel(logging.INFO)
    filehandler = logging.FileHandler(f'agent_log_{asctime().replace(" ", "").lower().replace(":", "")}.log',
                                      'a')
    formatter = logging.Formatter('%(asctime)-15s::%(levelname)s::%(filename)s::%(funcName)s::%(lineno)d::%(message)s')
    filehandler.setFormatter(formatter)
    logger = logging.getLogger("agent_logger")
    for hdlr in logger.handlers[:]:  # remove the existing file handlers
        if isinstance(hdlr, logging.FileHandler):
            logger.removeHandler(hdlr)
    logger.addHandler(filehandler)  # set the new handler
    logger.setLevel(logging.INFO)

    api_key = os.getenv('gpt_api_key')

    # GPT - 4 Turbo. The latest GPT - 4 model intended to reduce cases of “laziness” where the model doesn’t complete
    # a task. Returns a maximum of 4,096 output tokens. Link:
    # https://openai.com/blog/new-embedding-models-and-api-updates
    set_llm(key=api_key, model="gpt-4-0125-preview", temperature=0)

    print("Launching Gradio ChatInterface for searchbot...")

    demo = gr.ChatInterface(fn=google_search_chat,
                            title="Search Assistant", retry_btn=None, undo_btn=None, clear_btn=None,
                            theme="soft")
    demo.launch()
#   auth=('convo', 'session2024')