conjurer/conjurer_librarian/conjurer_librarian.py

"""
This module contains the implementation of the Librarian class and related functions for searching and refining queries.

Classes:
- Librarian: Represents a librarian object that performs search and refinement operations on queries.

Functions:
- flask_debug: Starts a Flask application in debug mode without using the reloader.
- waitress_run: Serves the Flask application using the Waitress WSGI server.
- BackgroundTaskSearch: Represents a background task for running the Librarian object asynchronously.
"""

import asyncio
import json
import logging
import netrc
import re
import threading
import time
from json.decoder import JSONDecodeError
from logging import handlers
from queue import Queue
from urllib.request import urlopen

import requests
import scrape_bot
import search_bot
from flask import Flask, jsonify, request
from habanero import Crossref
from waitress import serve

# Constants
NETRC_FILE = r"C:\Users\Activcom.pl\.netrc"
HOST_ADDRESS = "192.168.1.192"
PORT_ADDRESS = 5001
MAIN_BOT_ADDRESS = "http://192.168.1.191:5000"
SEND_RESULTS = "/conjurer"
BDSM_UUID_TEST = "96b7f85a-1142-4908-8986-62a2ea25a147"
MAX_CR_RESULTS = 150000
ENCODING = "utf-8"

app = Flask(__name__)

librarian_queue = Queue()
librarian_list = []


class Librarian(object):
    def __init__(self, _app, query, uuid) -> None:
        """
        Initializes a Librarian object.

        Args:
        - _app: The Flask application object.
        - query: The query to be searched.
        - uuid: The unique identifier for the search.

        Attributes:
        - cr: The Crossref object for performing the search.
        - query: The query to be searched.
        - uuid: The unique identifier for the search.
        - limit: The maximum number of search results to fetch.
        - fetched: The number of search results fetched so far.
        - hit: The number of search results that match the refinement criteria.
        - total: The total number of search results.
        - app: The Flask application object.
        - live_results: A list to store live search results.
        - final_result: A list to store the final refined search results.
        - not_in_db: A list to store search results that are not in the local database.
        - search_result_from_cr: A dictionary to store the search results from Crossref.
        - done: A flag indicating if the search is done.
        """
        netrc_mod = netrc.netrc(NETRC_FILE)
        authTokens = netrc_mod.authenticators("crossref")
        self.cr = Crossref(mailto=authTokens[0])
        self.query = query
        self.uuid = str(uuid)
        self.limit = MAX_CR_RESULTS
        self.fetched = 0
        self.hit = 0
        self.total = 0
        self.app = _app
        self.live_results = []
        self.final_result = []
        self.not_in_db = []
        self.search_result_from_cr = {}
        self.done = False

    async def search_crossref(self, query):
        """
        Performs a search on Crossref for the given query.

        Args:
        - query: The query to be searched.

        Returns:
        - result: The search result from Crossref.

        Raises:
        - None.
        """
        self.app.logger.info("STARTED SEARCH")
        result = self.cr.works(query=query, limit=1000)
        self.search_result_from_cr.update(result)
        self.total = result["message"]["total-results"]
        self.fetched += len(result["message"]["items"])
        self.app.logger.info(self.total)
        self.app.logger.info(self.fetched)
        while self.total > self.fetched and self.limit > self.fetched:
            result = self.cr.works(query=query, limit=1000)
            self.search_result_from_cr.update(result)
            self.total = result["message"]["total-results"]
            self.fetched += len(result["message"]["items"])
            self.app.logger.info(self.total)
            self.app.logger.info(self.fetched)
            time.sleep(0.1)
        self.app.logger.info("CROSSREF DONE")
        # TODO: add to sql with UUID
        # sql row - title, doi, author, uuid, last retrieved, is_available

        with open("rr_results.json", "r+", encoding="utf-8") as data_file:
            # First we load existing data into a dict.
            try:
                file_data = json.load(data_file)
            except JSONDecodeError:
                file_data = {}
            summarized_results = []
            self.app.logger.info("REFINE: Removing all derived works from the list")
            for item in self.search_result_from_cr["message"]["items"]:
                container = "NO"
                if "container-title" in item:
                    container = "YES"
                else:
                    summarized_results.append(
                        {
                            "DOI": item["DOI"],
                            "title": item["title"] if "title" in item else None,
                            "type": item["type"] if "type" in item else None,
                            "container": container,
                        }
                    )
            result = {
                self.uuid: {
                    "total_results": self.search_result_from_cr["message"][
                        "total-results"
                    ],
                    "on_page": 1,
                    "summary": summarized_results,
                    "results": self.search_result_from_cr["message"]["items"],
                }
            }
            if file_data:
                file_data.update(result)
            else:
                file_data = result

            self.app.logger.info("REFINE: Dumping to file")
            data_file.truncate(0)
            data_file.seek(0)
            json.dump(file_data, data_file, indent=4)
            return result

    async def check_dois_in_local_db(self, doi):
        """
        Checks if the given DOI exists in the local database.

        Args:
        - doi: The DOI to be checked.

        Returns:
        - result: The search result from the backend app.

        Raises:
        - None.
        """
        self.app.logger.info("REFINE: Running search in the backend app")
        coro = asyncio.to_thread(
            search_bot.search_for_doi, doi, self.live_results, self.app.logger
        )
        result = await coro
        return result

    async def check_if_exists_brute_force(self, page_url):
        """
        Checks if the given page URL exists using brute force method.

        Args:
        - page_url: The URL of the page to be checked.

        Returns:
        - exists: A boolean indicating if the page exists.

        Raises:
        - ValueError: If the URL does not start with 'http:' or 'https:'.
        """
        self.app.logger.error("REFINE: Brute force search!")
        if not page_url.startswith(("http:", "https:")):
            raise ValueError("URL must start with 'http:' or 'https:'")
        # trunk-ignore(bandit/B310)
        with urlopen(page_url) as response:
            data = response.read()
            text = data.decode("utf-8")
            for line in text.splitlines():
                if re.match(
                    r".*Unfortunately, Sci-Hub doesn't have the requested document.*",
                    line,
                ):
                    return False
                if m := re.match(
                    r".*<embed type=\"application.pdf\"\s*src=\"(.*\.pdf)", line
                ):
                    direct_download_link = "https:" + str(m.group(1))
                    self.app.logger.info(direct_download_link)
                    return True
        return False

    async def check_if_exists(self, doi, brute_force=False):
        """
        Checks if the given DOI exists.

        Args:
        - doi: The DOI to be checked.
        - brute_force: A flag indicating if brute force method should be used.

        Returns:
        - result: The search result.

        Raises:
        - None.
        """
        result = {}
        result = await self.check_dois_in_local_db(doi)
        if brute_force and not result:
            for item in doi:
                item_link = "https://sci-hub.se/" + item[0]
                tmp = {
                    "DOI": item["DOI"],
                    "exists": await self.check_if_exists_brute_force(item_link),
                    "data": item,
                }
                result.update(tmp)
        return result

    async def refine_query(self, unrefined_result):
        """
        Refines the search query based on the unrefined search result.

        Args:
        - unrefined_result: The unrefined search result.
        - brute_force: A flag indicating if brute force method should be used.

        Returns:
        - final_result: The final refined search result.

        Raises:
        - None.
        """
        temp = []
        refined_result = []
        for key in unrefined_result:
            self.app.logger.info("KEY:")
            self.app.logger.info(key)

        for item in unrefined_result[self.uuid]["summary"]:
            if item["container"] == "NO" and item["title"]:
                refined_result.append(item)

        for item in refined_result:
            temp.append([item["DOI"], item])

        refined_result = await self.check_if_exists(temp)

        for item in refined_result:
            if item["exists"]:
                self.final_result.append(item["data"])
            else:
                self.not_in_db.append(item["data"])
        self.hit = len(self.final_result)
        return self.final_result

    async def answer_query(self):
        """
        Answers the search query.

        Args:
        - None.

        Returns:
        - result: The search result.

        Raises:
        - None.
        """
        self.app.logger.info(f"Search started {self.uuid}")
        database = None
        query = self.query
        with open("cr_results.json", "r+", encoding="utf-8") as cr_file:
            try:
                database = json.load(cr_file)
            except JSONDecodeError:
                pass
            if database:
                for item in database.keys():
                    if item in self.uuid:
                        self.result = await self.refine_query(database)
                        temp = {}
                        for item in self.result:
                            temp[item["DOI"]] = {
                                "Title": item["title"],
                                "type": item["type"],
                            }
                        self.result = temp
                        temp2 = {}
                        for item in self.not_in_db:
                            temp2[item["DOI"]] = {
                                "Title": item["title"],
                                "type": item["type"],
                            }
                        self.not_in_db = temp2
                        return self.result
                answer = await self.search_crossref(query=query)
                database.update(answer)
                cr_file.truncate(0)
                cr_file.seek(0)
                json.dump(database, cr_file)
                self.result = await self.refine_query(answer)
            else:
                answer = await self.search_crossref(query=query)
                database = answer
                cr_file.truncate(0)
                cr_file.seek(0)
                json.dump(database, cr_file)
                self.result = await self.refine_query(answer)
                self.app.logger.info("Refined")
                self.done = True
            temp = {}
            for item in self.result:
                temp[item["DOI"]] = {"Title": item["title"], "type": item["type"]}
            self.result = temp
            temp2 = {}
            for item in self.not_in_db:
                temp2[item["DOI"]] = {"Title": item["title"], "type": item["type"]}
                self.not_in_db = temp2
            return self.result


# ============================= FLASK INTERNALS===============================


def flask_debug():
    """
    Starts a Flask application in debug mode without using the reloader.

    Args:
    - None.

    Returns:
    - None.

    Raises:
    - None.
    """
    # trunk-ignore(bandit/B201)
    app.run(debug=True, use_reloader=False, host=HOST_ADDRESS, port=PORT_ADDRESS)


def waitress_run():
    """
    Serves the Flask application using the Waitress WSGI server.

    Args:
    - None.

    Returns:
    - None.

    Raises:
    - None.
    """
    serve(app, host=HOST_ADDRESS, port=PORT_ADDRESS)


class BackgroundTaskSearch(threading.Thread):
    """
    A background task for searching and saving results to files.

    This class extends the `threading.Thread` class and is responsible for running
    the search task in the background. It retrieves queries from a queue, performs
    the search, and saves the results to files.

    Attributes:
        app (App): The application instance.
    """

    def run(self):
        """
        Run the background task.

        This method is called when the thread is started. It creates a new event loop,
        runs the `_run` method, and closes the event loop.
        """
        loop = asyncio.new_event_loop()
        loop.run_until_complete(self._run())
        loop.close()

    async def _run(self):
        """
        Perform the search task.

        This method is an asynchronous coroutine that runs in a loop. It retrieves a
        librarian from the queue, answers the query, and saves the results to files.
        It also sends the results to a remote server.

        The search task continues running indefinitely until the thread is stopped.
        """
        while True:
            database = None
            ndb_database = None
            librarian = librarian_queue.get()
            self.app.logger.info("STARTED")
            result = await librarian.answer_query()
            result = {librarian.uuid: result}
            self.app.logger.info("Saving to file")

            # Save results to "not_in_db.json" file
            with open("not_in_db.json", "r+", encoding="utf-8") as ndb_file:
                try:
                    ndb_database = json.load(ndb_file)
                except JSONDecodeError:
                    pass
                if ndb_database:
                    ndb_database.update(librarian.not_in_db)
                else:
                    ndb_database = librarian.not_in_db
                ndb_file.truncate(0)
                ndb_file.seek(0)
                json.dump(ndb_database, ndb_file)

            # Save results to "s_results.json" file
            with open("s_results.json", "r+", encoding="utf-8") as s_file:
                try:
                    database = json.load(s_file)
                except JSONDecodeError:
                    pass
                if database:
                    self.app.logger.info(database)
                    self.app.logger.info(result)
                    database.update(result)
                else:
                    database = result
                self.app.logger.info("DUMPING DATA")
                s_file.truncate(0)
                s_file.seek(0)
                json.dump(database, s_file)
            self.app.logger.info("FINISHED")

            self.app.logger.info(result)
            coroutine = asyncio.to_thread(
                requests.post,
                f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}",
                json=result,
                timeout=360,
            )
            self.app.logger.info("SENT")
            result = await coroutine
            self.app.logger.info(result.status_code)
            self.app.logger.info("SEND CONFIRMED")
            time.sleep(1)


# ==================================SERVER ROUTES==========================================
@app.route("/query", methods=["POST"])
async def query_database():
    """
    Endpoint for querying the database.

    This function receives a POST request containing a JSON payload with a query and a UUID.
    It creates a Librarian object with the query and UUID, and adds it to the librarian_queue and librarian_list.
    Finally, it returns a JSON response indicating the success of the operation, along with the query, UUID, and the current size of the librarian_queue.

    Returns:
        tuple: A tuple containing a JSON response and a status code.
    """
    record = json.loads(request.data)
    app.logger.info(record)
    app.logger.info(record["query"])
    app.logger.info(record["UUID"])
    uuid = record["UUID"]
    cl = Librarian(app, record["query"], uuid)
    librarian_queue.put(cl)
    librarian_list.append(cl)
    answer_data = (record["query"], record["UUID"], librarian_queue.qsize())
    return_data = (
        jsonify(isError=False, message="Success", statusCode=200, data=answer_data),
        200,
    )
    return return_data


@app.route("/get_partial_result", methods=["POST"])
async def get_partial():
    """
    Retrieves the partial result for a given UUID.

    Returns:
        A JSON response containing the partial result.
    """
    record = json.loads(request.data)
    app.logger.info(record)
    app.logger.info(record["UUID"])
    for lib in librarian_list:
        if lib.uuid == record["UUID"]:
            answer_data = lib.live_results
            break
    return_data = (
        jsonify(isError=False, message="Success", statusCode=200, data=answer_data),
        200,
    )
    return return_data


# =======================================MAIN===================================================
if __name__ == "__main__":
    app.logger.setLevel(logging.DEBUG)
    h1 = handlers.RotatingFileHandler(
        filename="E:\\logs\\librarian.log",
        encoding="utf-8",
        mode="a",
        maxBytes=6 * 1024 * 1024,
        backupCount=6,
    )

    app.logger.addHandler(h1)
    threads = []
    threads.append(threading.Thread(target=waitress_run))
    # threads.append(threading.Thread(target=flask_debug))
    bgtask = BackgroundTaskSearch()
    bgtask.app = app
    threads.append(bgtask)
    threads.append(threading.Thread(target=scrape_bot.scraper, args=(app.logger,)))
    i = 0
    for worker in threads:
        try:
            app.logger.info("App number: %s", i)
            i += 1
            worker.start()
        except RuntimeError as e:
            app.logger.error("Exploded")
            print(str(e))
    for worker in threads:
        worker.join()