import numpy as np

from asdf.extension import Converter


class NDArrayConverter(Converter):
    tags = [
        "tag:stsci.edu:asdf/core/ndarray-1.0.0",
        "tag:stsci.edu:asdf/core/ndarray-1.1.0",
        "tag:stsci.edu:asdf/core/ndarray-1.2.0",
    ]
    types = [
        np.ndarray,  # we use the type here so the extension can find the sub-classes
        "numpy.ma.MaskedArray",  # in numpy 2.0
        "numpy.ma.core.MaskedArray",
        "asdf.tags.core.ndarray.NDArrayType",
        "asdf.tags.core.stream.Stream",
    ]

    def to_yaml_tree(self, obj, tag, ctx):
        import numpy as np
        from numpy import ma

        from asdf import config, util
        from asdf._block.options import Options
        from asdf.tags.core.ndarray import NDArrayType, numpy_array_to_list, numpy_dtype_to_asdf_datatype
        from asdf.tags.core.stream import Stream

        data = obj

        if isinstance(obj, Stream):
            # previously, stream never passed on data, we can do that here
            ctx._blocks.set_streamed_write_block(data._array, data)

            result = {}
            result["source"] = -1
            result["shape"] = ["*", *data._shape]
            result["datatype"] = data._datatype
            result["byteorder"] = data._byteorder
            if data._strides is not None:
                result["strides"] = data._strides
            return result

        # sort out block writing options
        if isinstance(obj, NDArrayType) and isinstance(obj._source, str):
            # this is an external block, if we have no other settings, keep it as external
            options = ctx._blocks.options.lookup_by_object(data)
            if options is None:
                options = Options("external")
        else:
            options = ctx._blocks.options.get_options(data)

        # The ndarray-1.0.0 schema does not permit 0 valued strides.
        # Perhaps we'll want to allow this someday, to efficiently
        # represent an array of all the same value.
        if any(stride == 0 for stride in data.strides):
            data = np.ascontiguousarray(data)

        # Use the base array if that option is set or if the option
        # is unset and the AsdfConfig default is set
        cfg = config.get_config()
        if options.save_base or (options.save_base is None and cfg.default_array_save_base):
            base = util.get_array_base(data)
        else:
            base = data

        # The view computations that follow assume that the base array
        # is contiguous.  If not, we need to make a copy to avoid
        # writing a nonsense view.
        if not base.flags.forc:
            data = np.ascontiguousarray(data)
            base = util.get_array_base(data)

        shape = data.shape

        # possibly override options based on config settings
        if cfg.all_array_storage is not None:
            options.storage_type = cfg.all_array_storage
        if cfg.all_array_compression != "input":
            options.compression = cfg.all_array_compression
            options.compression_kwargs = cfg.all_array_compression_kwargs
        inline_threshold = cfg.array_inline_threshold

        if inline_threshold is not None and options.storage_type in ("inline", "internal"):
            if data.size < inline_threshold:
                options.storage_type = "inline"
            else:
                options.storage_type = "internal"
        ctx._blocks.options.set_options(data, options)

        # Compute the offset relative to the base array and not the
        # block data, in case the block is compressed.
        offset = data.ctypes.data - base.ctypes.data

        strides = None if data.flags.c_contiguous else data.strides
        dtype, byteorder = numpy_dtype_to_asdf_datatype(
            data.dtype,
            include_byteorder=(options.storage_type != "inline"),
        )

        result = {}

        result["shape"] = list(shape)
        if options.storage_type == "streamed":
            result["shape"][0] = "*"

        if options.storage_type == "inline":
            listdata = numpy_array_to_list(data)
            result["data"] = listdata
            result["datatype"] = dtype

        else:
            result["shape"] = list(shape)
            if options.storage_type == "streamed":
                result["shape"][0] = "*"
                result["source"] = -1
                ctx._blocks.set_streamed_write_block(base, data)
            else:
                result["source"] = ctx._blocks.make_write_block(base, options, obj)
            result["datatype"] = dtype
            result["byteorder"] = byteorder

            if offset > 0:
                result["offset"] = offset

            if strides is not None:
                result["strides"] = list(strides)

        if isinstance(data, ma.MaskedArray):
            if options.storage_type == "inline":
                ctx._blocks._set_array_storage(data.mask, "inline")

            result["mask"] = data.mask

        return result

    def from_yaml_tree(self, node, tag, ctx):
        import sys
        import weakref

        from asdf.tags.core import NDArrayType
        from asdf.tags.core.ndarray import asdf_datatype_to_numpy_dtype

        if isinstance(node, list):
            instance = NDArrayType(node, None, None, None, None, None, None)
            ctx._blocks._set_array_storage(instance, "inline")
            if not ctx._blocks._lazy_load:
                return instance._make_array()
            return instance

        if isinstance(node, dict):
            shape = node.get("shape", None)
            if "source" in node and "data" in node:
                msg = "Both source and data may not be provided at the same time"
                raise ValueError(msg)
            if "source" in node:
                source = node["source"]
                byteorder = node["byteorder"]
            else:
                source = node["data"]
                byteorder = sys.byteorder
            dtype = asdf_datatype_to_numpy_dtype(node["datatype"], byteorder) if "datatype" in node else None
            offset = node.get("offset", 0)
            strides = node.get("strides", None)
            mask = node.get("mask", None)

            if isinstance(source, int):
                # internal block
                data_callback = ctx.get_block_data_callback(source)
                instance = NDArrayType(source, shape, dtype, offset, strides, "A", mask, data_callback)
            elif isinstance(source, str):
                # external
                def data_callback(_attr=None, _ref=weakref.ref(ctx._blocks)):
                    if _attr not in (None, "cached_data", "data"):
                        raise AttributeError(f"_attr {_attr} is not supported")
                    blks = _ref()
                    if blks is None:
                        msg = "Failed to resolve reference to AsdfFile to read external block"
                        raise OSError(msg)
                    array = blks._load_external(source)
                    blks._set_array_storage(array, "external")
                    return array

                instance = NDArrayType(source, shape, dtype, offset, strides, "A", mask, data_callback)
            else:
                # inline
                instance = NDArrayType(source, shape, dtype, offset, strides, "A", mask)
                ctx._blocks._set_array_storage(instance, "inline")

            if not ctx._blocks._lazy_load:
                return instance._make_array()
            return instance

        msg = "Invalid ndarray description."
        raise TypeError(msg)

    def to_info(self, obj):
        return {"shape": obj.shape, "dtype": obj.dtype}
