"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Monday, 2nd Sept., 2024

    OBJECTIVE:

        To provide an easy way to assess images for adult content.
        Tried and implemented using HuggingFace behaviour_models.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# System-level activities:
import io

# For using the AI model:
import torch
from transformers import pipeline

# To download images from URLs:
import requests

# To read images:
from PIL import Image

# Common:
from shared import variables


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  CLASSES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class AssessImageNSFW:

    def __init__(
            self,
            model,
            nsfw_label = "nsfw",
            nsfw_threshold = 0.25,
    ):

        """
        This class assesses the input image and tells if it is blurry or clear.
        :param model: The model to use. Either the name of the HuggingFace repo, or the directory where the model is
            stored.
        :param nsfw_threshold: The max allowed blurriness (0 to 1 range).
        """

        # Note down the config:
        self.__model = model
        self.__device = "gpu" if torch.cuda.is_available() else "cpu"
        self.__nsfw_label = nsfw_label
        self.__nsfw_threshold = nsfw_threshold

        # Initialize the classifier:
        self.__classifier = pipeline(
            task = "image-classification",
            model = self.__model,
            device = self.__device
        )

    @staticmethod
    async def read_image(source):

        """
        Reads an image in whichever format it is provided and returns it as a PIL object.
        :param source: The image as either a path or a URL or a io.BytesIO object.
        :return: The image opened as a PIL object.
        """

        # If the input image is already a PIL image:
        if isinstance(source, Image.Image): return source

        # If a buffer is provided:
        elif isinstance(source, io.BytesIO):
            source.seek(0)
            return Image.open(source)

        # If a string is provided (local path or URL):
        elif isinstance(source, str):
            if source.startswith("http://") or source.startswith("https://"):
                response = await variables.http_client.get(source)
                return Image.open(io.BytesIO(response.content))
            else: return Image.open(source)

    async def classify(self, image):

        """
        To get the prediction of the model from the given input image.
        :param image: The image (as a PIL object or file path or a URL).
        :return: The dictionary of classes with their respective predictions.
        """

        image = await self.read_image(image)
        predictions = self.__classifier(image)
        classes = {p["label"]: p["score"] for p in predictions}
        return classes

    async def is_ok(
            self,
            image,
            nsfw_threshold = None
    ):

        """
        Checks if the image given to it can be used, or should be rejected.
        :param image: The image (as a PIL object or file path or a URL).
        :param nsfw_threshold: A custom threshold to test against. If not provided, the default will be taken that was
            provided when the instance was created.
        :return: True if the image is okay, else False.
        """

        image = await self.read_image(image)
        nsfw_threshold = nsfw_threshold or self.__nsfw_threshold
        classes = await self.classify(image)
        return True if classes[self.__nsfw_label] <= nsfw_threshold else False


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    import asyncio


    async def main():

        assessor = AssessImageNSFW(model = r"/path/to/model")
        usable = await assessor.is_ok(r"/path/to/image/img.jpg")
        classes = await assessor.classify(r"https://...")


    asyncio.run(main())