from backend.otel_decode import decode_any_value, decode_kv_attributes


def test_decode_any_value_array_and_kvlist():
    from opentelemetry.proto.common.v1 import common_pb2

    v = common_pb2.AnyValue(
        kvlist_value=common_pb2.KeyValueList(
            values=[
                common_pb2.KeyValue(
                    key="messages",
                    value=common_pb2.AnyValue(
                        array_value=common_pb2.ArrayValue(
                            values=[
                                common_pb2.AnyValue(
                                    kvlist_value=common_pb2.KeyValueList(
                                        values=[
                                            common_pb2.KeyValue(
                                                key="role",
                                                value=common_pb2.AnyValue(string_value="user"),
                                            ),
                                            common_pb2.KeyValue(
                                                key="content",
                                                value=common_pb2.AnyValue(string_value="hi"),
                                            ),
                                        ]
                                    )
                                )
                            ]
                        )
                    ),
                ),
                common_pb2.KeyValue(
                    key="tokens",
                    value=common_pb2.AnyValue(int_value=123),
                ),
            ]
        )
    )

    decoded = decode_any_value(v)
    assert decoded["tokens"] == 123
    assert decoded["messages"][0]["role"] == "user"
    assert decoded["messages"][0]["content"] == "hi"


def test_decode_kv_attributes_decodes_complex_values():
    from opentelemetry.proto.common.v1 import common_pb2

    kvs = [
        common_pb2.KeyValue(key="ai.prompt.messages", value=common_pb2.AnyValue(array_value=common_pb2.ArrayValue(values=[]))),
        common_pb2.KeyValue(key="ai.operationId", value=common_pb2.AnyValue(string_value="ai.generateText")),
        common_pb2.KeyValue(key="gen_ai.usage.input_tokens", value=common_pb2.AnyValue(int_value=10)),
    ]

    attrs = decode_kv_attributes(kvs)
    assert attrs["ai.prompt.messages"] == []
    assert attrs["ai.operationId"] == "ai.generateText"
    assert attrs["gen_ai.usage.input_tokens"] == 10

