# Copyright (c) 2017 Ansible Project
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

DOCUMENTATION = '''
    name: gcp_compute
    plugin_type: inventory
    short_description: Google Cloud Compute Engine inventory source
    requirements:
        - requests >= 2.18.4
        - google-auth >= 1.3.0
    extends_documentation_fragment:
        - constructed
        - inventory_cache
    description:
        - Get inventory hosts from Google Cloud Platform GCE.
        - Uses a YAML configuration file that ends with gcp_compute.(yml|yaml) or gcp.(yml|yaml).
    options:
        plugin:
            description: token that ensures this is a source file for the 'gcp_compute' plugin.
            required: True
            choices: ['gcp_compute']
        zones:
          description: A list of regions in which to describe GCE instances.
          default: all zones available to a given project
        projects:
          description: A list of projects in which to describe GCE instances.
        filters:
          description: >
            A list of filter value pairs. Available filters are listed here
            U(https://cloud.google.com/compute/docs/reference/rest/v1/instances/list).
            Each additional filter in the list will act be added as an AND condition
            (filter1 and filter2)
        hostnames:
          description: A list of options that describe the ordering for which
              hostnames should be assigned. Currently supported hostnames are
              'public_ip', 'private_ip', or 'name'.
          default: ['public_ip', 'private_ip', 'name']
        auth_kind:
            description:
                - The type of credential used.
        service_account_file:
            description:
                - The path of a Service Account JSON file if serviceaccount is selected as type.
        service_account_email:
            description:
                - An optional service account email address if machineaccount is selected
                  and the user does not wish to use the default email.
        vars_prefix:
            description: prefix to apply to host variables, does not include facts nor params
            default: ''
        use_contrib_script_compatible_sanitization:
          description:
            - By default this plugin is using a general group name sanitization to create safe and usable group names for use in Ansible.
              This option allows you to override that, in efforts to allow migration from the old inventory script.
            - For this to work you should also turn off the TRANSFORM_INVALID_GROUP_CHARS setting,
              otherwise the core engine will just use the standard sanitization on top.
            - This is not the default as such names break certain functionality as not all characters are valid Python identifiers
              which group names end up being used as.
          type: bool
          default: False
          version_added: '2.8'
'''

EXAMPLES = '''
plugin: gcp_compute
zones: # populate inventory with instances in these regions
  - us-east1-a
projects:
  - gcp-prod-gke-100
  - gcp-cicd-101
filters:
  - machineType = n1-standard-1
  - scheduling.automaticRestart = true AND machineType = n1-standard-1
scopes:
  - https://www.googleapis.com/auth/compute
service_account_file: /tmp/service_account.json
auth_kind: serviceaccount
keyed_groups:
  # Create groups from GCE labels
  - prefix: gcp
    key: labels
hostnames:
  # List host by name instead of the default public ip
  - name
compose:
  # Set an inventory parameter to use the Public IP address to connect to the host
  # For Private ip use "networkInterfaces[0].networkIP"
  ansible_host: networkInterfaces[0].accessConfigs[0].natIP
'''

from ansible.errors import AnsibleError, AnsibleParserError
from ansible.module_utils._text import to_native
from ansible.module_utils.gcp_utils import GcpSession, navigate_hash, GcpRequestException
from ansible.plugins.inventory import BaseInventoryPlugin, Constructable, Cacheable
import json


# The mappings give an array of keys to get from the filter name to the value
# returned by boto3's GCE describe_instances method.
class GcpMockModule(object):
    def __init__(self, params):
        self.params = params

    def fail_json(self, *args, **kwargs):
        raise AnsibleError(kwargs['msg'])


class InventoryModule(BaseInventoryPlugin, Constructable, Cacheable):

    NAME = 'gcp_compute'

    def __init__(self):
        super(InventoryModule, self).__init__()

        self.group_prefix = 'gcp_'

    def _populate_host(self, item):
        '''
            :param item: A GCP instance
        '''
        hostname = self._get_hostname(item)
        self.inventory.add_host(hostname)
        for key in item:
            try:
                self.inventory.set_variable(hostname, self.get_option('vars_prefix') + key, item[key])
            except (ValueError, TypeError) as e:
                self.display.warning("Could not set host info hostvar for %s, skipping %s: %s" % (hostname, key, to_native(e)))
        self.inventory.add_child('all', hostname)

    def verify_file(self, path):
        '''
            :param path: the path to the inventory config file
            :return the contents of the config file
        '''
        if super(InventoryModule, self).verify_file(path):
            if path.endswith(('gcp.yml', 'gcp.yaml')):
                return True
            elif path.endswith(('gcp_compute.yml', 'gcp_compute.yaml')):
                return True
        return False

    def self_link(self, params):
        '''
            :param params: a dict containing all of the fields relevant to build URL
            :return the formatted URL as a string.
        '''
        return "https://www.googleapis.com/compute/v1/projects/{project}/zones/{zone}/instances".format(**params)

    def fetch_list(self, params, link, query):
        '''
            :param params: a dict containing all of the fields relevant to build URL
            :param link: a formatted URL
            :param query: a formatted query string
            :return the JSON response containing a list of instances.
        '''
        module = GcpMockModule(params)
        auth = GcpSession(module, 'compute')
        response = auth.get(link, params={'filter': query})
        return self._return_if_object(module, response)

    def _get_zones(self, config_data):
        '''
            :param config_data: dict of info from inventory file
            :return an array of zones that this project has access to
        '''
        link = "https://www.googleapis.com/compute/v1/projects/{project}/zones".format(**config_data)
        zones = []
        zones_response = self.fetch_list(config_data, link, '')
        for item in zones_response['items']:
            zones.append(item['name'])
        return zones

    def _get_query_options(self, filters):
        '''
            :param config_data: contents of the inventory config file
            :return A fully built query string
        '''
        if not filters:
            return ''

        if len(filters) == 1:
            return filters[0]
        else:
            queries = []
            for f in filters:
                # For multiple queries, all queries should have ()
                if f[0] != '(' and f[-1] != ')':
                    queries.append("(%s)" % ''.join(f))
                else:
                    queries.append(f)

            return ' '.join(queries)

    def _return_if_object(self, module, response):
        '''
            :param module: A GcpModule
            :param response: A Requests response object
            :return JSON response
        '''
        # If not found, return nothing.
        if response.status_code == 404:
            return None

        # If no content, return nothing.
        if response.status_code == 204:
            return None

        try:
            response.raise_for_status
            result = response.json()
        except getattr(json.decoder, 'JSONDecodeError', ValueError) as inst:
            module.fail_json(msg="Invalid JSON response with error: %s" % inst)
        except GcpRequestException as inst:
            module.fail_json(msg="Network error: %s" % inst)

        if navigate_hash(result, ['error', 'errors']):
            module.fail_json(msg=navigate_hash(result, ['error', 'errors']))
        if result['kind'] != 'compute#instanceList' and result['kind'] != 'compute#zoneList':
            module.fail_json(msg="Incorrect result: {kind}".format(**result))

        return result

    def _format_items(self, items):
        '''
            :param items: A list of hosts
        '''
        for host in items:
            if 'zone' in host:
                host['zone_selflink'] = host['zone']
                host['zone'] = host['zone'].split('/')[-1]
            if 'machineType' in host:
                host['machineType_selflink'] = host['machineType']
                host['machineType'] = host['machineType'].split('/')[-1]

            if 'networkInterfaces' in host:
                for network in host['networkInterfaces']:
                    if 'network' in network:
                        network['network'] = self._format_network_info(network['network'])
                    if 'subnetwork' in network:
                        network['subnetwork'] = self._format_network_info(network['subnetwork'])

            host['project'] = host['selfLink'].split('/')[6]
        return items

    def _add_hosts(self, items, config_data, format_items=True):
        '''
            :param items: A list of hosts
            :param config_data: configuration data
            :param format_items: format items or not
        '''
        if not items:
            return
        if format_items:
            items = self._format_items(items)

        for host in items:
            self._populate_host(host)

            hostname = self._get_hostname(host)
            self._set_composite_vars(self.get_option('compose'), host, hostname)
            self._add_host_to_composed_groups(self.get_option('groups'), host, hostname)
            self._add_host_to_keyed_groups(self.get_option('keyed_groups'), host, hostname)

    def _format_network_info(self, address):
        '''
            :param address: A GCP network address
            :return a dict with network shortname and region
        '''
        split = address.split('/')
        region = ''
        if 'global' in split:
            region = 'global'
        else:
            region = split[8]
        return {
            'region': region,
            'name': split[-1],
            'selfLink': address
        }

    def _get_hostname(self, item):
        '''
            :param item: A host response from GCP
            :return the hostname of this instance
        '''
        hostname_ordering = ['public_ip', 'private_ip', 'name']
        if self.get_option('hostnames'):
            hostname_ordering = self.get_option('hostnames')

        for order in hostname_ordering:
            name = None
            if order == 'public_ip':
                name = self._get_publicip(item)
            elif order == 'private_ip':
                name = self._get_privateip(item)
            elif order == 'name':
                name = item[u'name']
            else:
                raise AnsibleParserError("%s is not a valid hostname precedent" % order)

            if name:
                return name

        raise AnsibleParserError("No valid name found for host")

    def _get_publicip(self, item):
        '''
            :param item: A host response from GCP
            :return the publicIP of this instance or None
        '''
        # Get public IP if exists
        for interface in item['networkInterfaces']:
            if 'accessConfigs' in interface:
                for accessConfig in interface['accessConfigs']:
                    if 'natIP' in accessConfig:
                        return accessConfig[u'natIP']
        return None

    def _get_privateip(self, item):
        '''
            :param item: A host response from GCP
            :return the privateIP of this instance or None
        '''
        # Fallback: Get private IP
        for interface in item[u'networkInterfaces']:
            if 'networkIP' in interface:
                return interface[u'networkIP']

    def parse(self, inventory, loader, path, cache=True):
        super(InventoryModule, self).parse(inventory, loader, path)

        config_data = {}
        config_data = self._read_config_data(path)

        if self.get_option('use_contrib_script_compatible_sanitization'):
            self._sanitize_group_name = self._legacy_script_compatible_group_sanitization

        # get user specifications
        if 'zones' in config_data:
            if not isinstance(config_data['zones'], list):
                raise AnsibleParserError("Zones must be a list in GCP inventory YAML files")

        # get user specifications
        if 'projects' not in config_data:
            raise AnsibleParserError("Projects must be included in inventory YAML file")

        if not isinstance(config_data['projects'], list):
            raise AnsibleParserError("Projects must be a list in GCP inventory YAML files")

        # add in documented defaults
        if 'filters' not in config_data:
            config_data['filters'] = None

        projects = config_data['projects']
        zones = config_data.get('zones')
        config_data['scopes'] = ['https://www.googleapis.com/auth/compute']

        query = self._get_query_options(config_data['filters'])

        # Cache logic
        if cache:
            cache = self.get_option('cache')
            cache_key = self.get_cache_key(path)
        else:
            cache_key = None

        cache_needs_update = False
        if cache:
            try:
                results = self.cache.get(cache_key)
                for project in results:
                    for zone in results[project]:
                        self._add_hosts(results[project][zone], config_data, False)
            except KeyError:
                cache_needs_update = True

        if not cache or cache_needs_update:
            cached_data = {}
            for project in projects:
                cached_data[project] = {}
                config_data['project'] = project
                if not zones:
                    zones = self._get_zones(config_data)
                for zone in zones:
                    config_data['zone'] = zone
                    link = self.self_link(config_data)
                    resp = self.fetch_list(config_data, link, query)
                    self._add_hosts(resp.get('items'), config_data)
                    cached_data[project][zone] = resp.get('items')

        if cache_needs_update:
            self.cache.set(cache_key, cached_data)

    @staticmethod
    def _legacy_script_compatible_group_sanitization(name):

        return name