"""
Utilities for embedded ADSF files in FITS.
"""
import io
import re

import numpy as np

from . import asdf, block, generic_io, util

try:
    from astropy.io import fits
except ImportError:
    raise ImportError("AsdfInFits requires astropy")


ASDF_EXTENSION_NAME = "ASDF"
FITS_SOURCE_PREFIX = "fits:"


__all__ = ["AsdfInFits"]


class _FitsBlock:
    def __init__(self, hdu):
        self._hdu = hdu

    def __repr__(self):
        return f"<FitsBlock {self._hdu.name},{self._hdu.ver}>"

    def __len__(self):
        return self._hdu.data.nbytes

    @property
    def data(self):
        return self._hdu.data

    @property
    def readonly(self):
        return False

    @property
    def array_storage(self):
        return "fits"

    @property
    def trust_data_dtype(self):
        # astropy.io.fits returns arrays in native byte order
        # when it has to apply scaling.  In that case, we don't
        # want to interpret the bytes as big-endian, since astropy
        # has already converted them properly.
        return True


class _EmbeddedBlockManager(block.BlockManager):
    def __init__(self, hdulist, asdffile):
        self._hdulist = hdulist

        super().__init__(asdffile)

    def get_block(self, source):
        if isinstance(source, str) and source.startswith(FITS_SOURCE_PREFIX):
            parts = re.match(
                # All printable ASCII characters are allowed in EXTNAME
                "((?P<name>[ -~]+),)?(?P<ver>[0-9]+)",
                source[len(FITS_SOURCE_PREFIX) :],
            )
            if parts is not None:
                ver = int(parts.group("ver"))
                if parts.group("name"):
                    pair = (parts.group("name"), ver)
                else:
                    pair = ver
                return _FitsBlock(self._hdulist[pair])
            else:
                raise ValueError(f"Can not parse source '{source}'")

        return super().get_block(source)

    def get_source(self, block):
        if isinstance(block, _FitsBlock):
            for i, hdu in enumerate(self._hdulist):
                if hdu is block._hdu:
                    if hdu.name == "":
                        return f"{FITS_SOURCE_PREFIX}{i}"
                    else:
                        return f"{FITS_SOURCE_PREFIX}{hdu.name},{hdu.ver}"
            raise ValueError("FITS block seems to have been removed")

        return super().get_source(block)

    def find_or_create_block_for_array(self, arr, ctx):

        base = util.get_array_base(arr)
        for hdu in self._hdulist:
            if hdu.data is None:
                continue
            if base is util.get_array_base(hdu.data):
                return _FitsBlock(hdu)

        return super().find_or_create_block_for_array(arr, ctx)


class AsdfInFits(asdf.AsdfFile):
    """
    Embed ASDF tree content in a FITS file.

    The YAML rendering of the tree is stored in a special FITS
    extension with the EXTNAME of ``ASDF``.  Arrays in the ASDF tree
    may refer to binary data in other FITS extensions by setting
    source to a string with the prefix ``fits:`` followed by an
    ``EXTNAME``, ``EXTVER`` pair, e.g. ``fits:SCI,0``.

    Examples
    --------
    Create a FITS file with ASDF structure, based on an existing FITS
    file::

        import numpy as np
        from astropy.io import fits

        from asdf import fits_embed

        hdulist = fits.HDUList([
            fits.PrimaryHDU(),
            fits.ImageHDU(np.arange(512, dtype=np.float32), name='SCI'),
            fits.ImageHDU(np.zeros(512, dtype=np.int16), name='DQ')])

        tree = {
            'model': {
                'sci': {
                    'data': hdulist['SCI'].data,
                    'wcs': 'WCS info'
                },
                'dq': {
                    'data': hdulist['DQ'].data,
                    'wcs': 'WCS info'
                }
            }
        }

        ff = fits_embed.AsdfInFits(hdulist, tree)
        ff.write_to('test.fits')  # doctest: +SKIP
    """

    def __init__(self, hdulist=None, tree=None, **kwargs):
        if hdulist is None:
            hdulist = fits.HDUList()
        super().__init__(tree=tree, **kwargs)
        self._blocks = _EmbeddedBlockManager(hdulist, self)
        self._hdulist = hdulist
        self._close_hdulist = False

    def __exit__(self, type, value, traceback):
        super().__exit__(type, value, traceback)
        if self._close_hdulist:
            self._hdulist.close()
        self._tree = {}

    def close(self):
        super().close()
        if self._close_hdulist:
            self._hdulist.close()
        self._tree = {}

    @classmethod
    def open(
        cls,
        fd,
        uri=None,
        validate_checksums=False,
        extensions=None,
        ignore_version_mismatch=True,
        ignore_unrecognized_tag=False,
        strict_extension_check=False,
        ignore_missing_extensions=False,
        **kwargs,
    ):
        """Creates a new AsdfInFits object based on given input data

        Parameters
        ----------
        fd : FITS HDUList instance, URI string, or file-like object
            May be an already opened instance of a FITS HDUList instance,
            string ``file`` or ``http`` URI, or a Python file-like object.

        uri : str, optional
            The URI for this ASDF file.  Used to resolve relative
            references against.  If not provided, will be
            automatically determined from the associated file object,
            if possible and if created from `asdf.open`.

        validate_checksums : bool, optional
            If `True`, validate the blocks against their checksums.
            Requires reading the entire file, so disabled by default.

        extensions : object, optional
            Additional extensions to use when reading and writing the file.
            May be any of the following: `asdf.extension.AsdfExtension`,
            `asdf.extension.Extension`, `asdf.extension.AsdfExtensionList`
            or a `list` extensions.

        ignore_version_mismatch : bool, optional
            When `True`, do not raise warnings for mismatched schema versions.

        strict_extension_check : bool, optional
            When `True`, if the given ASDF file contains metadata about the
            extensions used to create it, and if those extensions are not
            installed, opening the file will fail. When `False`, opening a file
            under such conditions will cause only a warning. Defaults to
            `False`.

        ignore_missing_extensions : bool, optional
            When `True`, do not raise warnings when a file is read that
            contains metadata about extensions that are not available. Defaults
            to `False`.

        validate_on_read : bool, optional
            DEPRECATED. When `True`, validate the newly opened file against tag
            and custom schemas.  Recommended unless the file is already known
            to be valid.
        """
        return cls._open_impl(
            fd,
            uri=uri,
            validate_checksums=validate_checksums,
            extensions=extensions,
            ignore_version_mismatch=ignore_version_mismatch,
            ignore_unrecognized_tag=ignore_unrecognized_tag,
            strict_extension_check=strict_extension_check,
            ignore_missing_extensions=ignore_missing_extensions,
            **kwargs,
        )

    @classmethod
    def _open_impl(
        cls,
        fd,
        uri=None,
        validate_checksums=False,
        extensions=None,
        ignore_version_mismatch=True,
        ignore_unrecognized_tag=False,
        strict_extension_check=False,
        ignore_missing_extensions=False,
        **kwargs,
    ):

        close_hdulist = False
        if isinstance(fd, fits.hdu.hdulist.HDUList):
            hdulist = fd
        else:
            uri = generic_io.get_uri(fd)
            try:
                hdulist = fits.open(fd)
                # Since we created this HDUList object, we need to be
                # responsible for cleaning up upon close() or __exit__
                close_hdulist = True
            except OSError:
                raise ValueError(f"Failed to parse given file '{uri}'. Is it FITS?")

        self = cls(
            hdulist,
            uri=uri,
            ignore_version_mismatch=ignore_version_mismatch,
            ignore_unrecognized_tag=ignore_unrecognized_tag,
        )

        self._close_hdulist = close_hdulist

        try:
            asdf_extension = hdulist[ASDF_EXTENSION_NAME]
        except (KeyError, IndexError, AttributeError):
            # This means there is no ASDF extension
            return self

        generic_file = generic_io.get_file(io.BytesIO(asdf_extension.data), mode="r", uri=uri)

        try:
            return cls._open_asdf(
                self,
                generic_file,
                validate_checksums=validate_checksums,
                extensions=extensions,
                strict_extension_check=strict_extension_check,
                ignore_missing_extensions=ignore_missing_extensions,
                **kwargs,
            )
        except RuntimeError:
            self.close()
            raise

    def _create_hdu(self, buff, use_image_hdu):
        # Allow writing to old-style ImageHDU for backwards compatibility
        if use_image_hdu:
            array = np.frombuffer(buff.getvalue(), np.uint8)
            return fits.ImageHDU(array, name=ASDF_EXTENSION_NAME)
        else:
            data = np.array(buff.getbuffer(), dtype=np.uint8)[None, :]
            fmt = f"{len(data[0])}B"
            column = fits.Column(array=data, format=fmt, name="ASDF_METADATA")
            return fits.BinTableHDU.from_columns([column], name=ASDF_EXTENSION_NAME)

    def _update_asdf_extension(
        self, all_array_storage=None, all_array_compression=None, pad_blocks=False, use_image_hdu=False, **kwargs
    ):
        if self.blocks.streamed_block is not None:
            raise ValueError("Can not save streamed data to ASDF-in-FITS file.")

        buff = io.BytesIO()
        super().write_to(
            buff,
            all_array_storage=all_array_storage,
            all_array_compression=all_array_compression,
            pad_blocks=pad_blocks,
            include_block_index=False,
            **kwargs,
        )

        if ASDF_EXTENSION_NAME in self._hdulist:
            del self._hdulist[ASDF_EXTENSION_NAME]
        self._hdulist.append(self._create_hdu(buff, use_image_hdu))

    def write_to(
        self,
        filename,
        all_array_storage=None,
        all_array_compression=None,
        pad_blocks=False,
        use_image_hdu=False,
        *args,
        **kwargs,
    ):
        if "auto_inline" in kwargs:
            asdf_kwargs = {"auto_inline": kwargs.pop("auto_inline")}
        else:
            asdf_kwargs = {}

        self._update_asdf_extension(
            all_array_storage=all_array_storage,
            all_array_compression=all_array_compression,
            pad_blocks=pad_blocks,
            use_image_hdu=use_image_hdu,
            **asdf_kwargs,
        )

        self._hdulist.writeto(filename, *args, **kwargs)

    def update(self, all_array_storage=None, all_array_compression=None, pad_blocks=False, **kwargs):
        raise NotImplementedError("In-place update is not currently implemented for ASDF-in-FITS")