from abc import ABC
from requests.exceptions import HTTPError
from ensembl_rest_client.base import cache as erc_cache
import requests
import json
import pandas as pd
import time
# import pprint as pp


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class _BaseRestClient(ABC):
    """Base level Ensembl REST client, do not use directly.

    Parameters
    ----------
    url : `str`, optional, default: `https://rest.ensembl.org`
        The main URL were the endpoints are attached. Note the human
        genome build 37. Can be accessed at https://grch37.ensembl.org
        . However, please note that not all methods will work the
        GRCh37 version.
    pandas : `bool`, optional, default: `False`
        Return all data as `pandas.DataFrames`
    retries : `int`, optional, default: `2`
        The number of times to retry the query if it should fail with a
        potentially recoverable status code
    max_requests_per_sec : `int`, optional, default: `15`
        A positive integer. If the rate exceeds this then a pause is
        initiated to bring the rate into line
    ping: `bool`, optional, default: `True`
        Perform a ping when the object is instantiated. If the ping fails
        then an HTTPError will be raised
    cache : `str` or `bool`, optional, default: `False`
        Should caching be used, if `False`, then no cache is used. If `True`
        then a cache file will be created and used that corresponds to the
        endpoint URL. If it is a string, then that is used as a specific file
        to cache the data. Important, caches are not safe across multiple
        processes (and possibly threads) they are implemented with the Python
        sheve. If you are using multi-processing please supply a filename for
        each process. Note, that the cache will be used if the endpoint, kwargs
        and method type (post or get) are the same and was issued within the
        ``cache_seconds_cutoff``.
    cache_seconds_cutoff : `int`, optional, default `604800`
        The time in seconds that cache qeures must be within in order to be
        considered for use. The default is the number of seconds in 1 week.

    Raises
    ------
    HTTPError
        If the actual ping request fails or the ping returns a False value.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, url='https://rest.ensembl.org', pandas=False,
                 retries=2, max_requests_per_sec=15, ping=True,
                 cache=False, cache_seconds_cutoff=7 * erc_cache.DAY):
        # Make sure the URL has no trailing /
        self._url = url if not url.endswith('/') else url.rstrip('/')
        self.pandas = pandas
        self._max_requests_per_sec = max_requests_per_sec

        # Will hold the number of request that have been made after the
        # previous time check
        self._req_count = 0

        # Initialise the time as a starting point even though no requests have
        # been made yet
        self._last_req = time.time()
        self._retries = retries

        if self._retries < 0:
            raise ValueError("retries should be a positive integer or 0")

        if self._max_requests_per_sec < 1:
            raise ValueError(
                "max_requests_per_sec should be a positive integer"
            )

        # Get the cache and update/add method (or dummy if cache is not being
        # used)
        self.cache = erc_cache.get_cache(
            cache, seconds_cutoff=cache_seconds_cutoff
        )

        # Perform a ping check
        if ping is True:
            self.__class__.ping(url=self.url)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def ping(cls, url='https://rest.ensembl.org'):
        """
        Ensure the Ensembl REST server is up before starting the test.

        Raises
        ------
        HTTPError
            If the actual ping request fails or the ping returns a False value
        """
        ping_endpoint = '{0}/info/ping?'.format(url)
        r = requests.get(
            ping_endpoint,
            headers={"Content-Type": "application/json"}
        )

        if not r.ok:
            r.raise_for_status()

        decoded = r.json()
        if bool(int(decoded['ping'])) is False:
            raise HTTPError(
                "The response was good but the ping was 0",
                response=r
            )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def rest_query(self, method_type, endpoint, **kwargs):
        """
        Perform a query against the Ensembl REST server

        Parameters
        ----------
        method_type : str
            The query type, can either be GET or POST
        endpoint : str
            The query endpoint, this will be joined with the root URL
        **kwargs
            Arguments specific to the query, if it is a post query, this should
            contain a `maximum_post_size` kwarg. If it is a region query, it
            should contain a `slice_length` kwarg. These are checked prior to
            querying the server. Other kwargs are specific for the query

        Returns
        -------
        data : list or dict or :obj:`pandas.DataFrame`
            The rest query data

        Raises
        ------
        HTTPError
            If there is an unrecoverable error with the query or the number
            of tries is reached
        """
        # Make sure the endpoint has a starting / and join to the root URL
        # to create a query endpoint
        endpoint = endpoint if endpoint.startswith('/') \
            else '/{0}'.format(endpoint)
        query_endpoint = '{0}{1}'.format(self._url, endpoint)

        # Loop until we have reached the number of retries, the + 1 is the
        # initial query
        for i in range(self._retries + 1):
            # Check the cache
            cache_key = self.cache.get_key(
                (endpoint, method_type, kwargs)
            )
            cache_entry = self.cache.get_from_cache(cache_key)
            if cache_entry is not None:
                return cache_entry

            # Do we need to wait a bit before querying. It might make sense
            # to make this a class variable as multiple classes could be
            # querying
            self._check_rate_limit()

            try:
                if method_type == 'GET':
                    data = self._do_get_query(query_endpoint, **kwargs)
                    self.cache.add_to_cache(cache_key, data)
                    return data
                if method_type == 'POST':
                    data = self._do_post_query(query_endpoint, **kwargs)
                    self.cache.add_to_cache(cache_key, data)
                    return data
                else:
                    raise ValueError(
                        "unknown method type: '{0}'".format(method_type)
                    )
            except HTTPError as e:
                # If we get an error then we assess if we can recover from it
                # and retry
                if i == self._retries or \
                   self._assess_status_code(e.response) is False:
                    raise

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def url(self):
        """The rest base URL
        """
        return self._url

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _check_rate_limit(self):
        """
        """
        # check if we need to rate limit ourselves
        if self._req_count >= self._max_requests_per_sec:
            delta = time.time() - self._last_req
            if delta < 1:
                time.sleep(1 - delta)
            self._last_req = time.time()
            self._req_count = 0
        self._req_count += 1

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _assess_status_code(self, response):
        """
        Check the error status code and see if we can potentially recover from
        it

        Parameters
        ----------
        response : :obj:`requests.response`
            The request response object that has the status code

        Returns
        -------
        retry : bool
            If `True`, it means it is worth retrying the query if `False` then
            no we should just error out

        Raises
        ------
        ValueError
            If we get a status code that we do not know how to handle
        """
        if response.status_code == 400:
            return False
        elif response.status_code > 500:
            return False
        elif response.status_code == 404:
            return False
        elif response.status_code == 404:
            time.sleep(5)
            return True
        elif response.status_code == 408:
            time.sleep(5)
            return True
        elif response.status_code == 429:
            try:
                time.sleep(float(response.headers['retry-after']))
            except KeyError:
                time.sleep(5)
            return True
        else:
            raise ValueError(
                "unhandled response code: {0}".format(response.status_code)
            )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _do_get_query(self, endpoint, **kwargs):
        """
        Perform a GET query against the endpoint

        Parameters
        ----------
        endpoint : str
            A join of the base  URL and the specific endpoint
        **kwargs
            Arguments specific to the query

        Returns
        -------
        data : list or dict or :obj:`pandas.DataFrame`
            The rest query data
        """
        endpoint = '{0}?'.format(endpoint)
        kwargs, normalise = self._set_kwargs(kwargs)

        req_kwargs = dict(
            headers={"Content-Type": "application/json"}
        )
        if len(kwargs) > 0:
            req_kwargs['params'] = kwargs

        r = requests.get(endpoint, **req_kwargs)

        if not r.ok:
            r.raise_for_status()

        return self._return_data(r.json(), normalise, kwargs)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _do_post_query(self, endpoint, **kwargs):
        """
        Perform a GET query against the endpoint

        Parameters
        ----------
        endpoint : str
            A join of the base  URL and the specific endpoint
        **kwargs
            Arguments specific to the query

        Returns
        -------
        data : list or dict or :obj:`pandas.DataFrame`
            The rest query data
        """
        # pop out the non post kwargs, we will default it at 100, whatever
        # is left becomes post data
        max_post = kwargs.pop('maximum_post_size', 100)
        max_slice_len = kwargs.pop('slice_length', 1E6)
        # TODO: perform post/slice length/size checks

        kwargs, normalise = self._set_kwargs(kwargs)
        req_kwargs = dict(headers={"Content-Type": "application/json"})

        if len(kwargs) > 0:
            req_kwargs['data'] = json.dumps(kwargs)

        r = requests.post(endpoint, **req_kwargs)

        if not r.ok:
            r.raise_for_status()

        return self._return_data(r.json(), normalise, kwargs)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _set_kwargs(self, kwargs):
        """
        Remove any kwarg that is NoneType from the kwargs, as these will not
        have been defined by the user. TODO: Do I need to make ints/floats into
        strings or will requests do that? I probably need to convert bools to
        0/1 scale though

        Parameters
        ----------
        kwargs : dict
            The keyword arguments to process

        Returns
        -------
        kwargs : dict
            The processed kwargs, the return is not strictly necessary as these
            happen in place
        normalise : func
            A function to normalise the results data if pandas is True
        """
        # Extract the normalise function and fall back to the default pandas
        # one
        normalise = kwargs.pop('norm_func', default_pandas)

        keys = list(kwargs.keys())
        for k in keys:
            if kwargs[k] is None:
                del kwargs[k]

        return kwargs, normalise

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _return_data(self, data, normalise, kwargs):
        """
        Format the return data based on what the user wants

        Parameters
        ----------
        data : list or dict
            The data returned from the REST query to format
        normalise : func
            A normalisation function that is applied if pandas is True

        Returns
        -------
        data : list or dict or :obj:`pandas.DataFrame`
            The rest query data
        """
        if self.pandas is True:
            return normalise(data, **kwargs)
        return data


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def default_pandas(data, **kwargs):
    """
    A default pandas normaliser, this does not try to handle any nesting
    and just drops the data in a `pandas.DataFrame`

    Parameters
    ----------
    data : list or dict
        The data to convert into a `pandas.DataFrame`
    **kwargs
        Any keyword arguments passed to the query that the normalise function
        will need to know about to do it's job

    Returns
    -------
    pandas_data : :obj:`pandas.DataFrame`
        The data shoehorned into a `pandas.DataFrame`
    """
    if isinstance(data, dict):
        data = [data]
    return pd.DataFrame(data)
