DOCUMENTATION = """
    name: file_inventory
    version_added: 2.9
    short_description: File invenotry in YAML format
    description:
        - Uses YAML files with hosts to construct an inventory.
        - Hosts are defined once.
        - Host dictionaries are in a list under the key "hosts"
        - Virtual host dictionaries are in a list under the key "virtual_hosts"
        - Platform dictonaries are in a list under the key "platforms"
        - Group dictonaries are in a list under the key "group"
        - Files for hosts should be per application/system and named "hosts_<SYSTEM>.yaml"
        - Files for virtual hosts should be per application/system and named "virtual_hosts_<SYSTEM>.yaml"
        - Files for groups should be named "groups.yaml"
        - Files for platforms should be named "platforms.yaml"
    options:
        plugin:
            description: Name of the plugin
            required: true
            choices: ['file_inventory']
        inventory_directory:
            description: Directory with the YAML files relative to the playbook
            default: file_inventory
        environments:
            required: true
            description: List with which environments to get inventory for
        mandatory_host_keys:
            description: Mandatory keys for host object
            default:
                - name
                - platform
                - data_center
                - environment
                - system
                - interfaces
        mandatory_virtual_host_keys:
            description: Mandatory keys for virtual host object
            default:
                - url
                - type
                - data_center
                - environment
                - system
                - host
                - config
                - certificate_common_name
                - certificate_san
                - certificate_provider
                - certificate_file
                - certificate_key_file
                - ca_certificate_file
        extra_group_keys:
            description: Keys to use to set as groups
            default:
                - data_center
                - environment
                - system
"""

EXAMPLES = r"""
    # file_inventory.yaml, configuration file in YAML
    plugin: file_inventory
    environments:
      - dev
"""

import glob

from ansible.plugins.inventory import BaseInventoryPlugin
from ansible.errors import AnsibleError, AnsibleParserError
from ansible.parsing.utils.yaml import from_yaml


class InventoryModule(BaseInventoryPlugin):
    NAME = "file_inventory"

    def verify_file(self, path):
        """return true/false if this is a valid file for this plugin"""
        valid = False
        if super(InventoryModule, self).verify_file(path):
            if path.endswith((".yaml")):
                valid = True
        return valid

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

        self._read_config_data(path)
        try:
            self.plugin = self.get_option("plugin")
            self.inventory_directory = self.get_option("inventory_directory")
            self.environments = self.get_option("environments")
            self.mandatory_host_keys = self.get_option("mandatory_host_keys")
            self.mandatory_virtual_host_keys = self.get_option(
                "mandatory_virtual_host_keys"
            )
            self.extra_group_keys = self.get_option("extra_group_keys")

        except Exception as e:
            raise AnsibleParserError(f"All correct options required: {e}")

        self.structured_inventory = self._get_structured_inventory()
        self._populate()

    def _get_structured_inventory(self):
        inventory_files = glob.glob(f"./{self.inventory_directory}/*.yaml")

        inventory_data = {}
        for inv_file in inventory_files:
            with open(inv_file) as f:
                try:
                    inv_dict = from_yaml(f, vault_secrets=self.loader._vault.secrets)
                except Exception as e:
                    raise AnsibleParserError(
                        f"Could not load file as YAML: {inv_file} {e}"
                    )

                for k, v in inv_dict.items():
                    if k not in inventory_data:
                        inventory_data[k] = []
                    for i in v:
                        inventory_data[k].append(i)

        # verify that mandatory keys are set for hosts
        for host in inventory_data["hosts"]:
            for k in self.mandatory_host_keys:
                if k not in host or not host[k]:
                    e = f"ERROR: {host['name']}: Key not set or empty: {k}"
                    raise AnsibleError(e)

        # verify that mandatory keys are set for virtual hosts
        for vhost in inventory_data.get("virtual_hosts", []):
            for k in self.mandatory_virtual_host_keys:
                if k not in vhost or not vhost[k]:
                    e = f"ERROR: {vhost['url']}: Key not set or empty: {k}"
                    raise AnsibleError(e)

        return inventory_data

    def _populate(self):
        """Return the hosts and groups"""
        for host in self.structured_inventory["hosts"]:
            if host["environment"] in self.environments:
                hostname = host["name"]
                platform = host["platform"]
                platform_info = self._get_platform_info(platform)
                if not platform_info:
                    e = f"ERROR: {hostname}: Platform not found: {platform}"
                    raise AnsibleError(e)

                platformgroups = [
                    platform_info["groups"]["platform"],
                    platform_info["groups"]["platform_family"],
                    platform_info["groups"]["platform_dist"],
                ]
                platformgroups = ["platform_" + x for x in platformgroups]
                extra_hgs = [f"{x}_{host[x]}" for x in self.extra_group_keys]
                hgs = host["groups"] + platformgroups + extra_hgs
                for group in hgs:
                    self.inventory.add_group(group)
                    self.inventory.add_host(host=host["name"], group=group)

                # if it is a Zabbix Server get Zabbix configuration
                if ("zabbix_server" in host["groups"]) or (
                    "zabbix_infra" in host["groups"]
                ):
                    zhosts = self.get_zabbix_configuration_hosts()
                    zvhosts = self._get_zabbix_configuration_virtual_hosts()
                    zhosts += zvhosts
                    self.inventory.set_variable(
                        hostname, "zabbix_configuration_hosts", zhosts
                    )

                    zhgs = []
                    for zh in zhosts:
                        zhgs.extend([g for g in zh["host_groups"] if g not in zhgs])

                    # create groups for the Zabbix sub groups tree
                    for group in zhgs:
                        if group.find("/") > 0:
                            subgroups = group.split("/")
                            for x in range(1, len(subgroups)):
                                newgroup = "/".join(subgroups[:x])
                                if newgroup not in zhgs:
                                    zhgs.append(newgroup)

                    self.inventory.set_variable(
                        hostname, "zabbix_configuration_host_groups", zhgs
                    )

    def _get_platform_info(self, name):
        data = {}
        for _os in self.structured_inventory["platforms"]:
            for family in _os["families"]:
                for dist in family["distributions"]:
                    if dist["name"] == name:
                        data["groups"] = {
                            "platform": _os["name"],
                            "platform_family": family["name"],
                            "platform_dist": dist["name"],
                        }

                        templ = [
                            _os.get("zabbix_templates"),
                            family.get("zabbix_templates"),
                            dist.get("zabbix_templates"),
                        ]
                        data["zabbix_templates"] = [
                            item for subl in templ if subl for item in subl
                        ]

                        break
                else:
                    continue
                break
            else:
                continue
            break
        return data

    def get_zabbix_configuration_hosts(self):
        data = []
        for host in self.structured_inventory["hosts"]:
            # check if host should be monitored by Zabbix
            # default is that it should be monitored but can be
            # overriden by setting "zabbix_monitoring" to False
            zabbix_monitoring = host.get("zabbix_monitoring", True)
            hostname = host["name"]
            zhost = {"name": hostname}

            if not zabbix_monitoring:
                zhost = {
                    "name": host["name"],
                    "state": "absent",
                    "host_groups": [],
                    "templates": [],
                    "interfaces": [],
                }
                data.append(zhost)

            elif zabbix_monitoring:
                # set host groups from the group tags
                zhost["host_groups"] = list(host.get("groups", []))

                # set the host templates from the host groups
                zhost["templates"] = []
                for zgroup in zhost["host_groups"]:
                    for inv_group in self.structured_inventory["groups"]:
                        if inv_group["name"] == zgroup:
                            zhost["templates"].extend(
                                inv_group.get("zabbix_templates", [])
                            )
                            break

                # set os groups and templates
                platform = host["platform"]
                platform_info = self._get_platform_info(platform)
                if not platform_info:
                    e = f"ERROR: {hostname}: Platform not found: {platform}"
                    raise AnsibleError(e)

                # set Zabbix sub groups and templates for platform
                platform_group = "platform/{0}/{1}/{2}".format(
                    platform_info["groups"]["platform"],
                    platform_info["groups"]["platform_family"],
                    platform_info["groups"]["platform_dist"],
                )
                zhost["host_groups"].append(platform_group)
                zhost["templates"].extend(list(platform_info["zabbix_templates"]))

                # set Zabbix sub groups for system
                system_env_group = "system/{0}/{1}".format(
                    host["system"], host["environment"]
                )
                zhost["host_groups"].append(system_env_group)

                # set Zabbix sub groups for environment
                if host["environment"] == "prod":
                    env_group = "environment/prod"
                else:
                    env_group = "environment/noprod/{0}".format(host["environment"])
                zhost["host_groups"].append(env_group)

                # set the interfaces
                zhost["interfaces"] = []
                z_int_types = {
                    "agent": {"num": "1", "port": "10050"},
                    "snmp": {"num": "2", "port": "161"},
                    "ipmi": {"num": "3", "port": "0"},
                    "jmx": {"num": "4", "port": "0"},
                }

                host_ints = {}
                for interface in host["interfaces"]:
                    # get interface type or default to 'agent'
                    try:
                        int_type = interface.get("zabbix_type", "agent")
                    except Exception:
                        e = f"ERROR: {hostname}: Interface not correct: {interface}"
                        raise AnsibleError(e)

                    if int_type not in z_int_types:
                        e = f"ERROR: {hostname}: Zabbix interface type not found: {int_type}"
                        raise AnsibleError(e)

                    if int_type not in host_ints:
                        host_ints[int_type] = []

                    host_ints[int_type].append(dict(interface))
                #  host_ints[interface].append(dict(details))

                for int_type in host_ints:

                    int_type_num = z_int_types[int_type]["num"]
                    z_ints = []
                    for interface in host_ints[int_type]:
                        int_fqdn = interface.get("fqdn", hostname)
                        int_ip = interface.get("ip")
                        if not int_ip:
                            e = f"ERROR: {hostname}: Interface IP not specified: {interface}"
                            raise AnsibleError(e)
                        z_int_port = z_int_types[int_type].get("port")
                        int_port = interface.get("port", z_int_port)
                        if not int_port:
                            e = f"ERROR: {hostname}: Port not specified and no default: {interface['ip']}"
                            raise AnsibleError(e)

                        z_int = {}
                        z_int["dns"] = int_fqdn
                        z_int["ip"] = int_ip
                        z_int["port"] = int_port
                        z_int["type"] = int_type_num
                        z_int["main"] = "0"
                        z_int["useip"] = "1"
                        if int_type == "snmp":
                            details = interface.get("details")
                            if details:
                                z_int["details"] = dict(details)

                        z_ints.append(z_int)

                    # make first interface of the type the main interface
                    z_ints[0]["main"] = "1"

                    zhost["interfaces"].extend(z_ints)

                # set the Zabbix host macros
                macros = host.get("zabbix_macros")
                if macros:
                    zhost["macros"] = list(macros)

                # set the Zabbix host tls_psk_identity
                tls_psk_identity = host.get("tls_psk_identity")
                if tls_psk_identity:
                    zhost["tls_psk_identity"] = tls_psk_identity
                # set the Zabbix host tls_connect
                tls_connect = host.get("tls_connect")
                if tls_connect:
                    zhost["tls_connect"] = tls_connect
                # set the Zabbix host tls_connect
                tls_accept = host.get("tls_accept")
                if tls_accept:
                    zhost["tls_accept"] = tls_accept
                # set the Zabbix host tls_psk
                tls_psk = host.get("tls_psk")
                if tls_psk:
                    zhost["tls_psk"] = tls_psk
                # set the Zabbix host owner_country
                owner_country = host.get("owner_country")
                if owner_country:
                    zhost["owner_country"] = owner_country
                    # Assign a tag country=countryname for each host
                    # zhost["tags"]={}
                    # zhost["tags"]["tag"]="country"
                    # zhost["tags"]["value"]=owner_country

                # set the Zabbix host description
                # There is an space in " " because when not setting the
                # description the Ansible module for zabbix_host doesn't
                # blank the field if it is not set.
                zhost["description"] = host.get("description", " ")

                data.append(zhost)

        return data

    def _get_zabbix_configuration_virtual_hosts(self):
        data = []
        for host in self.structured_inventory.get("virtual_hosts", []):
            url = host["url"]
            hostname = url.split("/")[2]
            url_scheme = url.split(":")[0]

            # set port if defined otherwise use default ports
            def_ports = {"https": "443", "ftps": "21"}
            try:
                int_port = hostname.split(":")[1]
                hostname = hostname.split(":")[0]
            except IndexError:
                int_port = def_ports[url_scheme]
            zhost = {"name": hostname, "host_groups": [], "templates": []}

            # set groups and templates from virtual host type
            platform = host["type"]
            platform_info = self._get_platform_info(platform)
            if not platform_info:
                e = f"ERROR: {hostname}: Virtual host type not found in platforms.yaml: {platform}"
                raise AnsibleError(e)

            # set Zabbix sub groups and templates for platform
            platform_group = "platform/{0}/{1}/{2}".format(
                platform_info["groups"]["platform"],
                platform_info["groups"]["platform_family"],
                platform_info["groups"]["platform_dist"],
            )
            zhost["host_groups"].append(platform_group)
            zhost["templates"].extend(list(platform_info["zabbix_templates"]))

            # set Zabbix sub groups for system
            system_env_group = "system/{0}/{1}".format(
                host["system"], host["environment"]
            )
            zhost["host_groups"].append(system_env_group)

            # set Zabbix sub groups for environment
            if host["environment"] == "prod":
                env_group = "environment/prod"
            else:
                env_group = "environment/noprod/{0}".format(host["environment"])
            zhost["host_groups"].append(env_group)

            z_int = {}
            z_int["dns"] = hostname
            z_int["ip"] = ""
            z_int["port"] = int_port
            z_int["type"] = "1"
            z_int["main"] = "1"
            z_int["useip"] = 0

            zhost["interfaces"] = [z_int]

            # set the Zabbix host macros
            macros = []
            if url_scheme == "ftps":
                macros.append(
                    {"macro": "{$TLS_CERTIFICATES_INFO_STARTTLS}", "value": "ftp"},
                )
            if not int_port == "443":
                macros.append(
                    {"macro": "{$TLS_CERTIFICATES_INFO_PORT}", "value": int_port}
                )

            zhost["macros"] = list(macros)

            # set the Zabbix host description for virtual hosts
            try:
                vhost_info = {
                    "url": url,
                    "type": platform,
                    "host": host["host"],
                    "config": host["config"],
                    "certificate_common_name": host["certificate_common_name"],
                    "certificate_san": host["certificate_san"],
                    "certificate_provider": host["certificate_provider"],
                    "certificate_file": host["certificate_file"],
                    "certificate_key_file": host["certificate_key_file"],
                    "ca_certificate_file": host["certificate_file"],
                }

                description = ""
                for k, v in vhost_info.items():
                    description += f"{k}: {v}\n"
                    zhost["description"] = description.rstrip()

            except Exception as e:
                raise AnsibleParserError(f"All correct fields required: {e}")

            data.append(zhost)

        return data
