# Copyright (c) 2021, Victor Martinez <VictorMartinezRubio@gmail.com>
# GNU General Public License v3.0+ (see LICENSES/GPL-3.0-or-later.txt or https://www.gnu.org/licenses/gpl-3.0.txt)
# SPDX-License-Identifier: GPL-3.0-or-later

from __future__ import annotations

DOCUMENTATION = r"""
author: Victor Martinez (@v1v)  <VictorMartinezRubio@gmail.com>
name: elastic
type: notification
short_description: Create distributed traces for each Ansible task in Elastic APM
version_added: 3.8.0
description:
  - This callback creates distributed traces for each Ansible task in Elastic APM.
  - You can configure the plugin with environment variables.
  - See U(https://www.elastic.co/guide/en/apm/agent/python/current/configuration.html).
options:
  hide_task_arguments:
    default: false
    type: bool
    description:
      - Hide the arguments for a task.
    env:
      - name: ANSIBLE_OPENTELEMETRY_HIDE_TASK_ARGUMENTS
  apm_service_name:
    default: ansible
    type: str
    description:
      - The service name resource attribute.
    env:
      - name: ELASTIC_APM_SERVICE_NAME
  apm_server_url:
    type: str
    description:
      - Use the APM server and its environment variables.
    env:
      - name: ELASTIC_APM_SERVER_URL
  apm_secret_token:
    type: str
    description:
      - Use the APM server token.
    env:
      - name: ELASTIC_APM_SECRET_TOKEN
  apm_api_key:
    type: str
    description:
      - Use the APM API key.
    env:
      - name: ELASTIC_APM_API_KEY
  apm_verify_server_cert:
    default: true
    type: bool
    description:
      - Verifies the SSL certificate if an HTTPS connection.
    env:
      - name: ELASTIC_APM_VERIFY_SERVER_CERT
  traceparent:
    type: str
    description:
      - The L(W3C Trace Context header traceparent,https://www.w3.org/TR/trace-context-1/#traceparent-header).
    env:
      - name: TRACEPARENT
requirements:
  - elastic-apm (Python library)
"""


EXAMPLES = r"""
examples: |-
  Enable the plugin in ansible.cfg:
    [defaults]
    callbacks_enabled = community.general.elastic

  Set the environment variable:
    export ELASTIC_APM_SERVER_URL=<your APM server URL)>
    export ELASTIC_APM_SERVICE_NAME=your_service_name
    export ELASTIC_APM_API_KEY=your_APM_API_KEY
"""

import getpass
import socket
import time
import uuid
from collections import OrderedDict
from contextlib import closing
from os.path import basename

from ansible.errors import AnsibleError, AnsibleRuntimeError
from ansible.module_utils.ansible_release import __version__ as ansible_version
from ansible.plugins.callback import CallbackBase

ELASTIC_LIBRARY_IMPORT_ERROR: ImportError | None
try:
    from elasticapm import Client, capture_span, instrument, label, trace_parent_from_string
except ImportError as imp_exc:
    ELASTIC_LIBRARY_IMPORT_ERROR = imp_exc
else:
    ELASTIC_LIBRARY_IMPORT_ERROR = None


class TaskData:
    """
    Data about an individual task.
    """

    def __init__(self, uuid, name, path, play, action, args):
        self.uuid = uuid
        self.name = name
        self.path = path
        self.play = play
        self.host_data = OrderedDict()
        self.start = time.time()
        self.action = action
        self.args = args

    def add_host(self, host):
        if host.uuid in self.host_data:
            if host.status == "included":
                # concatenate task include output from multiple items
                host.result = f"{self.host_data[host.uuid].result}\n{host.result}"
            else:
                return

        self.host_data[host.uuid] = host


class HostData:
    """
    Data about an individual host.
    """

    def __init__(self, uuid, name, status, result):
        self.uuid = uuid
        self.name = name
        self.status = status
        self.result = result
        self.finish = time.time()


class ElasticSource:
    def __init__(self, display):
        self.ansible_playbook = ""
        self.session = str(uuid.uuid4())
        self.host = socket.gethostname()
        try:
            self.ip_address = socket.gethostbyname(socket.gethostname())
        except Exception:
            self.ip_address = None
        self.user = getpass.getuser()

        self._display = display

    def start_task(self, tasks_data, hide_task_arguments, play_name, task):
        """record the start of a task for one or more hosts"""

        uuid = task._uuid

        if uuid in tasks_data:
            return

        name = task.get_name().strip()
        path = task.get_path()
        action = task.action
        args = None

        if not task.no_log and not hide_task_arguments:
            args = ", ".join((f"{k}={v}" for k, v in task.args.items()))

        tasks_data[uuid] = TaskData(uuid, name, path, play_name, action, args)

    def finish_task(self, tasks_data, status, result):
        """record the results of a task for a single host"""

        task_uuid = result._task._uuid

        if hasattr(result, "_host") and result._host is not None:
            host_uuid = result._host._uuid
            host_name = result._host.name
        else:
            host_uuid = "include"
            host_name = "include"

        task = tasks_data[task_uuid]

        task.add_host(HostData(host_uuid, host_name, status, result))

    def generate_distributed_traces(
        self,
        tasks_data,
        status,
        end_time,
        traceparent,
        apm_service_name,
        apm_server_url,
        apm_verify_server_cert,
        apm_secret_token,
        apm_api_key,
    ):
        """generate distributed traces from the collected TaskData and HostData"""

        tasks = []
        parent_start_time = None
        for task in tasks_data.values():
            if parent_start_time is None:
                parent_start_time = task.start
            tasks.append(task)

        apm_cli = self.init_apm_client(
            apm_server_url, apm_service_name, apm_verify_server_cert, apm_secret_token, apm_api_key
        )
        if apm_cli:
            with closing(apm_cli):
                instrument()  # Only call this once, as early as possible.
                if traceparent:
                    parent = trace_parent_from_string(traceparent)
                    apm_cli.begin_transaction("Session", trace_parent=parent, start=parent_start_time)
                else:
                    apm_cli.begin_transaction("Session", start=parent_start_time)
                # Populate trace metadata attributes
                label(ansible_version=ansible_version)
                label(ansible_session=self.session, ansible_host_name=self.host, ansible_host_user=self.user)
                if self.ip_address is not None:
                    label(ansible_host_ip=self.ip_address)

                for task_data in tasks:
                    for host_data in task_data.host_data.values():
                        self.create_span_data(apm_cli, task_data, host_data)

                apm_cli.end_transaction(name=__name__, result=status, duration=end_time - parent_start_time)

    def create_span_data(self, apm_cli, task_data, host_data):
        """create the span with the given TaskData and HostData"""

        name = f"[{host_data.name}] {task_data.play}: {task_data.name}"

        message = "success"
        status = "success"
        enriched_error_message = None
        if host_data.status == "included":
            rc = 0
        else:
            res = host_data.result._result
            rc = res.get("rc", 0)
            if host_data.status == "failed":
                message = self.get_error_message(res)
                enriched_error_message = self.enrich_error_message(res)
                status = "failure"
            elif host_data.status == "skipped":
                if "skip_reason" in res:
                    message = res["skip_reason"]
                else:
                    message = "skipped"
                status = "unknown"

        with capture_span(
            task_data.name,
            start=task_data.start,
            span_type="ansible.task.run",
            duration=host_data.finish - task_data.start,
            labels={
                "ansible.task.args": task_data.args,
                "ansible.task.message": message,
                "ansible.task.module": task_data.action,
                "ansible.task.name": name,
                "ansible.task.result": rc,
                "ansible.task.host.name": host_data.name,
                "ansible.task.host.status": host_data.status,
            },
        ) as span:
            span.outcome = status
            if "failure" in status:
                exception = AnsibleRuntimeError(
                    message=f"{task_data.action}: {name} failed with error message {enriched_error_message}"
                )
                apm_cli.capture_exception(exc_info=(type(exception), exception, exception.__traceback__), handled=True)

    def init_apm_client(self, apm_server_url, apm_service_name, apm_verify_server_cert, apm_secret_token, apm_api_key):
        if apm_server_url:
            return Client(
                service_name=apm_service_name,
                server_url=apm_server_url,
                verify_server_cert=False,
                secret_token=apm_secret_token,
                api_key=apm_api_key,
                use_elastic_traceparent_header=True,
                debug=True,
            )

    @staticmethod
    def get_error_message(result):
        if result.get("exception") is not None:
            return ElasticSource._last_line(result["exception"])
        return result.get("msg", "failed")

    @staticmethod
    def _last_line(text):
        lines = text.strip().split("\n")
        return lines[-1]

    @staticmethod
    def enrich_error_message(result):
        message = result.get("msg", "failed")
        exception = result.get("exception")
        stderr = result.get("stderr")
        return f'message: "{message}"\nexception: "{exception}"\nstderr: "{stderr}"'


class CallbackModule(CallbackBase):
    """
    This callback creates distributed traces with Elastic APM.
    """

    CALLBACK_VERSION = 2.0
    CALLBACK_TYPE = "notification"
    CALLBACK_NAME = "community.general.elastic"
    CALLBACK_NEEDS_ENABLED = True

    def __init__(self, display=None):
        super().__init__(display=display)
        self.hide_task_arguments = None
        self.apm_service_name = None
        self.ansible_playbook = None
        self.traceparent = False
        self.play_name = None
        self.tasks_data = None
        self.errors = 0
        self.disabled = False

        if ELASTIC_LIBRARY_IMPORT_ERROR:
            raise AnsibleError(
                "The `elastic-apm` must be installed to use this plugin"
            ) from ELASTIC_LIBRARY_IMPORT_ERROR

        self.tasks_data = OrderedDict()

        self.elastic = ElasticSource(display=self._display)

    def set_options(self, task_keys=None, var_options=None, direct=None):
        super().set_options(task_keys=task_keys, var_options=var_options, direct=direct)

        self.hide_task_arguments = self.get_option("hide_task_arguments")

        self.apm_service_name = self.get_option("apm_service_name")
        if not self.apm_service_name:
            self.apm_service_name = "ansible"

        self.apm_server_url = self.get_option("apm_server_url")
        self.apm_secret_token = self.get_option("apm_secret_token")
        self.apm_api_key = self.get_option("apm_api_key")
        self.apm_verify_server_cert = self.get_option("apm_verify_server_cert")
        self.traceparent = self.get_option("traceparent")

    def v2_playbook_on_start(self, playbook):
        self.ansible_playbook = basename(playbook._file_name)

    def v2_playbook_on_play_start(self, play):
        self.play_name = play.get_name()

    def v2_runner_on_no_hosts(self, task):
        self.elastic.start_task(self.tasks_data, self.hide_task_arguments, self.play_name, task)

    def v2_playbook_on_task_start(self, task, is_conditional):
        self.elastic.start_task(self.tasks_data, self.hide_task_arguments, self.play_name, task)

    def v2_playbook_on_cleanup_task_start(self, task):
        self.elastic.start_task(self.tasks_data, self.hide_task_arguments, self.play_name, task)

    def v2_playbook_on_handler_task_start(self, task):
        self.elastic.start_task(self.tasks_data, self.hide_task_arguments, self.play_name, task)

    def v2_runner_on_failed(self, result, ignore_errors=False):
        self.errors += 1
        self.elastic.finish_task(self.tasks_data, "failed", result)

    def v2_runner_on_ok(self, result):
        self.elastic.finish_task(self.tasks_data, "ok", result)

    def v2_runner_on_skipped(self, result):
        self.elastic.finish_task(self.tasks_data, "skipped", result)

    def v2_playbook_on_include(self, included_file):
        self.elastic.finish_task(self.tasks_data, "included", included_file)

    def v2_playbook_on_stats(self, stats):
        if self.errors == 0:
            status = "success"
        else:
            status = "failure"
        self.elastic.generate_distributed_traces(
            self.tasks_data,
            status,
            time.time(),
            self.traceparent,
            self.apm_service_name,
            self.apm_server_url,
            self.apm_verify_server_cert,
            self.apm_secret_token,
            self.apm_api_key,
        )

    def v2_runner_on_async_failed(self, result, **kwargs):
        self.errors += 1
