path: root/src/lib/szurubooru.py

"""
Szurubooru Library
~~~~~~~~~~~~~~~~~~

Szurubooru Library is used for communicating with a Szurubooru instance through python.

:copyright (c) 2023 Colin Wilk.
:license: MIT, see LICENSE for more details.
"""


from lib.dataclasses import Tag
import requests as r
from typing import List
import json
from urllib.parse import quote
import math
import functools


class SzurubooruException(Exception):
    """General Error returned from the Szurubooru REST API.

    The Szurubooru API returns errors together with an http code in a form like this:
    .. code-block:: JSON
        {
            "name": "Name of the error, e.g. 'PostNotFoundError'",
            "title": "Generic title of error message, e.g. 'Not found'",
            "description": "Detailed description of what went wrong, e.g.
                            'User `rr-` not found."
        }

    We map this json response to the SzurubooruException class or more specific children
    such as MissingRequiredFileError with the fields / args that are provided from the
    json response.

    Args:
        name: Name of the exception as defined by the Szurubooru REST API
              The name is also used in more specific Exceptions such as
              MissingRequiredFileError(SzurubooruException).
        title: Generic, human readable title of the error message.
        description: More detailed description of what went wrong. Often includes
                     input from the user that caused the error.
    """

    name: str
    title: str
    description: str

    def __init__(self, name: str, title: str, description: str) -> None:
        self.name = name
        self.title = title
        self.description = description
        super().__init__(description)

    @staticmethod
    def response_is_exception(res: r.Response) -> bool:
        """Checks if the Szurubooru REST API response is an error or not.

        Args:
            res (r.Response): the Response of the request from the requests module.

        Returns:
            bool: True if response is an error and False if it is not.
        """

        return res.status_code != 200

    @staticmethod
    def map_exception(res: r.Response) -> Exception:
        """Generates an exception for an error response from the Szurubooru REST API.

        Args:
            res (r.Response): the Response from the Szurubooru REST API that should be
                              mapped to a Python Exception.

        Raises:
            LookupError: Thrown when the Szurubooru REST API sent an error type (name)
                         that could not be mapped to one of the SzurubooruException
                         Python Exceptions. This can happen when there was an unexpected
                         e.g. an Internal Server Error in the Szurubooru REST API.

        Returns:
            Exception: Returns a more specific exception.
        """

        content: dict[str, str] = json.loads(res.content)
        name: str = content["name"]
        title: str = content["title"]
        description: str = content["description"]

        match name:
            case "MissingRequiredFileError":
                return MissingRequiredFileError(name, title, description)
            case "MissingRequiredParameterError":
                return MissingRequiredParameterError(name, title, description)
            case "InvalidParameterError":
                return InvalidParameterError(name, title, description)
            case "IntegrityError":
                return IntegrityError(name, title, description)
            case "SearchError":
                return SearchError(name, title, description)
            case "AuthError":
                return AuthError(name, title, description)
            case "PostNotFoundError":
                return PostNotFoundError(name, title, description)
            case "PostAlreadyFeaturedError":
                return PostAlreadyFeaturedError(name, title, description)
            case "PostAlreadyUploadedError":
                return PostAlreadyUploadedError(name, title, description)
            case "InvalidPostIdError":
                return InvalidPostIdError(name, title, description)
            case "InvalidPostSafetyError":
                return InvalidPostSafetyError(name, title, description)
            case "InvalidPostSourceError":
                return InvalidPostSourceError(name, title, description)
            case "InvalidPostContentError":
                return InvalidPostContentError(name, title, description)
            case "InvalidPostRelationError":
                return InvalidPostRelationError(name, title, description)
            case "InvalidPostNoteError":
                return InvalidPostNoteError(name, title, description)
            case "InvalidPostFlagError":
                return InvalidPostFlagError(name, title, description)
            case "InvalidFavoriteTargetError":
                return InvalidFavoriteTargetError(name, title, description)
            case "InvalidCommentIdError":
                return InvalidCommentIdError(name, title, description)
            case "CommentNotFoundError":
                return CommentNotFoundError(name, title, description)
            case "EmptyCommentTextError":
                return EmptyCommentTextError(name, title, description)
            case "InvalidScoreTargetError":
                return InvalidScoreTargetError(name, title, description)
            case "InvalidScoreValueError":
                return InvalidScoreValueError(name, title, description)
            case "TagCategoryNotFoundError":
                return TagCategoryNotFoundError(name, title, description)
            case "TagCategoryAlreadyExistsError":
                return TagCategoryAlreadyExistsError(name, title, description)
            case "TagCategoryIsInUseError":
                return TagCategoryIsInUseError(name, title, description)
            case "InvalidTagCategoryNameError":
                return InvalidTagCategoryNameError(name, title, description)
            case "InvalidTagCategoryColorError":
                return InvalidTagCategoryColorError(name, title, description)
            case "TagNotFoundError":
                return TagNotFoundError(name, title, description)
            case "TagAlreadyExistsError":
                return TagAlreadyExistsError(name, title, description)
            case "TagIsInUseError":
                return TagIsInUseError(name, title, description)
            case "InvalidTagNameError":
                return InvalidTagNameError(name, title, description)
            case "InvalidTagRelationError":
                return InvalidTagRelationError(name, title, description)
            case "InvalidTagCategoryError":
                return InvalidTagCategoryError(name, title, description)
            case "InvalidTagDescriptionError":
                return InvalidTagDescriptionError(name, title, description)
            case "UserNotFoundError":
                return UserNotFoundError(name, title, description)
            case "UserAlreadyExistsError":
                return UserAlreadyExistsError(name, title, description)
            case "InvalidUserNameError":
                return InvalidUserNameError(name, title, description)
            case "InvalidEmailError":
                return InvalidEmailError(name, title, description)
            case "InvalidPasswordError":
                return InvalidPasswordError(name, title, description)
            case "InvalidRankError":
                return InvalidRankError(name, title, description)
            case "InvalidAvatarError":
                return InvalidAvatarError(name, title, description)
            case "ProcessingError":
                return ProcessingError(name, title, description)
            case "ValidationError":
                return ValidationError(name, title, description)
            case _:
                raise LookupError(f'Unknown SzurubooruException: "{name}"')


class MissingRequiredFileError(SzurubooruException):
    pass


class MissingRequiredParameterError(SzurubooruException):
    pass


class InvalidParameterError(SzurubooruException):
    pass


class IntegrityError(SzurubooruException):
    pass


class SearchError(SzurubooruException):
    pass


class AuthError(SzurubooruException):
    pass


class PostNotFoundError(SzurubooruException):
    pass


class PostAlreadyFeaturedError(SzurubooruException):
    pass


class PostAlreadyUploadedError(SzurubooruException):
    pass


class InvalidPostIdError(SzurubooruException):
    pass


class InvalidPostSafetyError(SzurubooruException):
    pass


class InvalidPostSourceError(SzurubooruException):
    pass


class InvalidPostContentError(SzurubooruException):
    pass


class InvalidPostRelationError(SzurubooruException):
    pass


class InvalidPostNoteError(SzurubooruException):
    pass


class InvalidPostFlagError(SzurubooruException):
    pass


class InvalidFavoriteTargetError(SzurubooruException):
    pass


class InvalidCommentIdError(SzurubooruException):
    pass


class CommentNotFoundError(SzurubooruException):
    pass


class EmptyCommentTextError(SzurubooruException):
    pass


class InvalidScoreTargetError(SzurubooruException):
    pass


class InvalidScoreValueError(SzurubooruException):
    pass


class TagCategoryNotFoundError(SzurubooruException):
    pass


class TagCategoryAlreadyExistsError(SzurubooruException):
    pass


class TagCategoryIsInUseError(SzurubooruException):
    pass


class InvalidTagCategoryNameError(SzurubooruException):
    pass


class InvalidTagCategoryColorError(SzurubooruException):
    pass


class TagNotFoundError(SzurubooruException):
    pass


class TagAlreadyExistsError(SzurubooruException):
    pass


class TagIsInUseError(SzurubooruException):
    pass


class InvalidTagNameError(SzurubooruException):
    pass


class InvalidTagRelationError(SzurubooruException):
    pass


class InvalidTagCategoryError(SzurubooruException):
    pass


class InvalidTagDescriptionError(SzurubooruException):
    pass


class UserNotFoundError(SzurubooruException):
    pass


class UserAlreadyExistsError(SzurubooruException):
    pass


class InvalidUserNameError(SzurubooruException):
    pass


class InvalidEmailError(SzurubooruException):
    pass


class InvalidPasswordError(SzurubooruException):
    pass


class InvalidRankError(SzurubooruException):
    pass


class InvalidAvatarError(SzurubooruException):
    pass


class ProcessingError(SzurubooruException):
    pass


class ValidationError(SzurubooruException):
    pass


class Szurubooru:
    """Client class for communicating with the Szurubooru instance.

    Raises:
        SzurubooruException: When the Szurubooru REST API returns any kinds of errors.
                             See: SzurubooruException for mor information.
    """

    url: str
    token: str
    verify: bool
    header: dict[str, str]

    # Session Stats
    requests_get: int = 0
    requests_post: int = 0
    requests_put: int = 0
    requests_delete: int = 0

    def __init__(self, url: str, token: str, verify: bool = True):
        self.url = url
        self.token = token
        self.verify = verify
        self.header = {
            "Accept": "application/json",
            "Authorization": f"Token {self.token}",
        }

    @staticmethod
    def __raise_szurubooru_exception(func):
        @functools.wraps(func)
        def wrapper(self, **kwargs):
            res = func(self, **kwargs)
            if SzurubooruException.response_is_exception(res):
                raise SzurubooruException.map_exception(res)
            return res

        return wrapper

    @staticmethod
    def __count_stats(get: int = 0, post: int = 0, put: int = 0, delete: int = 0):
        def function_wrapper(func):
            @functools.wraps(func)
            def wrapper(self, **kwargs):
                self.requests_get += get
                self.requests_post += post
                self.requests_put += put
                self.requests_delete += delete
                return func(self, **kwargs)

            return wrapper

        return function_wrapper

    @__count_stats(get=1)
    @__raise_szurubooru_exception
    def __get(self, path: str) -> r.Response:
        return r.get(
            f"{self.url}{path}", timeout=15, verify=self.verify, headers=self.header
        )

    @__count_stats(post=1)
    @__raise_szurubooru_exception
    def __post(self, path: str, params: dict) -> r.Response:
        return r.post(
            f"{self.url}{path}",
            timeout=15,
            verify=self.verify,
            headers=self.header,
            data=json.dumps(params),
        )

    @__count_stats(put=1)
    @__raise_szurubooru_exception
    def __put(self, path: str, params: dict) -> r.Response:
        return r.put(
            f"{self.url}{path}",
            timeout=15,
            verify=self.verify,
            headers=self.header,
            data=json.dumps(params),
        )

    @__count_stats(delete=1)
    @__raise_szurubooru_exception
    def __delete(self, path: str, params: dict) -> r.Response:
        return r.delete(
            f"{self.url}{path}",
            timeout=15,
            verify=self.verify,
            headers=self.header,
            data=json.dumps(params),
        )

    def __map_tag_resource_response(self, res: r.Response, name: str = "") -> Tag:
        """Maps the requests response from the Szurubooru REST API to the Tag dataclass.

        This function is similar to __map_tag_resource and calls it internally.
        Instead of receiving a dictionary if the Szurubooru Tag it operates on the
        requests Response of the same form.

        To Illustrate, both code yield the same result:
            >>> __map_tag_resource_response(res, name)
            >>> __map_tag_resource(json.loads(res.content), name)



        Args:
            res (r.Response): The requests Response from an Szurubooru REST API endpoint
                              that returns a Szurubooru tag resource in a form like:
                              .. code-block:: JSON
                                {
                                    "version":      <version>,
                                    "names":        <names>,
                                    "category":     <category>,
                                    "implications": <implications>,
                                    "suggestions":  <suggestions>,
                                    "creationTime": <creation-time>,
                                    "lastEditTime": <last-edit-time>,
                                    "usages":       <usage-count>,
                                    "description":  <description>
                                }

            name (str, optional): If you provide a name here, the mapper will use that
                                  name for the Tag dataclass instead of fetching it from
                                  the alias names. This is useful if you have multiple
                                  aliases or names in Szurubooru.

        Returns:
            Tag: A Tag dataclass. The version parameter is included if the Szurubooru
                 REST API response contained a version field, otherwise it defaults to
                 -1. The name will be the first name returned by the Szurubooru REST API
                 or the name that was passed to the function as an optional parameter.
        """

        return self.__map_tag_resource(json.loads(res.content), name)

    def __map_tag_resource(self, c: dict, name: str = "") -> Tag:
        """Maps the dict from the Szurubooru REST API to the Tag dataclass.

        This function is similar to __map_tag_resource_response and is called by it.
        Instead of receiving a requests Response it operates on the dictionary of the
        same form.

        To Illustrate, both code yield the same result:
            >>> __map_tag_resource_response(res, name)
            >>> __map_tag_resource(json.loads(res.content), name)

        Args:
            c (dict): The dict object coming from an Szurubooru REST API Szurubooru tag
                      resource in a form like:
                      .. code-block:: JSON
                          {
                              "version":      <version>,
                              "names":        <names>,
                              "category":     <category>,
                              "implications": <implications>,
                              "suggestions":  <suggestions>,
                              "creationTime": <creation-time>,
                              "lastEditTime": <last-edit-time>,
                              "usages":       <usage-count>,
                              "description":  <description>
                          }
            name (str, optional): If you provide a name here, the mapper will use that
                                  name for the Tag dataclass instead of fetching it from
                                  the alias names. This is useful if you have multiple
                                  aliases or names in Szurubooru.

        Returns:
            Tag: A Tag dataclass. The version parameter is included if the Szurubooru
                 REST API response contained a version field, otherwise it defaults to
                 -1. The name will be the first name returned by the Szurubooru REST API
                 or the name that was passed to the function as an optional parameter.
        """

        return Tag(
            name=name if name != "" else c["names"][0],
            version=c.get("version", -1),
            category=c["category"],
            description=c["description"],
            implications=list(map(lambda a: a["names"][0], c["implications"])),
            suggestions=c["suggestions"],
        )

    def tag_exists(self, name: str) -> bool:
        """Checks if the tag exists on Szurubooru.

        Args:
            name (str): Name of the tag.

        Returns:
            bool: True if tag already exists. False if there was no matching tag on
                  Szurubooru.
        """

        try:
            res = self.__get(path=f"/api/tag/{quote(name)}")
        except TagNotFoundError:
            return False
        except Exception as e:
            raise e
        return res.status_code == 200

    def tag_len(self) -> int:
        """Checks how many tags there are on the Szurubooru instance.

        Returns:
            int: The number of tags that exist on the Szurubooru instance
        """

        return json.loads(self.__get(path="/api/tags").content)["total"]

    def tag_list(self, limit: int = 100, offset: int = 0, query: str = "") -> List[Tag]:
        """Lists tags on the Szurubooru instance.

        Args:
            limit (int, optional): Limits the amount of tags displayed on one page.
                                   Setting a smaller limit may help with query speed.
                                   Usually the maximum allowed limit is 100.
                                   Defaults to 100.
            offset (int, optional): Offset where to start the listing. Required for
                                    Paging the data when searching through many tags.
                                    If you have a page size of 100 you need to set an
                                    offset of 100 to see results on the 2. page.
                                    So your offset should be PAGE_NUM * LIMIT.
                                    Defaults to 0.
            query (str, optional): You can use a custom query to filter through tags
                                   better. If you use an empty string (the default) you
                                   will search through all tags on the Szurubooru
                                   instance.
                                   Defaults to ''.

        Returns:
            List[Tag]: List of Tag objects in the current page.
        """

        path = f"/api/tags?offset={offset}&limit={limit}&query={quote(query)}"
        results = json.loads(self.__get(path=path).content)["results"]
        return list(map(self.__map_tag_resource, results))

    def tag_list_all(self, query: str = "") -> List[Tag]:
        """Lists every available Tag (matching the query) on the Szurubooru instance.

        WARNING: This method calls the Szurubooru REST API multiple times depending on
        how many tags there are. It can take a while and can put significant load on the
        Szurubooru instance. Use with caution.

        In essence this method does the paging for you in tag_list and returns the full
        batch of Tags.

        Args:
            query (str, optional): You can use a custom query to filter through tags
                                   better. If you use an empty string (the default) you
                                   will search through all tags on the Szurubooru
                                   instance.
                                   Defaults to ''.

        Returns:
            List[Tag]: List of all Tag objects matching the query on the Szurubooru
                       instance.
        """

        tags = []
        for i in range(0, math.ceil(self.tag_len() / 100)):
            tags.extend(self.tag_list(offset=100 * i, limit=100, query=query))
        return tags

    def tag_get(self, name: str) -> Tag:
        """Get a tag form the Szurubooru instance by name.

        Args:
            name (str): Name of the tag.

        Raises:
            TagNotFoundError: Tag could not be found on the Szurubooru instance.

        Returns:
            Tag: Tag object of the requested tag.
        """

        res = self.__get(path=f"/api/tag/{quote(name)}")
        return self.__map_tag_resource_response(res, name=name)

    def tag_create(self, tag: Tag) -> Tag:
        """Create a tag on the Szurubooru instance.

        Args:
            tag (Tag): Tag dataclass representing the Tag that should be created.

        Returns:
            Tag: Resulting Tag dataclass that was returned by the Szurubooru REST API
                 representing the created tag as it was created on the Szurubooru
                 instance.
        """

        res = self.__post(
            path="/api/tags",
            params={
                "names": [tag.name],
                "category": tag.category,
                "description": tag.description,
                "implications": tag.implications,
                "suggestions": tag.suggestions,
            },
        )
        return self.__map_tag_resource_response(res, name=tag.name)

    def tag_update(self, tag: Tag) -> Tag:
        """Update a tag on the Szurubooru instance.

        Args:
            tag (Tag): Tag dataclass representing the Tag that should be updated. The
                       name is the identifying parameter for the Szurubooru instance.

        Returns:
            Tag: Resulting Tag dataclass that was returned by the Szurubooru REST API
                 representing the updated version of the tag as it is now on the
                 Szurubooru instance.
        """

        res = self.__put(
            path=f"/api/tag/{quote(tag.name)}",
            params={
                "version": tag.version,
                "category": tag.category,
                "description": tag.description,
                "implications": tag.implications,
                "suggestions": tag.suggestions,
            },
        )
        return self.__map_tag_resource_response(res, name=tag.name)

    def tag_delete(self, tag: Tag) -> None:
        """Deletes the given tag on the Szurubooru instance.

        Args:
            tag (Tag): Tag dataclass of the tag that should be reledte. Needs to contain
                       the current version number.
        """
        self.__delete(
            path=f"/api/tag/{quote(tag.name)}", params={"version": tag.version}
        )
        return

    def tag_merge(self, tag_to_merge: Tag, tag_to_merge_into: Tag) -> Tag:
        """Removes tag_to_merge and merges all uses into tag_to_merge_into.

        Removes source tag and merges all of its usages, suggestions and implications to
        the target tag. Other tag properties such as category and aliases do not get
        transferred and are discarded.

        Args:
            tag_to_merge (Tag): Tag dataclass representing the Tag that will be removed
                                and merged into the second Tag. The version parameter is
                                required.
            tag_to_merge_into (Tag): Tag dataclass representing the Tag that the first
                                     tag will be merged into. The version parameter is
                                     required.

        Returns:
            Tag: Resulting Tag dataclass that was returned by the Szurubooru REST API
                 representing the new (merged) tag on the Szurubooru instance.
        """

        res = self.__post(
            path="/api/tag-merge",
            params={
                "removeVersion": tag_to_merge.version,
                "remove": tag_to_merge.name,
                "mergeToVersion": tag_to_merge_into.version,
                "mergeTo": tag_to_merge_into.name,
            },
        )
        return self.__map_tag_resource_response(res, name=tag_to_merge_into.name)

    def tag_list_siblings(self, tag: Tag) -> List[dict[str, int]]:
        """Lists siblings of the given Tag

        Lists siblings of given tag, e.g. tags that were used in the same posts as the
        given tag. `occurrences` field signifies how many times a given sibling appears
        with given tag. Results are sorted by occurrences count and the list is
        truncated to the first 50 elements. Doesn't use paging.

        Args:
            tag (Tag): Tag that will queried. Only uses the tag.name parameter.

        Returns:
            List[dict[str, int]]: Returns a List of dictionaries with the primary tag
                                  name as the key and the number of occurrences as the
                                  value.
        """

        res = self.__get(path=f"/api/tag-siblings/{tag.name}")
        results = json.loads(res.content)["results"]
        return list(map(lambda a: {a["tag"]["names"][0]: a["occurrences"]}, results))