#!/usr/bin/env python
"""
Common functions that are used by many of the scripts
"""
from selenium.webdriver.chrome.options import Options
from selenium import webdriver
from collections import namedtuple
import inflection
import os
import re
import time
import pickle
import textwrap
import csv
import pprint as pp

MAX_WIDTH = 79
INDENT_CHARS = "    "

EndpointLink = namedtuple(
    'EndpointLink',
    ['name', 'key', 'endpoint', 'url', 'description']
)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def initialise_chrome(driver_path=None):
    """
    Initialise a chrome webdriver to download the files. Note that I can't make
    this headless at the moment as downloading does not seem to work in
    headless mode

    driver :str
        The path to the chrome webdriver (if it is not in your path)
    """

    chrome_options = Options()
    chrome_options.add_argument("--disable-extensions")
    chrome_options.add_argument("--disable-gpu")

    kwargs = {'options': chrome_options}
    if driver_path is not None:
        kwargs['executable_path'] = driver_path

    driver = webdriver.Chrome(**kwargs)

    return driver


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_endpoint_links(driver):
    """
    Get the links to all the endpoint webpages

    Parameters
    ----------
    driver : :obj:`selenium.Webdriver`
        The selenium webdriver

    Returns
    -------
    all_links : dict of list
        A each element of the list is an endpoint method
    """
    # Goto the rest landing page
    root_url = "http://rest.ensembl.org"
    driver.get(root_url)

    # All the links are in a big table
    main_table = driver.find_element_by_xpath(
        '//table[@class="table table-striped"]'
    )

    # Get the headings and the links associated with them
    headings = main_table.find_elements_by_xpath(
        './thead/tr/td/h3'
    )
    links = main_table.find_elements_by_xpath(
        './tbody'
    )

    # There should be the same amount of links and headings
    if len(links) != len(headings):
        raise ValueError("unmatched links and headings")

    all_links = {}
    for heading, link in zip(headings, links):
        heading_key = norm_string(heading.text)
        all_links[heading_key] = []
        for tr in link.find_elements_by_xpath('./tr'):
            link_data = [heading.text, heading_key]
            for idx, td in enumerate(tr.find_elements_by_xpath('./td')):
                if idx == 0:
                    link_data.append(td.text)
                    a = td.find_element_by_xpath('./a')
                    link_data.append(a.get_attribute('href'))
                else:
                    link_data.append(td.text)
            all_links[heading_key].append(EndpointLink(*link_data))
    return all_links


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def norm_string(value):
    """
    Normalise a string by making it lower case and replacing the spaces with
    underscores

    Parameters
    ----------
    value : str
        The string to normalise

    Returns
    -------
    norm_str
        The normalised string
    """
    return re.sub(r'\s+', '_', value.lower())


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_table_after(element):
    """
    Get the "sibling" table located after the element

    Parameters
    ----------
    element : :obj:`selenium.WebElement`
        The key element located before the table

    Returns
    -------
    table_element :  : :obj:`selenium.WebElement`
        The table element located after element, this will have the tag_name
        table
    """
    curr_tag = element.text
    next_table = element.find_element_by_xpath(
        './following-sibling::table'
    )
    back_check = next_table.find_element_by_xpath(
        './preceding-sibling::{0}'.format(element.tag_name)
    )

    # Make sure the table is the table that we expect
    if back_check.text != curr_tag:
        raise KeyError(
            "current: {0}, back check:{1}".format(
                curr_tag, back_check.text)
        )
    return next_table


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_resources(element):
    """
    Scrape the resources table. This has general information like return
    methods, query type GET/POST and max post size

    Parameters
    ----------
    element : :obj:`selenium.WebElement`
        The key element located before the table

    Returns
    -------
    resources : dict
        The keys are the resource name and the values are the resource value
    """
    resources = {}
    for tr in element.find_elements_by_xpath('./tbody/tr'):
        # print(tr)
        tds = tr.find_elements_by_xpath('./td')
        resources[norm_string(tds[0].text)] = tds[1].text

    try:
        resources['response_formats'] = \
            resources['response_formats'].split("\n")
    except KeyError:
        pass

    return resources


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_params(element):
    """
    Scrape a parameter section, these will be either an optional or required
    parameters

    Parameters
    ----------
    element : :obj:`selenium.WebElement`
        The key element located before the table

    Returns
    -------
    params : dict
        The keys name, type, desc, default, example_values
    """
    params = []
    for tr in element.find_elements_by_xpath('./tbody/tr'):
        param = {}
        tds = tr.find_elements_by_xpath('./td')

        param['name'] = tds[0].text
        param['type'] = tds[1].text
        param['desc'] = tds[2].text
        param['default'] = tds[3].text
        param['example_values'] = tds[4].text
        params.append(param)
    return params


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def parse_post_format(format_str):
    """
    """
    format_str = format_str.strip()

    if not re.match(r'^{.+?}$', format_str):
        if format_str == '':
            return {}

        raise TypeError(
            "unrecognised post format string: {0}".format(format_str)
        )
    elif re.match(r'^{\s*}$', format_str):
        return {}

    format_str = format_str.strip('{').strip('}')
    format_str = re.sub(r'[\s"]+', '', format_str)
    # format_str = re.split(r':,', format_str)

    format_str = re.sub(r'\[', '"ARRAY(', format_str)
    format_str = re.sub(r'\]', ')"', format_str)

    # I can't work out why this does not work??
    # quote = re.compile(
    #     r'''
    #     ,
    #     (?=
    #     (?:
    #     [^"]
    #     |
    #     "[^"]*"
    #     )*
    #     $)
    #     ''',
    #     re.MULTILINE
    # )
    # but this does??
    quote = re.compile(
        r''',(?=(?:[^'"`|]|'[^']*'|"[^"]*"|`[^`]*`|\|[^|]*\|)*$)'''
    )
    dict_delim = re.compile(
        r''':(?=(?:[^'"`|]|'[^']*'|"[^"]*"|`[^`]*`|\|[^|]*\|)*$)'''
    )

    format_str = quote.split(format_str)

    post_format = {}
    for i in format_str:
        key, value = dict_delim.split(i)
        # print('Key={0}'.format(key))
        # print('Value={0}'.format(value))
        if value == 'array':
            value = 'list'
        elif re.match(r'int(eger)?', value, re.IGNORECASE):
            value = 'int'
        elif value == 'long':
            value = 'int'
        elif value == 'float':
            value = 'float'
        elif value == 'string':
            value = 'str'
        elif value.startswith('"ARRAY('):
            value = 'list'
        else:
            for t in [int, float, str]:
                try:
                    value = type(t(value)).__name__
                    break
                except (ValueError, TypeError):
                    # raise
                    pass
        post_format[key] = value
    return post_format


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_post_message(element):
    """
    Scrape the resources table. This has general information like return
    methods, query type GET/POST and max post size

    Parameters
    ----------
    element : :obj:`selenium.WebElement`
        The key element located before the table

    Returns
    -------
    resources : dict
        The keys are the resource name and the values are the resource value
    """
    rows = element.find_elements_by_xpath('./tbody/tr')

    if len(rows) != 2:
        raise ValueError("unexpected row length: {0}".format(len(rows)))

    cols = rows[1].find_elements_by_xpath('./td')
    if len(cols) != 3:
        raise ValueError("unexpected col length: {0}".format(len(cols)))

    content_type = cols[0].text
    post_format = parse_post_format(cols[1].text)
    post_example = parse_post_format(cols[2].text)

    return {
        'content_type': content_type,
        'post_format': post_format,
        'post_example': post_example
    }


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_link_params(driver, link):
    """
    Scrape all the resource and parameter data for the link.

    Parameters
    ----------
    driver : :obj:`selenium.Webdriver`
        The selenium webdriver
    link : :obj:`ElementLink`
        The link to extract

    Returns
    -------
    all_params : dict
        The keys are resources, required, optional and optionally post_params
    """
    driver.get(link.url)

    # Get the resources
    resource = driver.find_element_by_xpath(
        '//h3[text() = "Resource Information"]'
    )
    resource_table = get_table_after(resource)
    resources = get_resources(resource_table)

    # Get the required parameters
    required = driver.find_element_by_xpath(
        '//h3[text() = "Required"]'
    )
    required_table = get_table_after(required)
    required_params = get_params(required_table)

    # Get the optional parameters
    optional = driver.find_element_by_xpath(
        '//h3[text() = "Optional"]'
    )
    optional_table = get_table_after(optional)
    optional_params = get_params(optional_table)

    time.sleep(2)
    # Get the required parameters
    try:
        message = driver.find_element_by_xpath(
            '//h2[text() = "Message"]'
        )
        message_table = get_table_after(message)
        message_params = get_post_message(message_table)

        return {
            'resources': resources,
            'required': required_params,
            'optional': optional_params,
            'post_options': message_params
        }
    except Exception:
        if resources['methods'] == 'POST':
            raise
        return {
            'resources': resources,
            'required': required_params,
            'optional': optional_params
        }


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def scrape_rest_attr(driver, links):
    """
    Parameters
    ----------
    driver : :obj:`selenium.Webdriver`
        The selenium webdriver
    links : dict of list
        Each element of the list is an `ElementLink`

    Returns
    -------
    rest_data : dict
        The keys are main rest endpoints and the values are dicts
    """
    rest_data = {}
    for key, values in links.items():
        rest_data[key] = []
        for i in values:
            time.sleep(1)
            params = get_link_params(driver, i)
            params['link'] = i
            rest_data[key].append(params)
        # break
    return rest_data


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def indent(level=1, value=""):
    """
    Indent a value by level

    Parameters
    ----------
    level : int, optional, default: 1
        The number of levels to indent
    value : str, optional, default: ""
        The value string to indent

    Returns
    -------
    indented_value : str
        The indented value
    """
    indent = INDENT_CHARS * level
    return "{0}{1}".format(indent, value)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def n_indent(level):
    """
    Return the number of indent characters for the required level

    Returns
    -------
    n_chars : int
        The number of indent spaces
    """
    return level * len(INDENT_CHARS)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_class_name(link_key):
    """
    Turn a link_key (rest "area") into a class name, this involved removing
    all underscores and making the first letter of all words uppercase

    Parameters
    ----------
    link_key : str
        The rest "area" that will correspond to a class name

    Returns
    -------
    class_name : str
        A camel-cased version of the link_key
    """
    return re.sub(r' ', '', re.sub(r'_', ' ', link_key).title())


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def write_docstring_quote(writer, level):
    """
    Write a quotation open/close for a docstring. The quote will have a
    carriage return after it

    Parameters
    ----------
    writer : :obj:`File`
        A file writer
    level : int
        The indent level
    """
    writer.write("{0}\n".format(indent(level, r'"""')))


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def write_class(writer, link_key):
    """
    Write the class definition line

    Parameters
    ----------
    writer : :obj:`File`
        A file writer
    link_key : str
        The rest "area" that will correspond to a class name
    """
    class_name = get_class_name(link_key)
    writer.write(
        "{0}\nclass {1}(_BaseRestClient):\n".format(
            "\n\n# {0}".format("@" * 77),
            class_name
        )
    )
    write_docstring_quote(writer, 1)
    writer.write(indent(1, "Perform {0} queries\n".format(class_name)))
    write_docstring_quote(writer, 1)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_method_name(link):
    """
    Using the link endpoint, create a method name, by removing all parameter
    locations, converting the / and spaces into underscores and making
    lowercase

    Parameters
    ----------
    link : :obj:`EndpointLink`
        The scraped endpoint link

    Returns
    -------
    method_name : str
        The method name derived from the REST endpoint
    """
    name = re.sub(r':\w+', '', link.endpoint)
    name = re.sub(r'[/\s]+', '_', name)
    print(name)
    name = re.sub(r'(\W|_)+$', '', name)
    print(name)
    return name.lower()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def write_separator_line(writer, level):
    """
    Write a method separator line:
    # ~~~~~~~~~~~~~~~~~~~~~~~~

    Parameters
    ----------
    writer : :obj:`File`
        A file writer
    level : int
        The indent level
    """
    lead_chars = '# '
    line_len = MAX_WIDTH - (len(lead_chars) + n_indent(level))
    sep_line = "\n{0}".format(
        indent(level, "{0}{1}\n".format(lead_chars, '~' * line_len))
    )
    writer.write(sep_line)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def sub_python_reserved(name):
    """
    If the name collides with a python function call then rename
    """
    if name == 'id':
        return 'query_id'
    elif name == 'format':
        return 'data_format'
    elif name == 'filter':
        return 'query_filter'
    elif name == 'type':
        return 'data_type'
    elif name == 'type':
        return 'data_type'
    elif name == 'class':
        return 'class_method'
    else:
        return name


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_arg_list(arg_list, add_self=False):
    """
    Build an argument or kw argument list for a method call

    Parameters
    ----------
    arg_list : list of dict
        All the arguments, these may be option or required arguments
    add_self : bool, optional, default: False
        Add a "self" argument to the start of the list

    Returns
    -------
    arg_list : list or str
        A list of the required/optional arguments to the method
    """
    args = []
    if add_self is True:
        args.append('self')

    for i in arg_list:
        # Some arguments collide with python core functions, so rename any that
        # do in a standard way
        i['code_name'] = sub_python_reserved(i['name'])

        # Now we want to make any camel-cased arguments more pythonic
        i['code_name'] = inflection.underscore(i['code_name'])

        # We are skipping callbacks at the moment as I am not sure how to use
        # them
        if i['code_name'] == 'callback':
            continue
        args.append(i['code_name'])
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def add_post_options(method_def):
    """
    Sometimes not all the post options are defined in the required/optional
    tables. This double checks them and makes sure that they are defined.
    If they are not defined they are placed in args  (this may change when I
    can test everything)

    Parameters
    ----------
    method_def : dict
        The method definition dictionary having the keys: resources, link,
        required, optional
    """
    # Merge the post options that have been extracted from the format and
    # examples sections, with the format overiding the examples
    post_options = {**method_def['post_options']['post_example'],
                    **method_def['post_options']['post_format']}

    for arg, dtype in post_options.items():
        # The default is that it has not been found
        found = False

        # Now first see if we already have it
        for i in method_def['required']+method_def['optional']:
            if i['name'] == arg:
                found = True
                break

        # If we can't find the post arg in any of the required/optionals
        # then we add it
        if found is False:
            method_def['required'].append(
                dict(
                    name=arg,
                    type=dtype,
                    default='-',
                    example_values='-',
                    desc='Post option not defined in required/optional args'
                )
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def write_method_def(writer, method_def):
    """
    Write the method definition line

    Parameters
    ----------
    writer : :obj:`File`
        A file writer
    method_def : dict
        The method definition dictionary having the keys: resources, link,
        required, optional
    """
    # pp.pprint(method_def)
    call_def = 'def'
    method_name = get_method_name(method_def['link'])
    print("M =", method_name)
    subsequent_indent = len(call_def) + len(method_name) + 2 + n_indent(1)

    tr = textwrap.TextWrapper(
        width=MAX_WIDTH,
        initial_indent=indent(1),
        subsequent_indent=(' ' * subsequent_indent)
    )

    args = build_arg_list(method_def['required'], add_self=True)
    kwargs = build_arg_list(method_def['optional'], add_self=False)

    arg_str = ", ".join(args)
    kwarg_str = ", ".join(["{0}=None".format(k) for k in kwargs])

    all_args = []
    if len(arg_str) > 1:
        all_args.append(arg_str)

    if len(kwarg_str) > 0:
        all_args.append(kwarg_str)

    wrapped = tr.wrap(
        "{0} {1}({2}):\n".format(
            call_def,
            method_name,
            ", ".join(all_args)
        )
    )
    writer.write("{0}\n".format("\n".join(wrapped)))


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_datatype(param_def):
    """
    Conver the REST "datatype" name into a python data type name

    Parameters
    ----------
    param_def : dict
        The information on the rest parameter

    Returns
    -------
    python_data_type_name : str
        The equivalent python data type name
    """
    rest_dt = param_def['type'].lower()
    allowed_values = ""

    if rest_dt == 'string':
        return 'str', allowed_values
    elif re.match(r'enum\((.+?)\)', rest_dt):
        values_regexp = re.match(r'enum\((.+?)\)', rest_dt)
        allowed_values = " Allowed values: {0}".format(
            values_regexp.group(1)
        )
        return 'str', allowed_values
    elif rest_dt.startswith('boolean'):
        return 'bool', allowed_values
    elif rest_dt == 'integer':
        return 'int', allowed_values
    elif rest_dt == '':
        return 'str', allowed_values
    else:
        return rest_dt, allowed_values


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_example_value_str(example):
    """
    Return a sensible example value string. If there is not one then return
    ""

    Parameters
    ----------
    example : str
        The scraped example value

    Returns
    -------
    example_str : str
        An example value string to place in the docstring
    """
    if example == '' or example == '-' or example == 'randomlygeneratedname':
        return ""
    else:
        return " Example values: '{0}'.".format(example)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_docstring_args(args, level, optional=False):
    """
    Generate a list of docstring arguments from the given arguments. This
    assumes that write_method_def has been called first.

    Parameters
    ----------
    args : list of dict
        The arguments to build the doctrings
    level : int
        The level the docstring will be written at
    optional : bool, optional, default: False
        Are the supplied arguments optional arguments, if so an optional
        string will be added to the doctring

    Returns
    -------
    docstring_args : list of str
        The docstring arguments
    """
    tr = textwrap.TextWrapper(
        width=MAX_WIDTH,
        initial_indent=indent(level+1),
        subsequent_indent=indent(level+1),
        break_long_words=False
    )

    optional_str = ""
    if optional is True:
        optional_str = " or NoneType, optional, default: NoneType"

    docstr_params = []
    for i in args:
        if i['code_name'] == 'callback':
            continue

        dtype, allowed_values = get_datatype(i)
        docstr_params.append(
            '{2}{0} : {1}{3}'.format(
                i['code_name'],
                dtype,
                indent(level),
                optional_str
            )
        )

        if not i['desc'].endswith('.'):
            i['desc'] = '{0}.'.format(i['desc'])

        # May need to adjust boolean example values here
        docstr_params.extend(
            tr.wrap(
                "{0}{1}{2}".format(
                    i['desc'],
                    get_example_value_str(i['example_values']),
                    allowed_values
                )
            )
        )
    return docstr_params


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def write_method_docstring(writer, method_def):
    """
    Write the method docstring line

    Parameters
    ----------
    writer : :obj:`File`
        A file writer
    method_def : dict
        The method definition dictionary having the keys: resources, link,
        required, optional
    """
    # pp.pprint(method_def)

    tr = textwrap.TextWrapper(
        width=MAX_WIDTH,
        initial_indent=indent(2),
        subsequent_indent=indent(2),
        break_long_words=False
    )

    docsting_quote = indent(2, '"""')
    desc = "{0}. {1} method. Queries [{2}]({3}).".format(
        re.sub(r'\.$', '', method_def['link'].description),
        method_def['resources']['methods'],
        re.sub(r'(GET|POST)\s+', '', method_def['link'].endpoint),
        method_def['link'].url
    )

    for msg, key in [('Max post size', 'maximum_post_size'),
                     ('Max slice length', 'slice_length')]:
        try:
            desc = '{0} {2}: {1}.'.format(
                desc,
                method_def['resources'][key],
                msg
            )
        except KeyError:
            pass
    desc = "\n".join(tr.wrap(desc))

    docstr_params = get_docstring_args(
        method_def['required'], 2, optional=False
    )
    docstr_params.extend(
        get_docstring_args(
            method_def['optional'], 2, optional=True
        )
    )

    if len(docstr_params) > 0:
        writer.write(
            "{0}\n{1}\n\n{2}Parameters\n{2}----------\n{3}\n{0}\n".format(
                docsting_quote,
                desc,
                indent(2),
                "\n".join(docstr_params)
            )
        )
    else:
        writer.write(
            "{0}\n{1}\n{0}\n".format(
                docsting_quote,
                desc
            )
        )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_endpoint(method_def):
    """
    """
    level = 2
    endpoint = re.sub(r'(GET|POST)\s*', '', method_def['link'].endpoint)

    params = []
    consumed = []
    for i in re.findall(r'(:\w+:?/|:\w+:?$)', endpoint):
        i = i.rstrip('/')
        var_name = i.strip(':')
        endpoint = re.sub(
            r'{0}'.format(i),
            '{{{0}}}'.format(var_name),
            endpoint
        )
        # Now find the code_name in the args
        found = False
        for arg in method_def['required']:
            if arg['name'] == var_name:
                consumed.append(arg['code_name'])
                params.append(
                    indent(
                        level+1,
                        value='{0}={1}'.format(var_name, arg['code_name'])
                    )
                )
                found = True
                break
        if found is False:
            return [indent(
                level,
                "raise NotImplementedError('this method call is broken')"
            )], consumed

    code = []
    if len(params) > 0:
        code.append(
            indent(level, value="endpoint = '{0}'.format(".format(endpoint))
        )
        code.append(',\n'.join(params))
        code.append(indent(level, value=")"))
    else:
        code.append(
            indent(level, value="endpoint = '{0}'".format(endpoint))
        )

    return code, consumed


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def write_code(writer, method_def):
    """
    Write the code for the REST method

    Parameters
    ----------
    writer : :obj:`File`
        A file writer
    method_def : dict
        The method definition dictionary having the keys: resources, link,
        required, optional
    """
    level = 2
    code = [indent(
        level,
        value="method_type = '{0}'".format(method_def['resources']['methods'])
        )]

    call_kwargs = [
        indent(level+1, 'method_type'),
        indent(level+1, 'endpoint')
    ]
    for var_name in ['maximum_post_size', 'slice_length']:
        try:
            value = method_def['resources'][var_name]
            call_kwargs.append(
                indent(level+1, '{0}={1}'.format(var_name, value))
                )
            code.append(
                indent(
                    level, value="{0} = {1}".format(var_name, value)
                )
            )
        except KeyError:
            pass

    endpoint, consumed = get_endpoint(method_def)
    code.extend(endpoint)
    code.append("")

    for i in method_def['required']+method_def['optional']:
        if i['code_name'] =='callback' or i['code_name'] in consumed:
            continue

        call_kwargs.append(
            indent(level+1, '{0}={1}'.format(i['name'], i['code_name']))
        )
    writer.write("\n".join(code))
    writer.write("\n")
    writer.write(
        "{0}return self.rest_query(\n{1}\n{0})\n".format(
            indent(level),
            ",\n".join(call_kwargs)
        )
    )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def write_method(writer, method_def):
    """
    Write a method definition to the module

    Parameters
    ----------
    writer : :obj:`File`
        A file writer
    method_def : dict
        The method definition dictionary having the keys: resources, link,
        required, optional
    """
    # We only output the method if it can respond in json format
    if 'json' not in method_def['resources']['response_formats']:
        return

    # The method separator line
    write_separator_line(writer, 1)

    # Now write the method definition
    write_method_def(writer, method_def)

    # Write the method docstring
    write_method_docstring(writer, method_def)
    write_code(writer, method_def)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
if __name__ == '__main__':
    driver = initialise_chrome(
        driver_path=os.path.join(
            os.environ['HOME'], 'Software/bin/chromedriver'
        ),
    )
    links = get_endpoint_links(driver)
    rest_data = scrape_rest_attr(driver, links)
    pickle.dump(rest_data, open('rest_params.p', 'wb'))
    rest_data = pickle.load(open('rest_params.p', 'rb'))
    # pp.pprint(rest_data)

    # with open('ensembl_rest_client.py', 'wt') as module_writer:
    #     for query_type, methods in rest_data.items():
    #         write_class(module_writer, query_type)
    #         print(query_type)
    #         for m in methods:
    #             try:
    #                 # Now we want to add in any post features that were not
    #                 # in the required/optional args
    #                 add_post_options(m)
    #             except KeyError:
    #                 # Not a post method
    #                 pass

    #             write_method(module_writer, m)

    pp.pprint(rest_data['information'])
