import json

import pytest

import sentry_sdk
from sentry_sdk.tracing import Transaction, Span
from sentry_sdk.tracing_utils import (
    compute_tracestate_value,
    extract_sentrytrace_data,
    extract_tracestate_data,
    reinflate_tracestate,
)
from sentry_sdk.utils import from_base64, to_base64


try:
    from unittest import mock  # python 3.3 and above
except ImportError:
    import mock  # python < 3.3


def test_tracestate_computation(sentry_init):
    sentry_init(
        dsn="https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012",
        environment="dogpark",
        release="off.leash.park",
    )

    sentry_sdk.set_user({"id": 12312013, "segment": "bigs"})

    transaction = Transaction(
        name="/interactions/other-dogs/new-dog",
        op="greeting.sniff",
        trace_id="12312012123120121231201212312012",
    )

    # force lazy computation to create a value
    transaction.to_tracestate()

    computed_value = transaction._sentry_tracestate.replace("sentry=", "")
    # we have to decode and reinflate the data because we can guarantee that the
    # order of the entries in the jsonified dict will be the same here as when
    # the tracestate is computed
    reinflated_trace_data = json.loads(from_base64(computed_value))

    assert reinflated_trace_data == {
        "trace_id": "12312012123120121231201212312012",
        "environment": "dogpark",
        "release": "off.leash.park",
        "public_key": "dogsarebadatkeepingsecrets",
        "user": {"id": 12312013, "segment": "bigs"},
        "transaction": "/interactions/other-dogs/new-dog",
    }


def test_doesnt_add_new_tracestate_to_transaction_when_none_given(sentry_init):
    sentry_init(
        dsn="https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012",
        environment="dogpark",
        release="off.leash.park",
    )

    transaction = Transaction(
        name="/interactions/other-dogs/new-dog",
        op="greeting.sniff",
        # sentry_tracestate=< value would be passed here >
    )

    assert transaction._sentry_tracestate is None


def test_adds_tracestate_to_transaction_when_to_traceparent_called(sentry_init):
    sentry_init(
        dsn="https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012",
        environment="dogpark",
        release="off.leash.park",
    )

    transaction = Transaction(
        name="/interactions/other-dogs/new-dog",
        op="greeting.sniff",
    )

    # no inherited tracestate, and none created in Transaction constructor
    assert transaction._sentry_tracestate is None

    transaction.to_tracestate()

    assert transaction._sentry_tracestate is not None


def test_adds_tracestate_to_transaction_when_getting_trace_context(sentry_init):
    sentry_init(
        dsn="https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012",
        environment="dogpark",
        release="off.leash.park",
    )

    transaction = Transaction(
        name="/interactions/other-dogs/new-dog",
        op="greeting.sniff",
    )

    # no inherited tracestate, and none created in Transaction constructor
    assert transaction._sentry_tracestate is None

    transaction.get_trace_context()

    assert transaction._sentry_tracestate is not None


@pytest.mark.parametrize(
    "set_by", ["inheritance", "to_tracestate", "get_trace_context"]
)
def test_tracestate_is_immutable_once_set(sentry_init, monkeypatch, set_by):
    monkeypatch.setattr(
        sentry_sdk.tracing,
        "compute_tracestate_entry",
        mock.Mock(return_value="sentry=doGsaREgReaT"),
    )

    sentry_init(
        dsn="https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012",
        environment="dogpark",
        release="off.leash.park",
    )

    # for each scenario, get to the point where tracestate has been set
    if set_by == "inheritance":
        transaction = Transaction(
            name="/interactions/other-dogs/new-dog",
            op="greeting.sniff",
            sentry_tracestate=("sentry=doGsaREgReaT"),
        )
    else:
        transaction = Transaction(
            name="/interactions/other-dogs/new-dog",
            op="greeting.sniff",
        )

        if set_by == "to_tracestate":
            transaction.to_tracestate()
        if set_by == "get_trace_context":
            transaction.get_trace_context()

    assert transaction._sentry_tracestate == "sentry=doGsaREgReaT"

    # user data would be included in tracestate if it were recomputed at this point
    sentry_sdk.set_user({"id": 12312013, "segment": "bigs"})

    # value hasn't changed
    assert transaction._sentry_tracestate == "sentry=doGsaREgReaT"


@pytest.mark.parametrize("sampled", [True, False, None])
def test_to_traceparent(sentry_init, sampled):

    transaction = Transaction(
        name="/interactions/other-dogs/new-dog",
        op="greeting.sniff",
        trace_id="12312012123120121231201212312012",
        sampled=sampled,
    )

    traceparent = transaction.to_traceparent()

    trace_id, parent_span_id, parent_sampled = traceparent.split("-")
    assert trace_id == "12312012123120121231201212312012"
    assert parent_span_id == transaction.span_id
    assert parent_sampled == (
        "1" if sampled is True else "0" if sampled is False else ""
    )


def test_to_tracestate(sentry_init):
    sentry_init(
        dsn="https://dogsarebadatkeepingsecrets@squirrelchasers.ingest.sentry.io/12312012",
        environment="dogpark",
        release="off.leash.park",
    )

    # it correctly uses the value from the transaction itself or the span's
    # containing transaction
    transaction_no_third_party = Transaction(
        trace_id="12312012123120121231201212312012",
        sentry_tracestate="sentry=doGsaREgReaT",
    )
    non_orphan_span = Span()
    non_orphan_span._containing_transaction = transaction_no_third_party
    assert transaction_no_third_party.to_tracestate() == "sentry=doGsaREgReaT"
    assert non_orphan_span.to_tracestate() == "sentry=doGsaREgReaT"

    # it combines sentry and third-party values correctly
    transaction_with_third_party = Transaction(
        trace_id="12312012123120121231201212312012",
        sentry_tracestate="sentry=doGsaREgReaT",
        third_party_tracestate="maisey=silly",
    )
    assert (
        transaction_with_third_party.to_tracestate()
        == "sentry=doGsaREgReaT,maisey=silly"
    )

    # it computes a tracestate from scratch for orphan transactions
    orphan_span = Span(
        trace_id="12312012123120121231201212312012",
    )
    assert orphan_span._containing_transaction is None
    assert orphan_span.to_tracestate() == "sentry=" + compute_tracestate_value(
        {
            "trace_id": "12312012123120121231201212312012",
            "environment": "dogpark",
            "release": "off.leash.park",
            "public_key": "dogsarebadatkeepingsecrets",
        }
    )


@pytest.mark.parametrize("sampling_decision", [True, False])
def test_sentrytrace_extraction(sampling_decision):
    sentrytrace_header = "12312012123120121231201212312012-0415201309082013-{}".format(
        1 if sampling_decision is True else 0
    )
    assert extract_sentrytrace_data(sentrytrace_header) == {
        "trace_id": "12312012123120121231201212312012",
        "parent_span_id": "0415201309082013",
        "parent_sampled": sampling_decision,
    }


@pytest.mark.parametrize(
    ("incoming_header", "expected_sentry_value", "expected_third_party"),
    [
        # sentry only
        ("sentry=doGsaREgReaT", "sentry=doGsaREgReaT", None),
        # sentry only, invalid (`!` isn't a valid base64 character)
        ("sentry=doGsaREgReaT!", None, None),
        # stuff before
        ("maisey=silly,sentry=doGsaREgReaT", "sentry=doGsaREgReaT", "maisey=silly"),
        # stuff after
        ("sentry=doGsaREgReaT,maisey=silly", "sentry=doGsaREgReaT", "maisey=silly"),
        # stuff before and after
        (
            "charlie=goofy,sentry=doGsaREgReaT,maisey=silly",
            "sentry=doGsaREgReaT",
            "charlie=goofy,maisey=silly",
        ),
        # multiple before
        (
            "charlie=goofy,maisey=silly,sentry=doGsaREgReaT",
            "sentry=doGsaREgReaT",
            "charlie=goofy,maisey=silly",
        ),
        # multiple after
        (
            "sentry=doGsaREgReaT,charlie=goofy,maisey=silly",
            "sentry=doGsaREgReaT",
            "charlie=goofy,maisey=silly",
        ),
        # multiple before and after
        (
            "charlie=goofy,maisey=silly,sentry=doGsaREgReaT,bodhi=floppy,cory=loyal",
            "sentry=doGsaREgReaT",
            "charlie=goofy,maisey=silly,bodhi=floppy,cory=loyal",
        ),
        # only third-party data
        ("maisey=silly", None, "maisey=silly"),
        # invalid third-party data, valid sentry data
        ("maisey_is_silly,sentry=doGsaREgReaT", "sentry=doGsaREgReaT", None),
        # valid third-party data, invalid sentry data
        ("maisey=silly,sentry=doGsaREgReaT!", None, "maisey=silly"),
        # nothing valid at all
        ("maisey_is_silly,sentry=doGsaREgReaT!", None, None),
    ],
)
def test_tracestate_extraction(
    incoming_header, expected_sentry_value, expected_third_party
):
    assert extract_tracestate_data(incoming_header) == {
        "sentry_tracestate": expected_sentry_value,
        "third_party_tracestate": expected_third_party,
    }


# TODO (kmclb) remove this parameterization once tracestate is a real feature
@pytest.mark.parametrize("tracestate_enabled", [True, False])
def test_iter_headers(sentry_init, monkeypatch, tracestate_enabled):
    monkeypatch.setattr(
        Transaction,
        "to_traceparent",
        mock.Mock(return_value="12312012123120121231201212312012-0415201309082013-0"),
    )
    monkeypatch.setattr(
        Transaction,
        "to_tracestate",
        mock.Mock(return_value="sentry=doGsaREgReaT,charlie=goofy"),
    )
    monkeypatch.setattr(
        sentry_sdk.tracing,
        "has_tracestate_enabled",
        mock.Mock(return_value=tracestate_enabled),
    )

    transaction = Transaction(
        name="/interactions/other-dogs/new-dog",
        op="greeting.sniff",
    )

    headers = dict(transaction.iter_headers())
    assert (
        headers["sentry-trace"] == "12312012123120121231201212312012-0415201309082013-0"
    )
    if tracestate_enabled:
        assert "tracestate" in headers
        assert headers["tracestate"] == "sentry=doGsaREgReaT,charlie=goofy"
    else:
        assert "tracestate" not in headers


@pytest.mark.parametrize(
    "data",
    [  # comes out with no trailing `=`
        {"name": "Maisey", "birthday": "12/31/12"},
        # comes out with one trailing `=`
        {"dogs": "yes", "cats": "maybe"},
        # comes out with two trailing `=`
        {"name": "Charlie", "birthday": "11/21/12"},
    ],
)
def test_tracestate_reinflation(data):
    encoded_tracestate = to_base64(json.dumps(data)).strip("=")
    assert reinflate_tracestate(encoded_tracestate) == data