# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Various utilities and cookbook-like things.
"""


# STDLIB
import codecs
import contextlib
import io
import re
import gzip

from distutils import version


__all__ = [
    'convert_to_writable_filelike',
    'stc_reference_frames',
    'coerce_range_list_param',
    ]


@contextlib.contextmanager
def convert_to_writable_filelike(fd, compressed=False):
    """
    Returns a writable file-like object suitable for streaming output.

    Parameters
    ----------
    fd : file path string or writable file-like object
        May be:

            - a file path, in which case it is opened, and the file
              object is returned.

            - an object with a :meth:``write`` method, in which case that
              object is returned.

    compressed : bool, optional
        If `True`, create a gzip-compressed file.  (Default is `False`).

    Returns
    -------
    fd : writable file-like object
    """
    if isinstance(fd, str):
        if fd.endswith('.gz') or compressed:
            with gzip.GzipFile(fd, 'wb') as real_fd:
                encoded_fd = io.TextIOWrapper(real_fd, encoding='utf8')
                yield encoded_fd
                encoded_fd.flush()
                real_fd.flush()
                return
        else:
            with open(fd, 'wt', encoding='utf8') as real_fd:
                yield real_fd
                return
    elif hasattr(fd, 'write'):
        assert callable(fd.write)

        if compressed:
            fd = gzip.GzipFile(fileobj=fd)

        # If we can't write Unicode strings, use a codecs.StreamWriter
        # object
        needs_wrapper = False
        try:
            fd.write('')
        except TypeError:
            needs_wrapper = True

        if not hasattr(fd, 'encoding') or fd.encoding is None:
            needs_wrapper = True

        if needs_wrapper:
            yield codecs.getwriter('utf-8')(fd)
            fd.flush()
        else:
            yield fd
            fd.flush()

        return
    else:
        raise TypeError("Can not be coerced to writable file-like object")


# <http://www.ivoa.net/Documents/REC/DM/STC-20071030.html>
stc_reference_frames = set([
    'FK4', 'FK5', 'ECLIPTIC', 'ICRS', 'GALACTIC', 'GALACTIC_I', 'GALACTIC_II',
    'SUPER_GALACTIC', 'AZ_EL', 'BODY', 'GEO_C', 'GEO_D', 'MAG', 'GSE', 'GSM',
    'SM', 'HGC', 'HGS', 'HEEQ', 'HRTN', 'HPC', 'HPR', 'HCC', 'HGI',
    'MERCURY_C', 'VENUS_C', 'LUNA_C', 'MARS_C', 'JUPITER_C_III',
    'SATURN_C_III', 'URANUS_C_III', 'NEPTUNE_C_III', 'PLUTO_C', 'MERCURY_G',
    'VENUS_G', 'LUNA_G', 'MARS_G', 'JUPITER_G_III', 'SATURN_G_III',
    'URANUS_G_III', 'NEPTUNE_G_III', 'PLUTO_G', 'UNKNOWNFrame'])


def coerce_range_list_param(p, frames=None, numeric=True):
    """