# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import logging
import os

import torch

open_source: bool = True


def _load_library(filename: str, version: str, no_throw: bool = False) -> None:
    """Load a shared library from the given filename."""
    try:
        library_path = os.path.join(os.path.dirname(__file__), filename)
        torch.ops.load_library(library_path)
        torch.classes.load_library(library_path)
        logging.info(f"Successfully loaded: '{filename}'")

    except Exception as error:
        logging.error(f"Could not load the library '{filename}'!\n\n\n{error}\n\n\n")
        if not no_throw:
            raise error


try:
    # Export the version string from the version file auto-generated by setup.py
    from .version import __target__, __variant__, __version__  # noqa: F401, E402
except Exception:
    __variant__: str = "INTERNAL"
    __version__: str = "INTERNAL"
    __target__: str = "default"

_default_libraries = [
    "mslk",
]

libraries_to_load = {
    "default": _default_libraries,
}

for library in libraries_to_load.get(__target__, []):
    # NOTE: In all cases, we want to throw an error if we cannot load the
    # library.  However, this appears to break the OSS documentation build,
    # where the Python documentation doesn't show up in the generated docs.
    #
    # To work around this problem, we introduce a fake build variant called
    # `docs` and we only throw a library load error when the variant is not
    # `docs`.  For more information, see:
    #
    #   https://github.com/pytorch/FBGEMM/pull/3477
    #   https://github.com/pytorch/FBGEMM/pull/3717
    _load_library(f"{library}.so", __version__, __variant__ == "docs")
