#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import base64
import copy
import functools
import glob
import io
import json
import os
from collections import OrderedDict

from polygraphy import constants, mod, util
from polygraphy.logger import G_LOGGER

np = mod.lazy_import("numpy")
torch = mod.lazy_import("torch>=1.13.0")


def legacy_str_from_type(typ):
    return "__polygraphy_encoded_" + typ.__name__


def str_from_type(typ):
    return typ.__name__


class BaseCustomImpl:
    """
    Base class for Polygraphy's JSON encoder/decoder.
    """

    @classmethod
    def register(cls, typ, alias=None):
        """
        Decorator that registers JSON encoding/decoding functions for types.

        Args:
            typ (type): The type to register
            alias (str):
                    An alias under which to also register the decoder function.
                    This can be used to retain backwards-compatibility when a class
                    name changes.

        For the documentation that follows, assume we have a class:
        ::

            class Dummy:
                def __init__(self, x):
                    self.x = x

        ========
        Encoders
        ========

        Encoder functions should accept instances of the specified type and
        return dictionaries.

        For example:
        ::

            @Encoder.register(Dummy)
            def encode(dummy):
                return {"x": dummy.x}


        To use the custom encoder, use the `to_json` helper:
        ::

            d = Dummy(x=1)
            d_json = to_json(d)


        ========
        Decoders
        ========

        Decoder functions should accept dictionaries, and return instances of the
        type.

        For example:
        ::

            @Decoder.register(Dummy)
            def decode(dct):
                return Dummy(x=dct["x"])


        To use the custom decoder, use the `from_json` helper:
        ::

            from_json(d_json)


        Args:
            typ (type): The type of the class for which to register the function.
        """

        def register_impl(func):
            def add(key, val):
                if key in cls.polygraphy_registered:
                    G_LOGGER.critical(
                        f"Duplicate serialization function for type: {key}.\nNote: Existing function: {cls.polygraphy_registered[key]}, New function: {func}"
                    )
                cls.polygraphy_registered[key] = val

            if cls == Encoder:

                def wrapped(obj):
                    dct = func(obj)
                    dct[constants.TYPE_MARKER] = str_from_type(typ)
                    return dct

                add(typ, wrapped)
                return wrapped
            elif cls == Decoder:

                def wrapped(dct):
                    if constants.TYPE_MARKER in dct:
                        del dct[constants.TYPE_MARKER]

                    type_name = legacy_str_from_type(typ)
                    if type_name in dct:
                        del dct[type_name]

                    return func(dct)

                add(legacy_str_from_type(typ), wrapped)
                add(str_from_type(typ), wrapped)
                if alias is not None:
                    add(alias, wrapped)
            else:
                G_LOGGER.critical("Cannot register for unrecognized class type: ")

        return register_impl


@mod.export()
class Encoder(BaseCustomImpl, json.JSONEncoder):
    """
    Polygraphy's custom JSON Encoder implementation.
    """

    polygraphy_registered = {}

    def default(self, o):
        if type(o) in self.polygraphy_registered:
            return self.polygraphy_registered[type(o)](o)
        return super().default(o)


@mod.export()
class Decoder(BaseCustomImpl):
    """
    Polygraphy's custom JSON Decoder implementation.
    """

    polygraphy_registered = {}

    def __call__(self, pairs):
        # The encoder will insert special key-value pairs into dictionaries encoded from
        # custom types. If we find one, then we know to decode using the corresponding custom
        # type function.
        dct = OrderedDict(pairs)

        # Handle legacy naming first - these keys should not be present in JSON generated by more recent versions of Polygraphy.
        for type_str, func in self.polygraphy_registered.items():
            if (
                type_str in dct and dct[type_str] == constants.LEGACY_TYPE_MARKER
            ):  # Found a custom type!
                return func(dct)

        type_name = dct.get(constants.TYPE_MARKER)
        if type_name is not None:
            if type_name not in self.polygraphy_registered:
                user_type_name = {
                    "Tensor": "torch.Tensor",
                    "ndarray": "np.ndarray",
                }.get(type_name, type_name)
                G_LOGGER.critical(
                    f"Could not decode serialized type: {user_type_name}. This could be because a required module is missing. "
                )
            return self.polygraphy_registered[type_name](dct)

        return dct


NUMPY_REGISTRATION_SUCCESS = False
TORCH_REGISTRATION_SUCCESS = False
COMMON_REGISTRATION_SUCCESS = False


def try_register_common_json(func):
    """
    Decorator that attempts to register common JSON encode/decode methods
    if the methods have not already been registered.

    This needs to be attempted multiple times because dependencies may become available in the
    middle of execution - for example, if using dependency auto-installation.
    """

    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        global NUMPY_REGISTRATION_SUCCESS
        if not NUMPY_REGISTRATION_SUCCESS and np.is_installed() and np.is_importable():
            # We define this alongside load_json/save_json so that it is guaranteed to be
            # imported before we need to encode/decode NumPy arrays.
            @Encoder.register(np.ndarray)
            def encode(array):
                outfile = io.BytesIO()
                np.save(outfile, array, allow_pickle=False)
                outfile.seek(0)
                data = base64.b64encode(outfile.read()).decode()
                return {"array": data}

            @Decoder.register(np.ndarray)
            def decode(dct):
                def load(mode="base64"):
                    if mode == "base64":
                        data = base64.b64decode(dct["array"].encode(), validate=True)
                    elif mode == "latin-1":
                        data = dct["array"].encode(mode)
                    else:
                        assert False, f"Unsupported mode: {mode}"
                    infile = io.BytesIO(data)
                    return np.load(infile, allow_pickle=False)

                try:
                    arr = load()
                except:
                    arr = load("latin-1")  # For backwards compatibility
                if isinstance(arr, np.ndarray):
                    return arr
                # For backwards compatibility with NPZ archives
                values = list(arr.values())
                if not values:
                    G_LOGGER.critical(
                        "Array payload is empty or invalid (NPZ archive contains no arrays)."
                    )
                return values[0]

            NUMPY_REGISTRATION_SUCCESS = True

        global TORCH_REGISTRATION_SUCCESS
        if (
            not TORCH_REGISTRATION_SUCCESS
            and torch.is_installed()
            and torch.is_importable()
        ):

            @Encoder.register(torch.Tensor)
            def encode(tensor):
                outfile = io.BytesIO()
                torch.save(tensor, outfile)
                outfile.seek(0)
                data = base64.b64encode(outfile.read()).decode()
                return {"tensor": data}

            @Decoder.register(torch.Tensor)
            def decode(dct):
                data = base64.b64decode(dct["tensor"].encode(), validate=True)
                infile = io.BytesIO(data)
                return torch.load(infile, weights_only=True)

            TORCH_REGISTRATION_SUCCESS = True

        global COMMON_REGISTRATION_SUCCESS
        if not COMMON_REGISTRATION_SUCCESS:
            # Pull in some common types so that we can get their associated serialization/deserialization
            # functions. This allows the user to avoid importing these manually.
            # Note: We can only do this here for submodules with no external dependencies.
            # That means, for example, nothing from `backend/` can be imported here.
            from polygraphy.common import FormattedArray
            from polygraphy.comparator import RunResults

            COMMON_REGISTRATION_SUCCESS = True
        return func(*args, **kwargs)

    return wrapped


@mod.export()
@try_register_common_json
def to_json(obj):
    """
    Encode an object to JSON.

    NOTE: For Polygraphy objects, you should use the ``to_json()`` method instead.

    Returns:
        str: A JSON representation of the object.
    """
    return json.dumps(obj, cls=Encoder, indent=constants.TAB)


@mod.export()
@try_register_common_json
def from_json(src):
    """
    Decode a JSON string to an object.

    NOTE: For Polygraphy objects, you should use the ``from_json()`` method instead.

    Args:
        src (str):
                The JSON representation of the object

    Returns:
        object: The decoded instance
    """
    return json.loads(src, object_pairs_hook=Decoder())


@mod.export()
@try_register_common_json
def save_json(obj, dest, description=None):
    """
    Encode an object as JSON and save it to a file.

    NOTE: For Polygraphy objects, you should use the ``save()`` method instead.

    Args:
        obj : The object to save.
        src (Union[str, file-like]): The path or file-like object to save to.
    """
    util.save_file(to_json(obj), dest, mode="w", description=description)


@mod.export()
@try_register_common_json
def load_json(src, description=None):
    """
    Loads a file and decodes the JSON contents.

    NOTE: For Polygraphy objects, you should use the ``load()`` method instead.

    Args:
        src (Union[str, file-like]): The path or file-like object to load from.

    Returns:
        object: The object, or `None` if nothing could be read.
    """
    return from_json(util.load_file(src, mode="r", description=description))


# Width of the zero-padded index in per-iteration file names (e.g. ``0000000000.json``). Fixed so
# that lexical ordering of the file names matches numeric index ordering; wide enough that no real
# dataset reaches it.
ITERATION_FILE_INDEX_WIDTH = 10


def list_iteration_files(dir_path):
    """
    Lists the ``*.json`` files in a directory in lexical order. Per-iteration files are written with
    zero-padded indices (see ``ITERATION_FILE_INDEX_WIDTH``), so lexical order matches index order.
    """
    return sorted(glob.glob(os.path.join(glob.escape(dir_path), "*.json")))


def is_single_file(path):
    """
    Whether a path refers to a single file (has a file extension) rather than a per-iteration
    directory. Depends only on the path string, not on what exists on disk.
    """
    return path is not None and os.path.splitext(path)[1] != ""


class IterationWriter:
    """
    Writes a sequence of per-iteration objects, the writing inverse of ``iterate_from_files``.
    A path with a file extension is a single file; an extensionless path is a directory of
    zero-padded ``<index>.json`` files (e.g. ``0000000000.json``, so lexical order matches index
    order), which must be empty or not yet exist (existing files are never overwritten).

    Append each iteration's object with ``append`` and finalize with ``flush`` (or use as a
    context manager). Directory mode writes each object immediately, keeping memory constant;
    single-file mode accumulates the objects and combines them on ``flush`` (so prefer a directory
    for large datasets).
    """

    def __init__(self, path, combine=None, copy_on_accumulate=False, description=None):
        """
        Args:
            path (str): The destination file or directory. If None, writing is disabled.
            combine (Callable[[list], object]):
                    Combines the accumulated objects into the single object written in single-file
                    mode. Defaults to the list of objects.
            copy_on_accumulate (bool):
                    Whether to deep-copy objects retained for single-file output. Needed when the
                    caller reuses buffers between iterations (as in streaming inference).
                    Defaults to False.
            description (str): A description of what is being saved, used for logging.
        """
        self.path = path
        self.single_file = is_single_file(path)
        self._combine = combine if combine is not None else (lambda items: items)
        self._copy = copy_on_accumulate
        self._description = description
        self._accumulated = []
        self._index = 0
        if path is not None:
            desc = self._description or "data"
            if self.single_file:
                G_LOGGER.info(f"Saving {desc} to a single file: {path}")
            else:
                G_LOGGER.info(
                    f"Saving {desc} to a directory of per-iteration files: {path}"
                )
                # Per-iteration files are written by index, so never overwrite existing content.
                if os.path.exists(path):
                    if not os.path.isdir(path):
                        G_LOGGER.critical(
                            f"Cannot save {desc} to '{path}': a file already exists there, but an "
                            f"extensionless path is a directory of per-iteration files. Use a path "
                            f"with a file extension to save a single file, or a different directory."
                        )
                    if os.listdir(path):
                        G_LOGGER.critical(
                            f"Cannot save {desc} to '{path}': the directory is not empty. "
                            f"Per-iteration files require an empty or new directory (so nothing is "
                            f"overwritten). Use a different path, or remove its contents first."
                        )
                os.makedirs(path, exist_ok=True)

    def append(self, item):
        if self.path is None:
            return
        if self.single_file:
            self._accumulated.append(copy.deepcopy(item) if self._copy else item)
        else:
            save_json(
                item,
                os.path.join(
                    self.path, f"{self._index:0{ITERATION_FILE_INDEX_WIDTH}d}.json"
                ),
                description=self._description,
            )
            self._index += 1

    def flush(self):
        # single_file is only True when path has an extension, so path is not None.
        if self.single_file:
            save_json(
                self._combine(self._accumulated),
                self.path,
                description=self._description,
            )

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.flush()


def iterate_from_files(paths, load_file, allow_dirs=True):
    """
    Lazily yields per-iteration objects from one or more file/directory paths. The shared traversal
    behind ``StreamingDataLoader`` and ``RunResults.load_streaming``. Memory stays constant when
    each iteration is its own file (a directory of per-iteration files); a single file is loaded in
    full by ``load_file`` before its iterations are yielded.

    Args:
        paths (Union[str, Sequence[str]]): A path, or sequence of paths. A directory's iterations
                are its ``*.json`` files (in index order); a file is split by ``load_file``.
        load_file (Callable[[str], Iterator]): A generator yielding the iterations in a single file.
        allow_dirs (bool): Whether a directory path is permitted (see ``StreamingDataLoader``).
    """
    # A single str/os.PathLike is treated as one path, not iterated character-by-character.
    if isinstance(paths, (str, os.PathLike)):
        paths = [paths]
    for path in paths:
        if os.path.isdir(path):
            if not allow_dirs:
                G_LOGGER.critical(
                    f"Could not load '{path}': it is a directory, but a file was expected. To read "
                    f"a directory of per-iteration files, use the streaming reader (e.g. "
                    f"--load-inputs/--load-outputs, or StreamingDataLoader/RunResults.load_streaming)."
                )
            for file_path in list_iteration_files(path):
                yield from load_file(file_path)
        else:
            yield from load_file(path)


def iterate_json_list(path, description=None):
    """
    Loads a single JSON file and yields its iterations: each element if the file holds a JSON list,
    otherwise the single decoded object. Used by ``StreamingDataLoader``, where a file may hold
    either one feed_dict or a list of feed_dicts.
    """
    obj = load_json(path, description=description)
    if isinstance(obj, list):
        yield from obj
    else:
        yield obj


@mod.export()
def add_json_methods(description=None):
    """
    Decorator that adds 4 JSON helper methods to a class:

    - to_json(): Convert to JSON string
    - from_json(): Convert from JSON string
    - save(): Convert to JSON and save to file
    - load(): Load from file and convert from JSON

    Args:
        description (str):
                A description of what is being saved or loaded.
    """

    def add_json_methods_impl(cls):
        # JSON methods

        def check_decoded(obj):
            if not isinstance(obj, cls):
                G_LOGGER.critical(
                    f"Provided JSON cannot be decoded into a {cls.__name__}.\nNote: JSON was decoded into a {type(obj)}:\n{obj}"
                )
            return obj

        def _to_json_method(self):
            """
            Encode this instance as a JSON object.

            Returns:
                str: A JSON representation of this instance.
            """
            return to_json(self)

        def _from_json_method(src):
            return check_decoded(from_json(src))

        _from_json_method.__doc__ = f"""
            Decode a JSON object and create an instance of this class.

            Args:
                src (str):
                        The JSON representation of the object

            Returns:
                {cls.__name__}: The decoded instance

            Raises:
                PolygraphyException:
                        If the JSON cannot be decoded to an instance of {cls.__name__}
            """

        cls.to_json = _to_json_method
        cls.from_json = staticmethod(_from_json_method)

        # Save/Load methods

        def _save_method(self, dest):
            """
            Encode this instance as a JSON object and save it to the specified path
            or file-like object.

            Args:
                dest (Union[str, file-like]):
                      The path or file-like object to write to.

            """
            save_json(self, dest, description=description)

        def _load_method(src):
            return check_decoded(load_json(src, description=description))

        _load_method.__doc__ = f"""
            Loads an instance of this class from a JSON file.

            Args:
                src (Union[str, file-like]): The path or file-like object to read from.

            Returns:
                {cls.__name__}: The decoded instance

            Raises:
                PolygraphyException:
                        If the JSON cannot be decoded to an instance of {cls.__name__}
            """

        cls.save = _save_method
        cls.load = staticmethod(_load_method)

        return cls

    return add_json_methods_impl