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 os
import threading
from json.decoder import JSONDecodeError
from logging import handlers
from pathlib import Path
from queue import Queue
from typing import Dict, Optional

import requests
import lib_paths
import scrape_bot
import search_bot
# import search_bot2 as search_bot
from flask import Flask, jsonify, request, abort
from habanero import Crossref
from waitress import serve

try:
    import netrc
except ImportError:  # pragma: no cover
    netrc = None

# Constants


def _env(name: str, default: str) -> str:
    return os.getenv(name, default)


def _env_path(name: str, default: str) -> Path:
    return Path(os.getenv(name, default)).expanduser().resolve()


BASE_DIR = Path(
    os.getenv("CONJURER_LIBRARIAN_BASE", str(Path(__file__).resolve().parent))
)
NETRC_FILE = _env_path("CONJURER_NETRC_FILE", str(Path.home() / ".netrc"))
HOST_ADDRESS = _env("CONJURER_LIBRARIAN_HOST", "0.0.0.0")
PORT_ADDRESS = int(_env("CONJURER_LIBRARIAN_PORT", "5001"))
MAIN_BOT_ADDRESS = _env("CONJURER_MAIN_BOT", "http://127.0.0.1:5000")
SEND_RESULTS = _env("CONJURER_LIBRARIAN_RESULTS_ENDPOINT", "/conjurer")
MAX_CR_RESULTS = int(_env("CONJURER_LIBRARIAN_MAX_RESULTS", "500"))
ENCODING = _env("CONJURER_ENCODING", "utf-8")
API_KEY = os.getenv("CONJURER_API_KEY")
LOGFILE_PATH = _env_path(
    "CONJURER_LIBRARIAN_LOG", str(BASE_DIR / "librarian.log")
)

app = Flask(__name__)

librarian_queue = Queue()
librarian_list = []


def _service_headers() -> Dict[str, str]:
    if API_KEY:
        return {"X-Conjurer-Api-Key": API_KEY}
    return {}


def _authorize_request() -> None:
    if API_KEY and request.headers.get("X-Conjurer-Api-Key") != API_KEY:
        abort(401)


# trunk-ignore(pylint/R0902)
class Librarian(object):
    """
    Represents a librarian object that performs search and refinement operations on queries.
    """

    def __init__(self, _app, query, uuid, _deep_search) -> 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.
        """
        mailto_contact: Optional[str] = os.getenv("CONJURER_CROSSREF_MAILTO")
        if netrc:
            try:
                netrc_mod = netrc.netrc(str(NETRC_FILE))
                auth_tokens = netrc_mod.authenticators("crossref")
                if auth_tokens:
                    mailto_contact = auth_tokens[0]
            except (FileNotFoundError, netrc.NetrcParseError):
                logging.getLogger("conjurer_librarian").warning(
                    "Crossref credentials missing in netrc %s", NETRC_FILE
                )
        if not mailto_contact:
            raise RuntimeError(
                "Crossref credentials not configured. Set CONJURER_CROSSREF_MAILTO or add to netrc."
            )
        self.cr = Crossref(
            mailto=mailto_contact,
            ua_string=f"Conjurer project. mailto:{mailto_contact}"
        )
        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
        self.deep_search = _deep_search

    async def search_crossref(self, query, deep_search=False):
        """
        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")
        
        if not deep_search:
            query_limit = MAX_CR_RESULTS if MAX_CR_RESULTS < 1000 else 1000
            cr_result = self.cr.works(query=query, limit=query_limit)
            self.search_result_from_cr.update(cr_result)
            self.total = cr_result["message"]["total-results"]
            self.fetched += len(cr_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:
                tmp_result = self.cr.works(query=query, limit=query_limit, offset=self.fetched)
                cr_result["message"]["items"].extend(tmp_result["message"]["items"])
                self.total = tmp_result["message"]["total-results"]
                self.fetched = len(cr_result["message"]["items"])
                self.app.logger.info(self.total)
                self.app.logger.info(self.fetched)
                await asyncio.sleep(0.1)

        else:
            cr_result = self.cr.works(query=query, cursor_max=15000, cursor='*', progress_bar = True)
            result = cr_result[0]
            for item in cr_result[1:]:
                result["message"]["items"].extend(item["message"]["items"])
                self.total = item["message"]["total-results"]
                self.fetched = len(result["message"]["items"])
                self.app.logger.info(self.total)
                self.app.logger.info(self.fetched)
            cr_result = result
            self.app.logger.info("Total, fetched:")
            self.app.logger.info(self.total)
            self.app.logger.info(self.fetched)
            self.search_result_from_cr.update(cr_result)
            self.total = cr_result["message"]["total-results"]
            self.app.logger.info("CROSSREF DONE")

            self.app.logger.info("CROSSREF DONE")
        with open(lib_paths.CR_RESULTS, "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 = {}
            data_file.truncate(0)
            data_file.seek(0)
            tmp = {self.uuid : self.search_result_from_cr}
            if file_data:
                file_data.update(tmp)
            else:
                file_data = tmp
            json.dump(file_data, data_file, indent=4)
        return cr_result


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

        Args:
        - unrefined_result: The unrefined search result.

        Returns:
        - refined_result: The refined search result.

        Raises:
        - None.
        """
        summarized_results = []
        self.app.logger.info("REFINE: Removing all derived works from the list")
        for item in unrefined_result["message"]["items"]:
            summarized_results.append(
                {
                    "DOI": item["DOI"],
                    "title": item["title"] if "title" in item else None,
                    "type": item["type"] if "type" in item else None,
                }
            )

        partial_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"],
            }
        }
        self.app.logger.info("REFINE: Dumping to file")
        temp = []
        refined_result = {}
        for key in partial_result:
            self.app.logger.info("KEY:")
            self.app.logger.info(key)

        for item in partial_result[self.uuid]["summary"]:
            if item["title"]:
                temp.append(item)

        for item in temp:
            refined_result[item["DOI"]]=  item
        with open(lib_paths.RR_RESULTS, "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 = {}
            data_file.truncate(0)
            data_file.seek(0)
            tmp = {self.uuid: refined_result}
            if file_data:
                file_data.update(tmp)
            else:
                file_data = tmp
            json.dump(file_data, data_file, indent=4)
        return refined_result

    async def check_if_exists(self, refined_result):
        """
        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 = {}
        self.app.logger.info("REFINE: Running search in the backend app")
        dois = []
        for item, value in refined_result.items():
            dois.append([item, value])
        coro = asyncio.to_thread(
            search_bot.search_for_doi, dois, self.live_results, self.app.logger
        )
        result = await coro
        result_list = []
        result_no_db = []
        for item in result:
            if item["exists"]:
                result_list.append(item)
            else:
                result_no_db.append(item)
        self.hit = len(result)
        return result_list, result_no_db


    async def answer_query(self, deep_search=False):
        """
        Answers the search query.

        Args:
        - None.

        Returns:
        - result: The search result.

        Raises:
        - None.
        """
        self.app.logger.info(f"Search started {self.uuid}")
        cr_result = await self.search_crossref(query=self.query, deep_search=deep_search)
        refined_result = await self.refine_search(cr_result)
        answer, negative_answer = await self.check_if_exists(refined_result)

        self.app.logger.info("Returning result")
        self.app.logger.info(answer)
        self.app.logger.info(negative_answer)

        for item in answer:
            self.final_result[item["DOI"]] = {"Title": item["data"]["title"], "type": item["data"]["type"]}
        for item in negative_answer:
            self.not_in_db[item["DOI"]] = {"Title": item["data"]["title"], "type": item["data"]["type"]}
        self.app.logger.info("Returning result case2")
        self.app.logger.info(self.final_result)
        return self.final_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(librarian.deep_search)
            result = {librarian.uuid: result}
            self.app.logger.info("Saving to file")

            # Save results to "not_in_db.json" file
            with open(lib_paths.NOT_IN_DB, "r+", encoding="utf-8") as ndb_file:
                ndb_database = {}
                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(lib_paths.S_RESULTS, "r+", encoding="utf-8") as s_file:
                database = {}
                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")

            # Send the result back to the bot. Log EXACTLY what goes out (target,
            # uuid, how many DOIs and which) so the librarian log makes it plain a
            # result was sent and what was in it.
            payload = result  # shape: {uuid: {DOI: {"Title": ..., "type": ...}}}
            hits = payload.get(librarian.uuid, {}) if isinstance(payload, dict) else {}
            target = f"{MAIN_BOT_ADDRESS}{SEND_RESULTS}"
            self.app.logger.info(
                "SENDING result for %s to %s: %d DOI(s): %s",
                librarian.uuid,
                target,
                len(hits),
                list(hits.keys()),
            )
            # A failed send must NOT kill this worker - otherwise a bot that is
            # momentarily down stalls every future query until the librarian is
            # restarted. Log and carry on to the next queued search.
            try:
                response = await asyncio.to_thread(
                    requests.post,
                    target,
                    json=payload,
                    headers=_service_headers(),
                    timeout=360,
                )
                if response.status_code == 200:
                    self.app.logger.info(
                        "SENT result for %s -> HTTP 200 (bot accepted)", librarian.uuid
                    )
                else:
                    self.app.logger.warning(
                        "SENT result for %s but bot returned HTTP %s: %s",
                        librarian.uuid,
                        response.status_code,
                        response.text[:500],
                    )
            except requests.exceptions.RequestException as exc:
                self.app.logger.error(
                    "FAILED to send result for %s to %s: %s",
                    librarian.uuid,
                    target,
                    exc,
                )
            await asyncio.sleep(1)


# ==================================SERVER ROUTES==========================================
@app.route("/query", methods=["POST"])
async def query_database():
    _authorize_request()
    """
    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"]
    deep_search = record["deep_search"]
    cl = Librarian(app, record["query"], uuid, deep_search)
    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():
    _authorize_request()
    """
    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)
    LOGFILE_PATH.parent.mkdir(parents=True, exist_ok=True)
    h1 = handlers.RotatingFileHandler(
        filename=str(LOGFILE_PATH),
        encoding=ENCODING,
        mode="a",
        maxBytes=6 * 1024 * 1024,
        backupCount=6,
    )

    app.logger.addHandler(h1)
    threads = []
    threads.append(threading.Thread(target=waitress_run, daemon=True))
    # threads.append(threading.Thread(target=flask_debug))
    bgtask = BackgroundTaskSearch()
    bgtask.app = app
    bgtask.daemon = True
    threads.append(bgtask)
    threads.append(
        threading.Thread(
            target=scrape_bot.scraper, args=(app.logger,), daemon=True
        )
    )
    i = 0
    try:
        for worker in threads:
            app.logger.info("App number: %s", i)
            i += 1
            worker.start()
        for worker in threads:
            worker.join()
    except KeyboardInterrupt:
        app.logger.info("Shutdown requested - exiting librarian service")