import os
import threading
import shutil
from datetime import timedelta, datetime
from flask import Flask, render_template, request, session, jsonify, url_for, redirect
from haystack.document_store.elasticsearch import *
from haystack.preprocessor.utils import convert_files_to_dicts
from haystack.preprocessor.cleaning import clean_wiki_text
from haystack import Finder
from haystack.retriever.sparse import ElasticsearchRetriever
from haystack.reader.transformers import TransformersReader
from elasticsearch import Elasticsearch

es = (
    Elasticsearch()
)  # Replace with Elasticsearch(["http://elasticsearch:9200/"], verify_certs=True) to build docker image
session_time = 60  # Session Timeout in Minutes
app = Flask(__name__)
app.secret_key = "cbqa_123"
app.permanent_session_lifetime = timedelta(minutes=session_time)
user_id = 0  # User ID to keep track w.r.t sessions and context data
current_users = dict()  # Used to store user id with time of login
user_doc_store = dict()  # Document store object of the user id
user_settings = dict()  # User settings for GPU and Pre-trained models choice

# Handles pre-processing the context and uploads the pre-processed context to Elasticsearch
# Each user is assigned with a separate Elasticsearch index starting with "user_{user_id}"
# Documents & textual context are deleted from them temp folder named with user_id under users dir after uploading to Es
def pre_process(user_id_key):
    uploads_dir = "users/" + str(user_id_key) + "/uploads/"
    try:
        es_result = es.search(
            index="user_" + str(user_id_key), body={"query": {"match_all": {}}}
        )
        no_docs = len(es_result["hits"]["hits"])

    except Exception as e:
        print(e)
        print("\n no documents in es")

    processed = convert_files_to_dicts(
        dir_path=uploads_dir, clean_func=clean_wiki_text, split_paragraphs=True
    )

    for doc in range(len(processed)):
        try:
            # print("\n Checking for duplicate docs ..")
            add_doc = True
            for each_doc in range(no_docs):
                doc_text = es_result["hits"]["hits"][each_doc]["_source"]["text"]
                doc_name = es_result["hits"]["hits"][each_doc]["_source"]["name"]
                doc_id = es_result["hits"]["hits"][each_doc]["_id"]
                if (
                    processed[doc]["meta"]["name"] == "context_file.txt"
                    and doc_name == "context_file.txt"
                ):
                    # print("Deleting context file to update with new changes ..")
                    es.delete(
                        index="user_" + str(user_id_key), doc_type="_doc", id=doc_id
                    )

                if processed[doc]["text"] == doc_text:
                    # print("\n There is a duplicate, So this document is not added ..")
                    add_doc = False
                    os.remove(uploads_dir + str(processed[doc]["meta"]["name"]))
                    break
            if add_doc:
                # print("\n No duplicates found, so adding this to es..")
                processed_lst = [processed[doc]]
                user_doc_store[user_id_key].write_documents(processed_lst)
                os.remove(uploads_dir + str(processed[doc]["meta"]["name"]))

        except Exception as e:
            print(e)
            # print("\n no documents in es")
            processed_lst = [processed[doc]]
            user_doc_store[user_id_key].write_documents(processed_lst)
            os.remove(uploads_dir + str(processed[doc]["meta"]["name"]))


# Handles setting up reader and retriever
def set_finder(user_id_key):
    if user_settings[user_id_key]["model"] == "roberta":
        model_path = (
            "deepset/roberta-base-squad2"  # Path of the models hosted in Hugging Face
        )
    elif user_settings[user_id_key]["model"] == "bert":
        model_path = "deepset/bert-large-uncased-whole-word-masking-squad2"
    elif user_settings[user_id_key]["model"] == "distilbert":
        model_path = "distilbert-base-uncased-distilled-squad"
    else:
        model_path = "illuin/camembert-base-fquad"

    retriever = ElasticsearchRetriever(document_store=user_doc_store[user_id_key])

    if user_settings[user_id_key]["gpu"] == "on":
        try:
            reader = TransformersReader(
                model_name_or_path=model_path, tokenizer=model_path, use_gpu=0
            )
        except Exception as e:
            print(e)
            print("GPU not available. Inferencing on CPU")
            reader = TransformersReader(
                model_name_or_path=model_path, tokenizer=model_path, use_gpu=-1
            )

    else:
        reader = TransformersReader(
            model_name_or_path=model_path, tokenizer=model_path, use_gpu=-1
        )

    finder = Finder(reader, retriever)

    return finder


# Handles deletion of context data completely from the server after the session time ends and deletes user id from dict
def user_session_timer():
    global current_users, session_time
    seconds_in_day = 24 * 60 * 60
    print("\n User tracker thread started @ ", datetime.now())
    while True:
        for user_id_key in current_users.copy():
            current_time = datetime.now()
            user_time = current_users[user_id_key]
            difference = current_time - user_time
            time_diff = divmod(
                difference.days * seconds_in_day + difference.seconds, 60
            )
            if time_diff[0] >= session_time:
                try:
                    del current_users[user_id_key]
                    del user_doc_store[user_id_key]
                    del user_settings[user_id_key]
                    shutil.rmtree("users/" + str(user_id_key))
                    es.indices.delete(
                        index="user_" + str(user_id_key), ignore=[400, 404]
                    )
                except OSError as e:
                    print("Error: %s - %s." % (e.filename, e.strerror))
                # print("\n Deleted user:", user_id_key, " @", datetime.now())


session_timer = threading.Thread(target=user_session_timer)
session_timer.start()


# Handles users w.r.t new session or already in session
@app.route("/")
def home():
    global user_id, current_users, session_time
    logging.info(
        "User connected at "
        + str(datetime.now())
        + " with IP: "
        + str(request.environ["REMOTE_ADDR"])
    )

    if "user" in session and session["user"] in current_users:
        user_id = session["user"]
        logged_on = current_users[user_id]
        current_time = datetime.now()
        diff_min_sec = (
            int(datetime.strftime(current_time, "%M"))
            - int(datetime.strftime(logged_on, "%M"))
        ) * 60
        diff_sec = int(datetime.strftime(current_time, "%S")) - int(
            datetime.strftime(logged_on, "%S")
        )
        diff_time = diff_min_sec + diff_sec
        time_left = (
            session_time * 60
        ) - diff_time  # For session timeout on client side
        return render_template("index.html", time_left=time_left)

    else:
        session.permanent = True
        current_time = datetime.now()
        user_id += 1
        current_users[user_id] = current_time
        session["user"] = user_id
        # print(current_users)
        if not os.path.exists("users/"):  # Creating user temp dir for uploading context
            os.makedirs("users/" + str(user_id))
            os.makedirs("users/" + str(user_id) + "/uploads")
        else:
            os.makedirs("users/" + str(user_id))
            os.makedirs("users/" + str(user_id) + "/uploads")
        user_doc_store[user_id] = ElasticsearchDocumentStore(
            host="localhost", index="user_" + str(user_id)
        )  # Change host = "elasticsearch" to build docker image
        user_settings[user_id] = {
            "gpu": "off",
            "model": "roberta",
        }  # Initial user settings

        logged_on = current_users[user_id]
        current_time = datetime.now()
        diff_min_sec = (
            int(datetime.strftime(current_time, "%M"))
            - int(datetime.strftime(logged_on, "%M"))
        ) * 60
        diff_sec = int(datetime.strftime(current_time, "%S")) - int(
            datetime.strftime(logged_on, "%S")
        )
        diff_time = diff_min_sec + diff_sec
        time_left = (
            session_time * 60
        ) - diff_time  # For session timeout on client side

        return render_template("index.html", time_left=time_left)


# Handles context documents uploads
@app.route("/upload_file", methods=["GET", "POST"])
def upload_file():
    global current_users

    if "user" in session:
        user_id_key = session["user"]
        if user_id_key in current_users:
            for f in request.files.getlist("file"):
                f.save(
                    os.path.join("users/" + str(user_id_key) + "/uploads", f.filename)
                )

            pre_process(user_id_key)

            return render_template("index.html")

        else:
            return redirect(url_for("session_timeout"))

    else:
        return redirect(url_for("session_timeout"))


# Handles context added through the textbox
@app.route("/context", methods=["POST"])
def context():
    if "user" in session:
        user_id_key = session["user"]
        if user_id_key in current_users:
            text_context = request.form["context"]
            context_file = open(
                "users/" + str(user_id_key) + "/uploads/context_file.txt", "w"
            )
            context_file.write(text_context)
            context_file.close()
            pre_process(user_id_key)
            return jsonify({"output": "" + text_context})

        else:
            return render_template("session_out.html")

    else:
        return redirect(url_for("session_timeout"))


# Provides extracted answers for the posted question
@app.route("/question", methods=["POST"])
def question():
    if "user" in session:
        user_id_key = session["user"]
        if user_id_key in current_users:
            query_question = request.form["question"]
            es_stats = es.indices.stats(index="user_" + str(user_id_key))
            user_index_size = es_stats["_all"]["primaries"]["store"]["size_in_bytes"]
            if (
                user_index_size == 208
            ):  # To check if index in Es is empty. 208 bytes is default index size without docs
                return jsonify({"error": "add context"})

            finder = set_finder(user_id_key)
            answers_dict = finder.get_answers(
                question=query_question, top_k_retriever=5, top_k_reader=5
            )
            unique_answers = list()
            output = list()
            if len(answers_dict["answers"]) > 0:
                for i in range(len(answers_dict["answers"])):
                    if (
                        answers_dict["answers"][i]["answer"] is not None
                        and answers_dict["answers"][i]["answer"] not in unique_answers
                    ):
                        temp_dict = answers_dict["answers"][i]
                        remove = (
                            "score",
                            "probability",
                            "offset_start",
                            "offset_end",
                            "document_id",
                        )
                        unique_answers.append(temp_dict["answer"])
                        if temp_dict["meta"]["name"] == "context_file.txt":
                            temp_dict["meta"]["name"] = "Textual Context"
                        temp_dict["meta"] = temp_dict["meta"]["name"]
                        output.append(temp_dict)
                        for key in remove:
                            if key in temp_dict:
                                del temp_dict[key]

            else:
                output = [
                    {"answer": "No Answers found ..", "context": " ", "meta": " "},
                ]

            return jsonify({"output": output})

        else:
            return render_template("session_out.html")


# Handles GPU setting changes.
@app.route("/gpu", methods=["POST"])
def gpu():
    if "user" in session:
        user_id_key = session["user"]
        if user_id_key in current_users:
            if user_settings[user_id_key]["gpu"] == "on":
                user_settings[user_id_key]["gpu"] = "off"
            else:
                user_settings[user_id_key]["gpu"] = "on"
    return jsonify({"output": "gpu status changed"})


# Handles pre-trained model choice setting changes.
@app.route("/models", methods=["POST"])
def models():
    if "user" in session:
        user_id_key = session["user"]
        if user_id_key in current_users:
            user_settings[user_id_key]["model"] = request.form["model"]
    return jsonify({"output": "model changed"})


# Handles session timeout redirection
@app.route("/session_timeout")
def session_timeout():
    return render_template("session_out.html")


# Handles removing of session identifier from session dict, This works only when app tab is open until session completes
@app.route("/session_out", methods=["POST"])
def session_out():
    session.pop("user", None)
    return redirect(url_for("session_timeout"))


# Comment the below block in case of building a docker image or running on WSGI server like gunicorn
if __name__ == "__main__":
    app.run(host="0.0.0.0")

[(15, 'elasticsearch.Elasticsearch', 'Elasticsearch', False, 'from elasticsearch import Elasticsearch\n'), (18, 'flask.Flask', 'Flask', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (20, 'datetime.timedelta', 'timedelta', False, 'from datetime import timedelta, datetime\n'), (145, 'threading.Thread', 'threading.Thread', False, 'import threading\n'), (41, 'haystack.preprocessor.utils.convert_files_to_dicts', 'convert_files_to_dicts', False, 'from haystack.preprocessor.utils import convert_files_to_dicts\n'), (94, 'haystack.retriever.sparse.ElasticsearchRetriever', 'ElasticsearchRetriever', False, 'from haystack.retriever.sparse import ElasticsearchRetriever\n'), (113, 'haystack.Finder', 'Finder', False, 'from haystack import Finder\n'), (325, 'flask.jsonify', 'jsonify', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (335, 'flask.jsonify', 'jsonify', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (341, 'flask.render_template', 'render_template', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (347, 'flask.session.pop', 'session.pop', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (109, 'haystack.reader.transformers.TransformersReader', 'TransformersReader', False, 'from haystack.reader.transformers import TransformersReader\n'), (122, 'datetime.datetime.now', 'datetime.now', False, 'from datetime import timedelta, datetime\n'), (163, 'datetime.datetime.now', 'datetime.now', False, 'from datetime import timedelta, datetime\n'), (175, 'flask.render_template', 'render_template', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (179, 'datetime.datetime.now', 'datetime.now', False, 'from datetime import timedelta, datetime\n'), (199, 'datetime.datetime.now', 'datetime.now', False, 'from datetime import timedelta, datetime\n'), (212, 'flask.render_template', 'render_template', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (348, 'flask.url_for', 'url_for', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (98, 'haystack.reader.transformers.TransformersReader', 'TransformersReader', False, 'from haystack.reader.transformers import TransformersReader\n'), (125, 'datetime.datetime.now', 'datetime.now', False, 'from datetime import timedelta, datetime\n'), (184, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (223, 'flask.request.files.getlist', 'request.files.getlist', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (230, 'flask.render_template', 'render_template', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (236, 'flask.url_for', 'url_for', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (252, 'flask.jsonify', 'jsonify', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (255, 'flask.render_template', 'render_template', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (258, 'flask.url_for', 'url_for', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (309, 'flask.jsonify', 'jsonify', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (312, 'flask.render_template', 'render_template', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (104, 'haystack.reader.transformers.TransformersReader', 'TransformersReader', False, 'from haystack.reader.transformers import TransformersReader\n'), (168, 'datetime.datetime.strftime', 'datetime.strftime', False, 'from datetime import timedelta, datetime\n'), (169, 'datetime.datetime.strftime', 'datetime.strftime', False, 'from datetime import timedelta, datetime\n'), (204, 'datetime.datetime.strftime', 'datetime.strftime', False, 'from datetime import timedelta, datetime\n'), (205, 'datetime.datetime.strftime', 'datetime.strftime', False, 'from datetime import timedelta, datetime\n'), (233, 'flask.url_for', 'url_for', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (273, 'flask.jsonify', 'jsonify', False, 'from flask import Flask, render_template, request, session, jsonify, url_for, redirect\n'), (165, 'datetime.datetime.strftime', 'datetime.strftime', False, 'from datetime import timedelta, datetime\n'), (166, 'datetime.datetime.strftime', 'datetime.strftime', False, 'from datetime import timedelta, datetime\n'), (201, 'datetime.datetime.strftime', 'datetime.strftime', False, 'from datetime import timedelta, datetime\n'), (202, 'datetime.datetime.strftime', 'datetime.strftime', False, 'from datetime import timedelta, datetime\n'), (155, 'datetime.datetime.now', 'datetime.now', False, 'from datetime import timedelta, datetime\n')]