# Copyright 2010 New Relic, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import json

import pytest
import webtest
from testing_support.fixtures import override_application_settings, validate_attributes
from testing_support.validators.validate_error_event_attributes import (
    validate_error_event_attributes,
)
from testing_support.validators.validate_transaction_event_attributes import (
    validate_transaction_event_attributes,
)
from testing_support.validators.validate_transaction_metrics import (
    validate_transaction_metrics,
)

from newrelic.api.application import application_instance
from newrelic.api.background_task import BackgroundTask, background_task
from newrelic.api.external_trace import ExternalTrace
from newrelic.api.time_trace import current_trace
from newrelic.api.transaction import (
    accept_distributed_trace_headers,
    accept_distributed_trace_payload,
    create_distributed_trace_payload,
    current_span_id,
    current_trace_id,
    current_transaction,
)
from newrelic.api.web_transaction import WSGIWebTransaction
from newrelic.api.wsgi_application import wsgi_application

distributed_trace_intrinsics = ["guid", "traceId", "priority", "sampled"]
inbound_payload_intrinsics = [
    "parent.type",
    "parent.app",
    "parent.account",
    "parent.transportType",
    "parent.transportDuration",
]

payload = {
    "v": [0, 1],
    "d": {
        "ac": "1",
        "ap": "2827902",
        "id": "7d3efb1b173fecfa",
        "pa": "5e5733a911cfbc73",
        "pr": 10.001,
        "sa": True,
        "ti": 1518469636035,
        "tr": "d6b4ba0c3a712ca",
        "ty": "App",
    },
}
parent_order = ["parent_type", "parent_account", "parent_app", "parent_transport_type"]
parent_info = {
    "parent_type": payload["d"]["ty"],
    "parent_account": payload["d"]["ac"],
    "parent_app": payload["d"]["ap"],
    "parent_transport_type": "HTTP",
}


@wsgi_application()
def target_wsgi_application(environ, start_response):
    status = "200 OK"
    output = b"hello world"
    response_headers = [("Content-type", "text/html; charset=utf-8"), ("Content-Length", str(len(output)))]

    txn = current_transaction()

    # Make assertions on the WSGIWebTransaction object
    assert txn._distributed_trace_state
    assert txn.parent_type == "App"
    assert txn.parent_app == "2827902"
    assert txn.parent_account == "1"
    assert txn.parent_span == "7d3efb1b173fecfa"
    assert txn.parent_transport_type == "HTTP"
    assert isinstance(txn.parent_transport_duration, float)
    assert txn._trace_id == "d6b4ba0c3a712ca"
    assert txn.priority == 10.001
    assert txn.sampled

    start_response(status, response_headers)
    return [output]


test_application = webtest.TestApp(target_wsgi_application)

_override_settings = {
    "trusted_account_key": "1",
    "distributed_tracing.enabled": True,
}


_metrics = [
    ("Supportability/DistributedTrace/AcceptPayload/Success", 1),
    ("Supportability/TraceContext/Accept/Success", None),
]


@override_application_settings(_override_settings)
@validate_transaction_metrics("", group="Uri", rollup_metrics=_metrics)
def test_distributed_tracing_web_transaction():
    headers = {"newrelic": json.dumps(payload)}
    response = test_application.get("/", headers=headers)
    assert "X-NewRelic-App-Data" not in response.headers


@pytest.mark.parametrize("span_events", (True, False))
@pytest.mark.parametrize("accept_payload", (True, False))
def test_distributed_trace_attributes(span_events, accept_payload):
    if accept_payload:
        _required_intrinsics = distributed_trace_intrinsics + inbound_payload_intrinsics
        _forgone_txn_intrinsics = []
        _forgone_error_intrinsics = []
        _exact_intrinsics = {
            "parent.type": "Mobile",
            "parent.app": "2827902",
            "parent.account": "1",
            "parent.transportType": "HTTP",
            "traceId": "d6b4ba0c3a712ca",
        }
        _exact_txn_attributes = {"agent": {}, "user": {}, "intrinsic": _exact_intrinsics.copy()}
        _exact_error_attributes = {"agent": {}, "user": {}, "intrinsic": _exact_intrinsics.copy()}
        _exact_txn_attributes["intrinsic"]["parentId"] = "7d3efb1b173fecfa"
        _exact_txn_attributes["intrinsic"]["parentSpanId"] = "c86df80de2e6f51c"

        _forgone_error_intrinsics.append("parentId")
        _forgone_error_intrinsics.append("parentSpanId")
        _forgone_txn_intrinsics.append("grandparentId")
        _forgone_error_intrinsics.append("grandparentId")

        _required_attributes = {"intrinsic": _required_intrinsics, "agent": [], "user": []}
        _forgone_txn_attributes = {"intrinsic": _forgone_txn_intrinsics, "agent": [], "user": []}
        _forgone_error_attributes = {"intrinsic": _forgone_error_intrinsics, "agent": [], "user": []}
    else:
        _required_intrinsics = distributed_trace_intrinsics
        _forgone_txn_intrinsics = _forgone_error_intrinsics = inbound_payload_intrinsics + [
            "grandparentId",
            "parentId",
            "parentSpanId",
        ]

        _required_attributes = {"intrinsic": _required_intrinsics, "agent": [], "user": []}
        _forgone_txn_attributes = {"intrinsic": _forgone_txn_intrinsics, "agent": [], "user": []}
        _forgone_error_attributes = {"intrinsic": _forgone_error_intrinsics, "agent": [], "user": []}
        _exact_txn_attributes = _exact_error_attributes = None

    _forgone_trace_intrinsics = _forgone_error_intrinsics

    test_settings = _override_settings.copy()
    test_settings["span_events.enabled"] = span_events

    @override_application_settings(test_settings)
    @validate_transaction_event_attributes(_required_attributes, _forgone_txn_attributes, _exact_txn_attributes)
    @validate_error_event_attributes(_required_attributes, _forgone_error_attributes, _exact_error_attributes)
    @validate_attributes("intrinsic", _required_intrinsics, _forgone_trace_intrinsics)
    @background_task(name="test_distributed_trace_attributes")
    def _test():
        txn = current_transaction()

        payload = {
            "v": [0, 1],
            "d": {
                "ty": "Mobile",
                "ac": "1",
                "ap": "2827902",
                "id": "c86df80de2e6f51c",
                "tr": "d6b4ba0c3a712ca",
                "ti": 1518469636035,
                "tx": "7d3efb1b173fecfa",
            },
        }
        payload["d"]["pa"] = "5e5733a911cfbc73"

        if accept_payload:
            result = accept_distributed_trace_payload(payload)
            assert result
        else:
            create_distributed_trace_payload()

        try:
            raise ValueError("cookies")
        except ValueError:
            txn.notice_error()

    _test()


_forgone_attributes = {
    "agent": [],
    "user": [],
    "intrinsic": (inbound_payload_intrinsics + ["grandparentId"]),
}


@override_application_settings(_override_settings)
@validate_transaction_event_attributes({}, _forgone_attributes)
@validate_error_event_attributes({}, _forgone_attributes)
@validate_attributes("intrinsic", {}, _forgone_attributes["intrinsic"])
@background_task(name="test_distributed_trace_attrs_omitted")
def test_distributed_trace_attrs_omitted():
    txn = current_transaction()
    try:
        raise ValueError("cookies")
    except ValueError:
        txn.notice_error()


# test our distributed_trace metrics by creating a transaction and then forcing
# it to process a distributed trace payload
@pytest.mark.parametrize("web_transaction", (True, False))
@pytest.mark.parametrize("gen_error", (True, False))
@pytest.mark.parametrize("has_parent", (True, False))
def test_distributed_tracing_metrics(web_transaction, gen_error, has_parent):
    def _make_dt_tag(pi):
        return "%s/%s/%s/%s/all" % tuple(pi[x] for x in parent_order)

    # figure out which metrics we'll see based on the test params
    # note: we'll always see DurationByCaller if the distributed
    # tracing flag is turned on
    metrics = ["DurationByCaller"]
    if gen_error:
        metrics.append("ErrorsByCaller")
    if has_parent:
        metrics.append("TransportDuration")

    tag = None
    dt_payload = copy.deepcopy(payload)

    # if has_parent is True, our metric name will be info about the parent,
    # otherwise it is Unknown/Unknown/Unknown/Unknown
    if has_parent:
        tag = _make_dt_tag(parent_info)
    else:
        # tag = _make_dt_tag(dict((x, "Unknown") for x in parent_order))
        tag = _make_dt_tag(dict((x, "Unknown") for x in parent_info.keys()))
        del dt_payload["d"]["tr"]

    # now run the test
    transaction_name = "test_dt_metrics_%s" % "_".join(metrics)
    _rollup_metrics = [
        ("%s/%s%s" % (x, tag, bt), 1) for x in metrics for bt in ["", "Web" if web_transaction else "Other"]
    ]

    def _make_test_transaction():
        application = application_instance()

        if not web_transaction:
            return BackgroundTask(application, transaction_name)

        environ = {"REQUEST_URI": "/trace_ends_after_txn"}
        tn = WSGIWebTransaction(application, environ)
        tn.set_transaction_name(transaction_name)
        return tn

    @override_application_settings(_override_settings)
    @validate_transaction_metrics(
        transaction_name, background_task=not (web_transaction), rollup_metrics=_rollup_metrics
    )
    def _test():
        with _make_test_transaction() as transaction:
            transaction.accept_distributed_trace_payload(dt_payload)

            if gen_error:
                try:
                    1 / 0
                except ZeroDivisionError:
                    transaction.notice_error()

    _test()


NEW_RELIC_ACCEPTED = [
    ("Supportability/DistributedTrace/AcceptPayload/Success", 1),
    ("Supportability/TraceContext/Accept/Success", None),
    ("Supportability/TraceContext/TraceParent/Accept/Success", None),
    ("Supportability/TraceContext/Accept/Success", None),
]
TRACE_CONTEXT_ACCEPTED = [
    ("Supportability/TraceContext/Accept/Success", 1),
    ("Supportability/TraceContext/TraceParent/Accept/Success", 1),
    ("Supportability/TraceContext/Accept/Success", 1),
    ("Supportability/DistributedTrace/AcceptPayload/Success", None),
]
NO_HEADERS_ACCEPTED = [
    ("Supportability/DistributedTrace/AcceptPayload/Success", None),
    ("Supportability/TraceContext/Accept/Success", None),
    ("Supportability/TraceContext/TraceParent/Accept/Success", None),
    ("Supportability/TraceContext/Accept/Success", None),
]
TRACEPARENT = "00-0af7651916cd43dd8448eb211c80319c-00f067aa0ba902b7-01"
TRACESTATE = "rojo=f06a0ba902b7,congo=t61rcWkgMzE"


@pytest.mark.parametrize(
    "traceparent,tracestate,newrelic,metrics",
    [
        (False, False, False, NO_HEADERS_ACCEPTED),
        (False, False, True, NEW_RELIC_ACCEPTED),
        (False, True, True, NEW_RELIC_ACCEPTED),
        (False, True, False, NO_HEADERS_ACCEPTED),
        (True, True, True, TRACE_CONTEXT_ACCEPTED),
        (True, False, False, TRACE_CONTEXT_ACCEPTED),
        (True, False, True, TRACE_CONTEXT_ACCEPTED),
        (True, True, False, TRACE_CONTEXT_ACCEPTED),
    ],
)
@override_application_settings(_override_settings)
def test_distributed_tracing_backwards_compatibility(traceparent, tracestate, newrelic, metrics):
    headers = []
    if traceparent:
        headers.append(("traceparent", TRACEPARENT))
    if tracestate:
        headers.append(("tracestate", TRACESTATE))
    if newrelic:
        headers.append(("newrelic", json.dumps(payload)))

    @validate_transaction_metrics(
        "test_distributed_tracing_backwards_compatibility", background_task=True, rollup_metrics=metrics
    )
    @background_task(name="test_distributed_tracing_backwards_compatibility")
    def _test():
        accept_distributed_trace_headers(headers)

    _test()


@background_task(name="test_current_trace_id_api_inside_transaction")
def test_current_trace_id_api_inside_transaction():
    trace_id = current_trace_id()
    assert len(trace_id) == 32
    assert trace_id == current_transaction().trace_id


def test_current_trace_id_api_outside_transaction():
    trace_id = current_trace_id()
    assert trace_id is None


@background_task(name="test_current_span_id_api_inside_transaction")
def test_current_span_id_inside_transaction():
    span_id = current_span_id()
    assert span_id == current_trace().guid


def test_current_span_id_outside_transaction():
    span_id = current_span_id()
    assert span_id is None


@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_outbound_dt_payload_generation(trusted_account_key):
    @override_application_settings(
        {
            "distributed_tracing.enabled": True,
            "account_id": "1",
            "trusted_account_key": trusted_account_key,
            "primary_application_id": "1",
        }
    )
    @background_task(name="test_outbound_dt_payload_generation")
    def _test_outbound_dt_payload_generation():
        transaction = current_transaction()
        payload = ExternalTrace.generate_request_headers(transaction)
        if trusted_account_key:
            assert payload
            # Ensure trusted account key present as vendor
            assert dict(payload)["tracestate"].startswith("1@nr=")
        else:
            assert not payload

    _test_outbound_dt_payload_generation()


@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_inbound_dt_payload_acceptance(trusted_account_key):
    @override_application_settings(
        {
            "distributed_tracing.enabled": True,
            "account_id": "1",
            "trusted_account_key": trusted_account_key,
            "primary_application_id": "1",
        }
    )
    @background_task(name="_test_inbound_dt_payload_acceptance")
    def _test_inbound_dt_payload_acceptance():
        transaction = current_transaction()

        payload = {
            "v": [0, 1],
            "d": {
                "ty": "Mobile",
                "ac": "1",
                "tk": "1",
                "ap": "2827902",
                "pa": "5e5733a911cfbc73",
                "id": "7d3efb1b173fecfa",
                "tr": "d6b4ba0c3a712ca",
                "ti": 1518469636035,
                "tx": "8703ff3d88eefe9d",
            },
        }

        result = transaction.accept_distributed_trace_payload(payload)
        if trusted_account_key:
            assert result
        else:
            assert not result

    _test_inbound_dt_payload_acceptance()
