from unittest import mock

import pytest
from fakeredis import FakeStrictRedis

import sentry_sdk
from sentry_sdk import capture_message, start_transaction
from sentry_sdk.consts import SPANDATA
from sentry_sdk.integrations.redis import RedisIntegration

MOCK_CONNECTION_POOL = mock.MagicMock()
MOCK_CONNECTION_POOL.connection_kwargs = {
    "host": "localhost",
    "port": 63791,
    "db": 1,
}


def test_basic(sentry_init, capture_events):
    sentry_init(integrations=[RedisIntegration()])
    events = capture_events()

    connection = FakeStrictRedis()

    connection.get("foobar")
    capture_message("hi")

    (event,) = events
    (crumb,) = event["breadcrumbs"]["values"]

    assert crumb == {
        "category": "redis",
        "message": "GET 'foobar'",
        "data": {
            "redis.key": "foobar",
            "redis.command": "GET",
            "redis.is_cluster": False,
            "db.operation": "GET",
        },
        "timestamp": crumb["timestamp"],
        "type": "redis",
    }


@pytest.mark.parametrize("span_streaming", [True, False])
@pytest.mark.parametrize(
    "is_transaction, send_default_pii, expected_first_ten",
    [
        (False, False, ["GET 'foo'", "SET 'bar' [Filtered]", "SET 'baz' [Filtered]"]),
        (True, True, ["GET 'foo'", "SET 'bar' 1", "SET 'baz' 2"]),
    ],
)
def test_redis_pipeline(
    sentry_init,
    capture_events,
    capture_items,
    is_transaction,
    send_default_pii,
    expected_first_ten,
    span_streaming,
):
    sentry_init(
        integrations=[RedisIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=send_default_pii,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    connection = FakeStrictRedis()

    if span_streaming:
        items = capture_items("span")
        with sentry_sdk.traces.start_span(name="custom parent"):
            pipeline = connection.pipeline(transaction=is_transaction)
            pipeline.get("foo")
            pipeline.set("bar", 1)
            pipeline.set("baz", 2)
            pipeline.execute()
        sentry_sdk.flush()

        assert len(items) == 2
        pipeline_span, parent_span = items[0].payload, items[1].payload

        assert parent_span["name"] == "custom parent"
        assert parent_span["is_segment"] is True

        assert pipeline_span["name"] == "redis.pipeline.execute"
        assert pipeline_span["attributes"]["sentry.op"] == "db.redis"
        assert pipeline_span["attributes"]["sentry.origin"] == "auto.db.redis"
        assert pipeline_span["attributes"][SPANDATA.DB_SYSTEM_NAME] == "redis"
    else:
        events = capture_events()
        with start_transaction():
            pipeline = connection.pipeline(transaction=is_transaction)
            pipeline.get("foo")
            pipeline.set("bar", 1)
            pipeline.set("baz", 2)
            pipeline.execute()

        (event,) = events
        (span,) = event["spans"]
        assert span["op"] == "db.redis"
        assert span["description"] == "redis.pipeline.execute"
        assert span["data"][SPANDATA.DB_SYSTEM] == "redis"
        assert span["data"]["redis.commands"] == {
            "count": 3,
            "first_ten": expected_first_ten,
        }
        assert span["tags"] == {
            "redis.transaction": is_transaction,
            "redis.is_cluster": False,
        }


@pytest.mark.parametrize("span_streaming", [True, False])
def test_sensitive_data(sentry_init, capture_events, capture_items, span_streaming):
    # fakeredis does not support the AUTH command, so we need to mock it
    with mock.patch(
        "sentry_sdk.integrations.redis.utils._COMMANDS_INCLUDING_SENSITIVE_DATA",
        ["get"],
    ):
        sentry_init(
            integrations=[RedisIntegration()],
            traces_sample_rate=1.0,
            send_default_pii=True,
            trace_lifecycle="stream" if span_streaming else "static",
        )

        connection = FakeStrictRedis()

        if span_streaming:
            items = capture_items("span")
            with sentry_sdk.traces.start_span(name="custom parent"):
                connection.get("this is super secret")
            sentry_sdk.flush()

            assert len(items) == 2
            redis_span, parent_span = items[0].payload, items[1].payload

            assert parent_span["name"] == "custom parent"
            assert redis_span["name"] == "GET [Filtered]"
            assert redis_span["attributes"][SPANDATA.DB_QUERY_TEXT] == "GET [Filtered]"
            assert redis_span["attributes"]["sentry.op"] == "db.redis"
        else:
            events = capture_events()
            with start_transaction():
                connection.get(
                    "this is super secret"
                )  # because fakeredis does not support AUTH we use GET instead

            (event,) = events
            spans = event["spans"]
            assert spans[0]["op"] == "db.redis"
            assert spans[0]["description"] == "GET [Filtered]"


@pytest.mark.parametrize("span_streaming", [True, False])
def test_pii_data_redacted(sentry_init, capture_events, capture_items, span_streaming):
    sentry_init(
        integrations=[RedisIntegration()],
        traces_sample_rate=1.0,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    connection = FakeStrictRedis()

    if span_streaming:
        items = capture_items("span")
        with sentry_sdk.traces.start_span(name="custom parent"):
            connection.set("somekey1", "my secret string1")
            connection.set("somekey2", "my secret string2")
            connection.get("somekey2")
            connection.delete("somekey1", "somekey2")
        sentry_sdk.flush()

        assert len(items) == 5
        set1, set2, get, delete, parent = [item.payload for item in items]

        assert parent["name"] == "custom parent"
        assert set1["name"] == "SET 'somekey1' [Filtered]"
        assert set1["attributes"][SPANDATA.DB_QUERY_TEXT] == "SET 'somekey1' [Filtered]"
        assert set1["attributes"]["sentry.op"] == "db.redis"
        assert set2["name"] == "SET 'somekey2' [Filtered]"
        assert set2["attributes"][SPANDATA.DB_QUERY_TEXT] == "SET 'somekey2' [Filtered]"
        assert get["name"] == "GET 'somekey2'"
        assert delete["name"] == "DEL 'somekey1' [Filtered]"
    else:
        events = capture_events()
        with start_transaction():
            connection.set("somekey1", "my secret string1")
            connection.set("somekey2", "my secret string2")
            connection.get("somekey2")
            connection.delete("somekey1", "somekey2")

        (event,) = events
        spans = event["spans"]
        assert spans[0]["op"] == "db.redis"
        assert spans[0]["description"] == "SET 'somekey1' [Filtered]"
        assert spans[1]["description"] == "SET 'somekey2' [Filtered]"
        assert spans[2]["description"] == "GET 'somekey2'"
        assert spans[3]["description"] == "DEL 'somekey1' [Filtered]"


@pytest.mark.parametrize("span_streaming", [True, False])
def test_pii_data_sent(sentry_init, capture_events, capture_items, span_streaming):
    sentry_init(
        integrations=[RedisIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    connection = FakeStrictRedis()

    if span_streaming:
        items = capture_items("span")
        with sentry_sdk.traces.start_span(name="custom parent"):
            connection.set("somekey1", "my secret string1")
            connection.set("somekey2", "my secret string2")
            connection.get("somekey2")
            connection.delete("somekey1", "somekey2")
        sentry_sdk.flush()

        assert len(items) == 5
        set1, set2, get, delete, parent = [item.payload for item in items]

        assert parent["name"] == "custom parent"
        assert set1["name"] == "SET 'somekey1' 'my secret string1'"
        assert (
            set1["attributes"][SPANDATA.DB_QUERY_TEXT]
            == "SET 'somekey1' 'my secret string1'"
        )
        assert set1["attributes"]["sentry.op"] == "db.redis"
        assert set2["name"] == "SET 'somekey2' 'my secret string2'"
        assert (
            set2["attributes"][SPANDATA.DB_QUERY_TEXT]
            == "SET 'somekey2' 'my secret string2'"
        )
        assert get["name"] == "GET 'somekey2'"
        assert delete["name"] == "DEL 'somekey1' 'somekey2'"
    else:
        events = capture_events()
        with start_transaction():
            connection.set("somekey1", "my secret string1")
            connection.set("somekey2", "my secret string2")
            connection.get("somekey2")
            connection.delete("somekey1", "somekey2")

        (event,) = events
        spans = event["spans"]
        assert spans[0]["op"] == "db.redis"
        assert spans[0]["description"] == "SET 'somekey1' 'my secret string1'"
        assert spans[1]["description"] == "SET 'somekey2' 'my secret string2'"
        assert spans[2]["description"] == "GET 'somekey2'"
        assert spans[3]["description"] == "DEL 'somekey1' 'somekey2'"


@pytest.mark.parametrize("span_streaming", [True, False])
def test_no_data_truncation_by_default(
    sentry_init, capture_events, capture_items, span_streaming
):
    sentry_init(
        integrations=[RedisIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    connection = FakeStrictRedis()
    long_string = "a" * 100000
    short_string = "b" * 10

    if span_streaming:
        items = capture_items("span")
        with sentry_sdk.traces.start_span(name="custom parent"):
            connection.set("somekey1", long_string)
            connection.set("somekey2", short_string)
        sentry_sdk.flush()

        assert len(items) == 3
        set1, set2, parent = [item.payload for item in items]

        assert parent["name"] == "custom parent"
        assert set1["name"] == f"SET 'somekey1' '{long_string}'"
        assert (
            set1["attributes"][SPANDATA.DB_QUERY_TEXT]
            == f"SET 'somekey1' '{long_string}'"
        )
        assert set1["attributes"]["sentry.op"] == "db.redis"
        assert set2["name"] == f"SET 'somekey2' '{short_string}'"
        assert (
            set2["attributes"][SPANDATA.DB_QUERY_TEXT]
            == f"SET 'somekey2' '{short_string}'"
        )
    else:
        events = capture_events()
        with start_transaction():
            connection.set("somekey1", long_string)
            connection.set("somekey2", short_string)

        (event,) = events
        spans = event["spans"]
        assert spans[0]["op"] == "db.redis"
        assert spans[0]["description"] == f"SET 'somekey1' '{long_string}'"
        assert spans[1]["description"] == f"SET 'somekey2' '{short_string}'"


@pytest.mark.parametrize("span_streaming", [True, False])
def test_data_truncation_custom(
    sentry_init, capture_events, capture_items, span_streaming
):
    sentry_init(
        integrations=[RedisIntegration(max_data_size=30)],
        traces_sample_rate=1.0,
        send_default_pii=True,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    connection = FakeStrictRedis()
    long_string = "a" * 100000
    short_string = "b" * 10
    expected_long = "SET 'somekey1' '%s..." % (
        long_string[: 30 - len("...") - len("SET 'somekey1' '")],
    )
    expected_short = "SET 'somekey2' '%s'" % (short_string,)

    if span_streaming:
        items = capture_items("span")
        with sentry_sdk.traces.start_span(name="custom parent"):
            connection.set("somekey1", long_string)
            connection.set("somekey2", short_string)
        sentry_sdk.flush()

        assert len(items) == 3
        set1, set2, parent = [item.payload for item in items]

        assert parent["name"] == "custom parent"
        assert set1["name"] == expected_long
        assert (
            set1["attributes"][SPANDATA.DB_QUERY_TEXT]
            == f"SET 'somekey1' '{long_string}'"
        )
        assert set1["attributes"]["sentry.op"] == "db.redis"
        assert set2["name"] == expected_short
        assert (
            set2["attributes"][SPANDATA.DB_QUERY_TEXT]
            == f"SET 'somekey2' '{short_string}'"
        )
    else:
        events = capture_events()
        with start_transaction():
            connection.set("somekey1", long_string)
            connection.set("somekey2", short_string)

        (event,) = events
        spans = event["spans"]
        assert spans[0]["op"] == "db.redis"
        assert spans[0]["description"] == expected_long
        assert spans[1]["description"] == expected_short


def test_breadcrumbs(sentry_init, capture_events):
    sentry_init(
        integrations=[RedisIntegration(max_data_size=30)],
        send_default_pii=True,
    )
    events = capture_events()

    connection = FakeStrictRedis()

    long_string = "a" * 100000
    connection.set("somekey1", long_string)
    short_string = "b" * 10
    connection.set("somekey2", short_string)

    capture_message("hi")

    (event,) = events
    crumbs = event["breadcrumbs"]["values"]

    assert crumbs[0] == {
        "message": "SET 'somekey1' 'aaaaaaaaaaa...",
        "type": "redis",
        "category": "redis",
        "data": {
            "db.operation": "SET",
            "redis.is_cluster": False,
            "redis.command": "SET",
            "redis.key": "somekey1",
        },
        "timestamp": crumbs[0]["timestamp"],
    }
    assert crumbs[1] == {
        "message": "SET 'somekey2' 'bbbbbbbbbb'",
        "type": "redis",
        "category": "redis",
        "data": {
            "db.operation": "SET",
            "redis.is_cluster": False,
            "redis.command": "SET",
            "redis.key": "somekey2",
        },
        "timestamp": crumbs[1]["timestamp"],
    }


@pytest.mark.parametrize("span_streaming", [True, False])
def test_db_connection_attributes_client(
    sentry_init, capture_events, capture_items, span_streaming
):
    sentry_init(
        traces_sample_rate=1.0,
        integrations=[RedisIntegration()],
        trace_lifecycle="stream" if span_streaming else "static",
    )

    if span_streaming:
        items = capture_items("span")
        with sentry_sdk.traces.start_span(name="custom parent"):
            connection = FakeStrictRedis(connection_pool=MOCK_CONNECTION_POOL)
            connection.get("foobar")
        sentry_sdk.flush()

        assert len(items) == 2
        redis_span, parent_span = items[0].payload, items[1].payload

        assert parent_span["name"] == "custom parent"
        assert redis_span["name"] == "GET 'foobar'"
        attrs = redis_span["attributes"]
        assert attrs["sentry.op"] == "db.redis"
        assert attrs[SPANDATA.DB_QUERY_TEXT] == "GET 'foobar'"
        assert attrs[SPANDATA.DB_SYSTEM_NAME] == "redis"
        assert attrs[SPANDATA.DB_DRIVER_NAME] == "redis-py"
        assert attrs[SPANDATA.DB_NAMESPACE] == "1"
        assert attrs[SPANDATA.SERVER_ADDRESS] == "localhost"
        assert attrs[SPANDATA.SERVER_PORT] == 63791
    else:
        events = capture_events()
        with start_transaction():
            connection = FakeStrictRedis(connection_pool=MOCK_CONNECTION_POOL)
            connection.get("foobar")

        (event,) = events
        (span,) = event["spans"]

        assert span["op"] == "db.redis"
        assert span["description"] == "GET 'foobar'"
        assert span["data"][SPANDATA.DB_SYSTEM] == "redis"
        assert span["data"][SPANDATA.DB_DRIVER_NAME] == "redis-py"
        assert span["data"][SPANDATA.DB_NAME] == "1"
        assert span["data"][SPANDATA.SERVER_ADDRESS] == "localhost"
        assert span["data"][SPANDATA.SERVER_PORT] == 63791


@pytest.mark.parametrize("span_streaming", [True, False])
def test_db_connection_attributes_pipeline(
    sentry_init, capture_events, capture_items, span_streaming
):
    sentry_init(
        traces_sample_rate=1.0,
        integrations=[RedisIntegration()],
        trace_lifecycle="stream" if span_streaming else "static",
    )

    if span_streaming:
        items = capture_items("span")
        with sentry_sdk.traces.start_span(name="custom parent"):
            connection = FakeStrictRedis(connection_pool=MOCK_CONNECTION_POOL)
            pipeline = connection.pipeline(transaction=False)
            pipeline.get("foo")
            pipeline.set("bar", 1)
            pipeline.set("baz", 2)
            pipeline.execute()
        sentry_sdk.flush()

        assert len(items) == 2
        pipeline_span, parent_span = items[0].payload, items[1].payload

        assert parent_span["name"] == "custom parent"
        assert pipeline_span["name"] == "redis.pipeline.execute"
        attrs = pipeline_span["attributes"]
        assert attrs["sentry.op"] == "db.redis"
        assert attrs[SPANDATA.DB_SYSTEM_NAME] == "redis"
        assert attrs[SPANDATA.DB_DRIVER_NAME] == "redis-py"
        assert attrs[SPANDATA.DB_NAMESPACE] == "1"
        assert attrs[SPANDATA.SERVER_ADDRESS] == "localhost"
        assert attrs[SPANDATA.SERVER_PORT] == 63791
    else:
        events = capture_events()
        with start_transaction():
            connection = FakeStrictRedis(connection_pool=MOCK_CONNECTION_POOL)
            pipeline = connection.pipeline(transaction=False)
            pipeline.get("foo")
            pipeline.set("bar", 1)
            pipeline.set("baz", 2)
            pipeline.execute()

        (event,) = events
        (span,) = event["spans"]

        assert span["op"] == "db.redis"
        assert span["description"] == "redis.pipeline.execute"
        assert span["data"][SPANDATA.DB_SYSTEM] == "redis"
        assert span["data"][SPANDATA.DB_DRIVER_NAME] == "redis-py"
        assert span["data"][SPANDATA.DB_NAME] == "1"
        assert span["data"][SPANDATA.SERVER_ADDRESS] == "localhost"
        assert span["data"][SPANDATA.SERVER_PORT] == 63791


@pytest.mark.parametrize("span_streaming", [True, False])
def test_span_origin(sentry_init, capture_events, capture_items, span_streaming):
    sentry_init(
        integrations=[RedisIntegration()],
        traces_sample_rate=1.0,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    connection = FakeStrictRedis()

    if span_streaming:
        items = capture_items("span")
        with sentry_sdk.traces.start_span(name="custom parent"):
            # default case
            connection.set("somekey", "somevalue")

            # pipeline
            pipeline = connection.pipeline(transaction=False)
            pipeline.get("somekey")
            pipeline.set("anotherkey", 1)
            pipeline.execute()
        sentry_sdk.flush()

        assert len(items) == 3
        set_span, pipeline_span, parent_span = [item.payload for item in items]

        assert parent_span["name"] == "custom parent"
        assert parent_span["attributes"]["sentry.origin"] == "manual"
        assert set_span["attributes"]["sentry.origin"] == "auto.db.redis"
        assert (
            set_span["attributes"][SPANDATA.DB_QUERY_TEXT] == "SET 'somekey' [Filtered]"
        )
        assert pipeline_span["attributes"]["sentry.origin"] == "auto.db.redis"
    else:
        events = capture_events()
        with start_transaction(name="custom_transaction"):
            # default case
            connection.set("somekey", "somevalue")

            # pipeline
            pipeline = connection.pipeline(transaction=False)
            pipeline.get("somekey")
            pipeline.set("anotherkey", 1)
            pipeline.execute()

        (event,) = events

        assert event["contexts"]["trace"]["origin"] == "manual"

        for span in event["spans"]:
            assert span["origin"] == "auto.db.redis"
