#!/usr/bin/python
# 6/16/2017 - remove PyQuery dependency
# 5/19/2016 - update to allow for authentication based on api-key, rather
# than username/pw
# See https://documentation.uts.nlm.nih.gov/rest/authentication.html for
# full explanation

import requests
import json

# import lxml.html as lh
from lxml.html import fromstring

URI = "https://utslogin.nlm.nih.gov"
QUERY_URI = "https://uts-ws.nlm.nih.gov"

# option 1 - username/pw authentication at /cas/v1/tickets
# auth_endpoint = "/cas/v1/tickets/"
# option 2 - api key authentication at /cas/v1/api-key
AUTH_ENDPOINT = "/cas/v1/api-key"

# Query end
QUERY_END = "NONE"


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class Authentication(object):
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    # def def __init__(self, args):
    def __init__(self, username=None, password=None, apikey=None):
        """
        Initialise

        Parameters
        ----------
        username :str
            The username
        password :str
            The password
        """
        # TODO: Make flexible by adding user/password or API key
        self.username = username
        self.password = password
        self.apikey = apikey

        self.use_api_key = True
        self.service = "http://umlsks.nlm.nih.gov"

        if (
            self.apikey is None
            and sum([i is not None for i in [self.username, self.password]])
            == 2
        ):
            self.use_api_key = False
        elif self.apikey is None:
            raise ValueError(
                "you must supply an API key or a username/password"
            )
        # Get the granting ticket
        self.gettgt()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def gettgt(self):
        """
        This gets a granting ticket that is available for 8 hours
        """
        params = {"username": self.username, "password": self.password}
        if self.use_api_key is True:
            params = {"apikey": self.apikey}

        h = {
            "Content-type": "application/x-www-form-urlencoded",
            "Accept": "text/plain",
            "User-Agent": "python",
        }
        r = requests.post(URI + AUTH_ENDPOINT, data=params, headers=h)
        response = fromstring(r.text)
        ## extract the entire URL needed from the HTML form (action attribute)
        # returned - looks similar to https://utslogin.nlm.nih.gov/cas/v1
        # /tickets/TGT-36471-aYqNLN2rFIJPXKzxwdTNC5ZT7z3B3cTAKfSc5ndHQcUxeaDOLN-cas
        # We make a POST call to this URL in the getst method
        self.tgt = response.xpath("//form/@action")[0]
        return self.tgt

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def getst(self):
        """
        Gets a service ticket, needed for every REST API call
        """
        params = {"service": self.service}
        h = {
            "Content-type": "application/x-www-form-urlencoded",
            "Accept": "text/plain",
            "User-Agent": "python",
        }
        r = requests.post(self.tgt, data=params, headers=h)
        st = r.text
        return st


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class UmlsRestQuery(Authentication):
    """
    A query class for the UMLS rest API
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, **kwargs):
        """ """
        self.page_size = kwargs.pop("page_size", 25)
        super().__init__(**kwargs)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def search(self, query_str, **kwargs):
        """
        Perform a general search, for kwargs see:
        https://documentation.uts.nlm.nih.gov/rest/search/

        Paramerters
        -----------
        **kwargs
            The search parameters

        Returns
        -------

        Examples
        --------

        """
        version = self._extract_version(kwargs)
        kwargs["string"] = query_str
        content_endpoint = "/rest/search/{0}".format(version)

        for i in self._do_query(content_endpoint, **kwargs):
            yield i

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _extract_version(self, kwargs):
        """
        Extract the UMLS version number from the keyword arguments

        Parameters
        ----------
        kwargs :dict
            The keyword arguments that may or may not contain 'version'. If it
            is not present then the version is set to 'current'

        Returns
        -------
        version :str
            The requested UMLS database version
        """
        return kwargs.pop("version", "current")

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _do_query(self, content_endpoint, **kwargs):
        """
        Perform the query and get the results
        """
        page_no = 0
        final_endpoint = "{0}{1}".format(QUERY_URI, content_endpoint)
        print(final_endpoint)
        while True:
            # Generate a new service ticket for each page if needed
            ticket = self.getst()
            page_no += 1
            query = kwargs.copy()
            query["ticket"] = ticket
            query["pageNumber"] = page_no
            query["pageSize"] = self.page_size

            # Check the response code
            r = requests.get(final_endpoint, params=query)
            r.raise_for_status()
            r.encoding = "utf-8"

            results = json.loads(r.text)
            # results = r.json()
            results = results["result"]

            # Last page
            # "result":{"classType":"searchResults","results":[{"ui":"NONE","name":"NO RESULTS"}]}}
            if results["results"][0]["ui"] == QUERY_END:
                break

            yield results["results"]
            # for page in results["results"]:
            #     yield page
