
import os

import shutil
import requests
import tarfile
import warnings
import numpy as np
from astropy.table import Table, Row
from astropy import coordinates
from astropy import units as u
from astropy.utils.decorators import deprecated, deprecated_renamed_argument

import pyvo

from astroquery import log
from ..query import BaseQuery, BaseVOQuery
from ..utils import commons, parse_coordinates
from ..exceptions import InvalidQueryError, NoResultsWarning
from . import conf

__all__ = ['Heasarc', 'HeasarcClass']


class HeasarcClass(BaseVOQuery, BaseQuery):
    """Class for accessing HEASARC data with VO protocol using the Xamin backend.


    """

    # we can move url to Config later
    VO_URL = conf.VO_URL
    TAR_URL = conf.TAR_URL
    S3_BUCKET = conf.S3_BUCKET
    timeout = conf.timeout

    def __init__(self):
        """Initialize some basic useful variables"""
        super().__init__()
        self._tap = None
        self._datalink = None
        self._meta_info = None
        self._session = None

    @property
    def tap(self):
        """TAP service"""
        if self._tap is None:
            self._tap = pyvo.dal.TAPService(
                f'{self.VO_URL}/tap', session=self._session
            )
            self._session = self._tap._session
        return self._tap

    @property
    def _meta(self):
        """Queries and holds meta-information about the catalogs.

        This is a table that holds useful information such as
        the list of default columns per catalog, the reasonable default
        search radius per table that is appropriate for a mission etc.
        Instead of making a server call for each catalog for that type information,
        we do a single one and then post-process the resulting table.

        These are not meant to be used directly by the user.
        """
        if self._meta_info is None:
            query = (
                "SELECT split_part(name, '.', 1) AS table, "
                "split_part(name, '.', 2) AS par, "
                "CAST(value AS DECIMAL) AS value "
                "FROM metainfo "
                "WHERE (type = 'parameter' and relation = 'order') "
                "OR relation LIKE 'defaultSearchRadius' "
                "ORDER BY value"
            )
            self._meta_info = self.query_tap(query).to_table()
            self._meta_info['value'] = np.array(
                self._meta_info['value'], np.float32)
            self._meta_info = self._meta_info[self._meta_info['value'] > 0]
        return self._meta_info

    def _get_default_columns(self, catalog_name):
        """Get a list of default columns for a catalog

        Parameters
        ----------
        catalog_name : str
            The name of catalog as a str

        Return
        ------
        a list of column names

        """
        meta = self._meta[
            (self._meta['table'] == catalog_name)
            & (self._meta['par'] != '')
        ]
        meta.sort('value')
        defaults = meta['par']
        return list(defaults)

    def get_default_radius(self, catalog_name):
        """Get a mission-appropriate default radius for a catalog

        Parameters
        ----------
        catalog_name : str
            The name of catalog as a str

        Returns
        -------
        The radius as `~astropy.units.Quantity`

        """
        meta = self._meta[
            (self._meta['table'] == catalog_name)
            & (self._meta['par'] == '')
        ]
        radius = np.float32(meta['value'][0]) * u.arcmin
        return radius

    def _set_session(self, session):
        """Set requests.Session to use when querying the data

        Parameters
        ----------
        session : `~requests.Session`
            The requests.Session to use

        """
        if not isinstance(session, requests.Session):
            raise ValueError('session is not a `~requests.Session` instance')

        self._session = session

    def list_catalogs(self, *, master=False, keywords=None):
        """Return a table of all available catalogs with two columns
        (name, description)

        Parameters
        ----------
        master : bool
            Select only master catalogs. Default is False
        keywords : str or list
            a str or a list of str of keywords used as search
            terms for catalogs. Words with a str separated by a space
            are AND'ed, while words in a list are OR'ed

        Returns
        -------
        `~astropy.table.Table` with columns: name, description

        """
        if keywords is not None:
            if not isinstance(keywords, list):
                keywords = [keywords]
            if not all([isinstance(wrd, str) for wrd in keywords]):
                raise ValueError('non-str found in keywords elements')

        # use 'mast' to include both 'master' and 'mastr'
        names, desc = [], []
        for lab, tab in self.tap.tables.items():
            if 'TAP' in lab or (master and 'mast' not in lab):
                continue
            if keywords is not None:
                matched = any(
                    [
                        all([wrd.lower() in f'{lab} {tab.description}'.lower()
                            for wrd in wrds.split()])
                        for wrds in keywords
                    ]
                )
                if not matched:
                    continue
            names.append(lab)
            desc.append(tab.description)
        return Table({'name': names, 'description': desc})

    @deprecated(
        since='0.4.8',
        alternative='list_catalogs',
    )
    def query_mission_list(self, *, cache=True, get_query_payload=False):
        """Returns a list of all available mission catalogs with descriptions.

        This method is deprecated, and is included only for limited
        backward compatibility with the old astroquery.Heasarc that uses
        the Browse interface. Please use `list_catalogs` instead.

        """
        return self.list_catalogs(master=False)

    def list_columns(self, catalog_name, full=False):
        """Return the columns available in catalog_name as a table

        Parameters
        ----------
        catalog_name : str
            The name of catalog as a str
        full : bool
            If True, return all columns, otherwise, return the standard list
            of columns

        Returns
        -------
        result : `~astropy.table.Table`
            A table with columns: name, description, unit

        """
        tables = self.tap.tables
        if catalog_name not in tables.keys():
            msg = (f'{catalog_name} is not available as a public catalog. '
                   'Try passing keywords to `~Heasarc.list_catalogs` to find '
                   'the catalog name')
            raise ValueError(msg)

        default_cols = self._get_default_columns(catalog_name)

        names, desc, unit = [], [], []
        for col in tables[catalog_name].columns:
            if full or col.name in default_cols:
                names.append(col.name)
                desc.append(col.description)
                unit.append(col.unit or '')
        cols = Table({'name': names, 'description': desc, 'unit': unit})

        return cols

    @deprecated(
        since='0.4.8',
        alternative='list_columns',
    )
    def query_mission_cols(self, mission, *, cache=True,
                           get_query_payload=False, **kwargs):
        """Query around a specific object within a given mission catalog

        NOTE: This method is deprecated, and is included only for limited
        backward compatibility with the old astroquery.Heasarc that uses
        the Browse interface. Please use `list_columns` instead.

        Parameters
        ----------
        mission : str
            Mission catalog (short name) to search from
        fields : str, optional
            Return format for columns from the server available options:
            * Standard      : Return default catalog columns
            * All (default) : Return all catalog columns
            * <custom>      : User defined csv list of columns to be returned
        cache : bool, optional
            Defaults to True. If set overrides global caching behavior.
            See :ref:`caching documentation <astroquery_cache>`.
        All other parameters have no effect

        """
        fields = kwargs.get('fields', 'All')
        full = fields != 'Standard'
        cols = self.list_columns(mission, full=full)
        cols = [col.upper() for col in cols['name'] if '__' not in col]
        return cols

    def query_tap(self, query, *, maxrec=None):
        """
        Send query to HEASARC's Xamin TAP using ADQL.
        Results in `~pyvo.dal.TAPResults` format.
        result.to_table gives `~astropy.table.Table` format.

        Parameters
        ----------
        query : str
            ADQL query to be executed
        maxrec : int
            maximum number of records to return

        Returns
        -------
        result : `~pyvo.dal.TAPResults`
            TAP query result.
        result.to_table : `~astropy.table.Table`
            TAP query result as `~astropy.table.Table`
        result.to_qtable : `~astropy.table.QTable`
            TAP query result as `~astropy.table.QTable`

        """
        log.debug(f'TAP query: {query}')
        self._saved_query = query
        return self.tap.search(query, language='ADQL', maxrec=maxrec)

    @deprecated_renamed_argument(
        ('mission', 'fields', 'resultmax', 'entry', 'coordsys', 'equinox',
         'displaymode', 'action', 'sortvar', 'cache'),
        ('catalog', 'columns', 'maxrec', None, None, None,
         None, None, None, None),
        since=['0.4.8']*10,
        arg_in_kwargs=(False, True, True, True, True, True,
                       True, True, True, False)
    )
    def query_region(self, position=None, catalog=None, radius=None, *,
                     spatial='cone', width=None, polygon=None, add_offset=False,
                     get_query_payload=False, columns=None, cache=False,
                     verbose=False, maxrec=None,
                     **kwargs):
        """Queries the HEASARC TAP server around a coordinate and returns a
        `~astropy.table.Table` object.

        Parameters
        ----------
        position : str, `astropy.coordinates` object
            Gives the position of the center of the cone or box if performing
            a cone or box search. Required if spatial is ``'cone'`` or
            ``'box'``. Ignored if spatial is ``'polygon'`` or ``'all-sky'``.
        catalog : str
            The catalog to query. To list the available catalogs, use
            :meth:`~astroquery.heasarc.HeasarcClass.list_catalogs`.
        spatial : str
            Type of spatial query: ``'cone'``, ``'box'``, ``'polygon'``, and
            ``'all-sky'``. Defaults to ``'cone'``.
        radius : str or `~astropy.units.Quantity` object, [optional for
            spatial == ``'cone'``].
            The string must be parsable by `~astropy.coordinates.Angle`. The
            appropriate `~astropy.units.Quantity` object from
            `astropy.units` may also be used. If None, a default value
            appropriate for the selected catalog is used. To see the default
            radius for the catalog, see
            `get_default_radius`.
        width : str, `~astropy.units.Quantity` object [Required for
            spatial == ``'box'``.]
            The string must be parsable by `~astropy.coordinates.Angle`. The
            appropriate `~astropy.units.Quantity` object from `astropy.units`
            may also be used.
        polygon : list, [Required for spatial is ``'polygon'``]
            A list of ``(ra, dec)`` pairs (as tuples), in decimal degrees,
            outlining the polygon to search in. It can also be a list of
            `astropy.coordinates` object or strings that can be parsed by
            `astropy.coordinates.ICRS`.
        add_offset: bool
            If True and spatial=='cone', add a search_offset column that
            indicates the separation (in arcmin) between the requested
            coordinate and the entry coordinates in the catalog. Default is False.
        get_query_payload : bool, optional
            If `True` then returns the generated ADQL query as str.
            Defaults to `False`.
        columns : str, optional
            Target column list with value separated by a comma(,).
            Use * for all the columns. The default is to return a subset
            of the columns that are generally the most useful.
        verbose : bool, optional
            If False, suppress vo warnings.
        maxrec : int, optional
            Maximum number of records


        Returns
        -------
        table : A `~astropy.table.Table` object.
        """
        # if verbose is False then suppress any VOTable related warnings
        if not verbose:
            commons.suppress_vo_warnings()

        if catalog is None:
            raise InvalidQueryError("catalog name is required! Use 'xray' "
                                    "to search the master X-ray catalog")

        if columns is None:
            columns = ', '.join(self._get_default_columns(catalog))
            if '__row' not in columns:
                columns += ',__row'

        if spatial.lower() == 'all-sky':
            where = ''
        elif spatial.lower() == 'polygon':
            try:
                coords_list = [parse_coordinates(coord).icrs
                               for coord in polygon]
            except TypeError:
                try:
                    coords_list = [coordinates.SkyCoord(*coord).icrs
                                   for coord in polygon]
                except u.UnitTypeError:
                    warnings.warn("Polygon endpoints are being interpreted as "
                                  "RA/Dec pairs specified in decimal degree "
                                  "units.")
                    coords_list = [
                        coordinates.SkyCoord(*coord, unit='deg').icrs
                        for coord in polygon
                    ]

            coords_str = [f'{coord.ra.deg},{coord.dec.deg}'
                          for coord in coords_list]
            where = (" WHERE CONTAINS(POINT('ICRS',ra,dec),"
                     f"POLYGON('ICRS',{','.join(coords_str)}))=1")
        else:
            coords_icrs = parse_coordinates(position).icrs
            ra, dec = coords_icrs.ra.deg, coords_icrs.dec.deg

            if spatial.lower() == 'cone':
                if radius is None:
                    radius = self.get_default_radius(catalog)
                elif isinstance(radius, str):
                    radius = coordinates.Angle(radius)
                where = (" WHERE CONTAINS(POINT('ICRS',ra,dec),CIRCLE("
                         f"'ICRS',{ra},{dec},{radius.to(u.deg).value}))=1")
                # add search_offset for the case of cone
                if add_offset:
                    columns += (",DISTANCE(POINT('ICRS',ra,dec), "
                                f"POINT('ICRS',{ra},{dec})) as search_offset")
            elif spatial.lower() == 'box':
                if isinstance(width, str):
                    width = coordinates.Angle(width)
                where = (" WHERE CONTAINS(POINT('ICRS',ra,dec),"
                         f"BOX('ICRS',{ra},{dec},{width.to(u.deg).value},"
                         f"{width.to(u.deg).value}))=1")
            else:
                raise ValueError("Unrecognized spatial query type. Must be one"
                                 " of 'cone', 'box', 'polygon', or 'all-sky'.")

        adql = f'SELECT {columns} FROM {catalog}{where}'

        if get_query_payload:
            return adql
        response = self.query_tap(query=adql, maxrec=maxrec)

        # save the response in case we want to use it later
        self._last_result = response
        self._last_catalog_name = catalog

        table = response.to_table()
        if add_offset:
            table['search_offset'].unit = u.arcmin
        if len(table) == 0:
            warnings.warn(
                NoResultsWarning("No matching rows were found in the query.")
            )
        return table

    @deprecated(
        since='0.4.8',
        alternative='query_region'
    )
    def query_object(self, object_name, mission, *,
                     cache=True, get_query_payload=False,
                     **kwargs):
        """Query around a specific object within a given mission catalog

        Parameters
        ----------
        object_name : str
            Object to query around. To set search radius use the 'radius'
            parameter.
        mission : str
            Mission catalog to search from
        cache : bool
            Defaults to True. If set overrides global caching behavior.
            See :ref:`caching documentation <astroquery_cache>`.
        **kwargs :
            see `~astroquery.heasarc.HeasarcClass._args_to_payload` for list
            of additional parameters that can be used to refine search query.

        """
        pos = coordinates.SkyCoord.from_name(object_name)
        return self.query_region(pos, catalog=mission, spatial='cone',
                                 get_query_payload=get_query_payload)

    def locate_data(self, query_result=None, catalog_name=None):
        """Get links to data products
        Use vo/datalinks to query the data products for some query_results.

        Parameters
        ----------
        query_result : `astropy.table.Table` or `astropy.table.Row`, optional
            A table that contain the search results. Typically as
            returned by query_region. If None, use the table from the
            most recent query_region call.
        catalog_name : str
            The catalog name for the which the query_result belongs to.
            If None, use the one from the most recent query_region call.

        Returns
        -------
        table : A `~astropy.table.Table` object.
        """
        if query_result is None:
            if (
                not hasattr(self, '_last_result')
                or self._last_result is None
            ):
                raise ValueError('query_result is None, and none '
                                 'found from a previous search')
            else:
                query_result = self._last_result

        if isinstance(query_result, Row):
            query_result = query_result.table[[query_result.index]]

        elif not isinstance(query_result, Table):
            raise TypeError(
                'query_result need to be an astropy.table.Table or '
                'astropy.table.Row'
            )

        # make sure we have a column __row
        if '__row' not in query_result.colnames:
            raise ValueError('No __row column found in query_result. '
                             'query_result needs to be the output of '
                             'query_region or a subset.')

        if catalog_name is None:
            catalog_name = self._last_catalog_name
        if not (
            isinstance(catalog_name, str)
            and catalog_name in self.tap.tables.keys()
        ):
            raise ValueError(f'Unknown catalog name: {catalog_name}')

        # datalink url
        dlink_url = f'{self.VO_URL}/datalink/{catalog_name}'
        query = pyvo.dal.adhoc.DatalinkQuery(
            baseurl=dlink_url,
            id=query_result['__row'],
            session=self._session
        )

        dl_result = pyvo.dal.DALResults(
            query.execute_votable(post=True),
            url=query.queryurl,
            session=query._session
        ).to_table()

        # include rows that have directory links (i.e. data) and those
        # that report errors (usually means there are no data products)
        dl_result = dl_result[np.ma.mask_or(
            dl_result['content_type'] == 'directory',
            dl_result['error_message'] != '',
            shrink=False
        )]
        dl_result = dl_result[['ID', 'access_url', 'content_length', 'error_message']]

        # add sciserver and s3 columns
        newcol = [
            f"/FTP/{row.split('FTP/')[1]}".replace('//', '/')
            if 'FTP' in row else ''
            for row in dl_result['access_url']
        ]
        dl_result.add_column(newcol, name='sciserver', index=2)
        newcol = [f"s3://{self.S3_BUCKET}/{row[5:]}" if row != '' else ''
                  for row in dl_result['sciserver']]
        dl_result.add_column(newcol, name='aws', index=3)

        return dl_result

    def enable_cloud(self, provider='aws', profile=None):
        """
        Enable downloading public files from the cloud.
        Requires the boto3 library to function.

        Parameters
        ----------
        provider : str
            Which cloud data provider to use. Currently, only 'aws' is
            supported.
        profile : str
            Profile to use to identify yourself to the cloud provider
            (usually in ~/.aws/config).
        """

        try:
            import boto3
            import botocore
        except ImportError:
            raise ImportError(
                'The cloud feature requires the boto3 package. '
                'Install it first.'
            )

        if profile is None:
            log.info('Enabling anonymous cloud data access ...')
            config = botocore.client.Config(
                signature_version=botocore.UNSIGNED)
            self.s3_resource = boto3.resource('s3', config=config)

        elif isinstance(profile, bool) and not profile:
            # profile is False, use system env credentials
            self.s3_resource = boto3.resource('s3')

        else:
            log.info(f'Enabling cloud data access with profile: {profile} ...')
            session = boto3.session.Session(profile_name=profile)
            self.s3_resource = session.resource(service_name='s3')

        self.s3_client = self.s3_resource.meta.client

    def download_data(self, links, host='heasarc', location='.'):
        """Download data products in links with a choice of getting the
        data from either the heasarc server, sciserver, or the cloud in AWS.


        Parameters
        ----------
        links : `astropy.table.Table` or `astropy.table.Row`
            The result from locate_data
        host : str
            The data host. The options are: heasarc (default), sciserver, aws.
            If host == 'sciserver', data is copied from the local mounted
            data drive.
            If host == 'aws', data is downloaded from Amazon S3 Open
            Data Repository.
        location : str
            local folder where the downloaded file will be saved.
            Default is current working directory

        Note that ff you are downloading large datasets (more 10 10GB),
        from the main heasarc server, it is recommended that you split
        it up, so that if the downloaded is interrupted, you do not need
        to start again.
        """

        if len(links) == 0:
            raise ValueError('Input links table is empty')

        if isinstance(links, Row):
            links = links.table[[links.index]]

        if host not in ['heasarc', 'sciserver', 'aws']:
            raise ValueError('host has to be one of heasarc, sciserver, aws')

        host_column = 'access_url' if host == 'heasarc' else host
        if host_column not in links.colnames:
            raise ValueError(
                f'No {host_column} column found in the table. Call '
                '`~locate_data` first'
            )

        # remove rows that dont have data, if any
        if 'error_message' in links.colnames:
            links = links[links['error_message'] == '']

        if host == 'heasarc':

            log.info('Downloading data from the HEASARC ...')
            self._download_heasarc(links, location)

        elif host == 'sciserver':

            log.info('Copying data on Sciserver ...')
            self._copy_sciserver(links, location)

        elif host == 'aws':

            log.info('Downloading data AWS S3 ...')
            self._download_s3(links, location)

    def _download_heasarc(self, links, location='.'):
        """Download data from the heasarc main server using xamin's tar servlet

        Do not call directly.
        Users should be using `~self.download_data` instead

        Parameters
        ----------
        links : `astropy.table.Table`
            The result from locate_data
        location : str
            local folder where the downloaded file will be saved.
            Default is current working directory

        """
        # The limit comes from the size of the string in the POST request
        if 'content_length' in links.columns:
            size = links['content_length'].sum() / 2**30
            if size > 10:
                warnings.warn(
                    f"The size of the requested file is large {size:.3f} GB. "
                    "If the download is interrupted, you may need to start "
                    "again. Consider downloading the data in chunks."
                )

        file_list = [f"/FTP/{link.split('FTP/')[1]}"
                     for link in links['access_url']]
        params = {
            'files': f'>{"&&>".join(file_list)}&&',
            'filter': ''
        }

        # get local_filepath name
        local_filepath = f'{location}/heasarc-data.tar'
        iname = 1
        while os.path.exists(local_filepath):
            local_filepath = f'{location}/heasarc-data.{iname}.tar'
            iname += 1

        log.info(f'Downloading to {local_filepath} ...')
        self._download_file(self.TAR_URL, local_filepath,
                            timeout=self.timeout,
                            continuation=False, cache=False, method="POST",
                            head_safe=False, data=params, verbose=False)

        # if all good and we have a tar file untar it
        if tarfile.is_tarfile(local_filepath):
            log.info(f'Untar {local_filepath} to {location} ...')
            tfile = tarfile.TarFile(local_filepath)
            tfile.extractall(path=location, filter="fully_trusted")
            tfile.close()
            os.remove(local_filepath)
        else:
            raise ValueError(
                'An error occurred when downloading the data. Retry again.'
            )

    def _copy_sciserver(self, links, location='.'):
        """Copy data from the local archive on sciserver

        Do not call directly.
        Users should be using `~self.download_data` instead

        """
        if not os.path.exists('/FTP/'):
            raise FileNotFoundError(
                'No data archive found. This should be run on Sciserver '
                'with the data drive mounted.'
            )

        # make sure the output folder exits
        os.makedirs(location, exist_ok=True)

        for link in links['sciserver']:
            link = str(link)
            log.info(f'Copying to {link} from the data drive ...')
            if not os.path.exists(link):
                raise ValueError(
                    f'No data found in {link}. '
                    'Make sure you are running this on Sciserver. '
                    'If you think data is missing, please contact the '
                    'Heasarc Help desk'
                )
            if os.path.isdir(link):
                download_dir = os.path.basename(link.strip('/'))
                shutil.copytree(link, f'{location}/{download_dir}')
            else:
                shutil.copy(link, location)

    def _download_s3(self, links, location='.'):
        """Download data from AWS S3
        Assuming open access.

        Do not call directly.
        Users should be using `~self.download_data` instead

        """
        keys_list = [link for link in links['aws']]
        if not hasattr(self, 's3_resource'):
            # all the data is public for now; no profile is needed
            self.enable_cloud(provider='aws', profile=None)

        def _s3_tree_download(client, bucket_name, path, local):
            """Download nested keys from s3"""
            response = client.list_objects_v2(Bucket=bucket_name, Prefix=path)
            content = response.get('Contents', [])
            for obj in content:
                key = obj['Key']
                path2 = '/'.join(path.strip('/').split('/')[:-1])
                dest = os.path.join(local, key[len(path2)+1:])
                os.makedirs(os.path.dirname(dest), exist_ok=True)
                client.download_file(bucket_name, key, dest)

        # loop through the requested links
        for key in keys_list:
            log.info(f'downloading {key}')
            path = key.replace(f's3://{self.S3_BUCKET}/', '')
            _s3_tree_download(self.s3_client, self.S3_BUCKET, path, location)


Heasarc = HeasarcClass()
