// <copyright file="BaggageTests.cs" company="OpenTelemetry Authors">
// Copyright The OpenTelemetry Authors
//
// 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.
// </copyright>

using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using Xunit;

namespace OpenTelemetry.Tests
{
    public class BaggageTests
    {
        private const string K1 = "Key1";
        private const string K2 = "Key2";
        private const string K3 = "Key3";

        private const string V1 = "Value1";
        private const string V2 = "Value2";
        private const string V3 = "Value3";

        [Fact]
        public void EmptyTest()
        {
            Assert.Empty(Baggage.GetBaggage());
            Assert.Empty(Baggage.Current.GetBaggage());
        }

        [Fact]
        public void SetAndGetTest()
        {
            var list = new List<KeyValuePair<string, string>>(2)
            {
                new KeyValuePair<string, string>(K1, V1),
                new KeyValuePair<string, string>(K2, V2),
            };

            Baggage.SetBaggage(K1, V1);
            var baggage = Baggage.Current.SetBaggage(K2, V2);
            Baggage.Current = baggage;

            Assert.NotEmpty(Baggage.GetBaggage());
            Assert.Equal(list, Baggage.GetBaggage(Baggage.Current));

            Assert.Equal(V1, Baggage.GetBaggage(K1));
            Assert.Equal(V1, Baggage.GetBaggage(K1.ToLower()));
            Assert.Equal(V1, Baggage.GetBaggage(K1.ToUpper()));
            Assert.Null(Baggage.GetBaggage("NO_KEY"));
            Assert.Equal(V2, Baggage.Current.GetBaggage(K2));

            Assert.Throws<ArgumentException>(() => Baggage.GetBaggage(null));
        }

        [Fact]
        public void SetExistingKeyTest()
        {
            var list = new List<KeyValuePair<string, string>>(2)
            {
                new KeyValuePair<string, string>(K1, V1),
            };

            Baggage.Current.SetBaggage(new KeyValuePair<string, string>(K1, V1));
            var baggage = Baggage.SetBaggage(K1, V1);
            Baggage.SetBaggage(new Dictionary<string, string> { [K1] = V1 }, baggage);

            Assert.Equal(list, Baggage.GetBaggage());
        }

        [Fact]
        public void SetNullValueTest()
        {
            var baggage = Baggage.Current;
            baggage = Baggage.SetBaggage(K1, V1, baggage);

            Assert.Equal(1, Baggage.Current.Count);
            Assert.Equal(1, baggage.Count);

            Baggage.Current.SetBaggage(K2, null);

            Assert.Equal(1, Baggage.Current.Count);

            Assert.Empty(Baggage.SetBaggage(K1, null).GetBaggage());

            Baggage.SetBaggage(K1, V1);
            Baggage.SetBaggage(new Dictionary<string, string>
            {
                [K1] = null,
                [K2] = V2,
            });
            Assert.Equal(1, Baggage.Current.Count);
            Assert.Contains(Baggage.GetBaggage(), kvp => kvp.Key == K2);
        }

        [Fact]
        public void RemoveTest()
        {
            var empty = Baggage.Current;
            var empty2 = Baggage.RemoveBaggage(K1);
            Assert.True(empty == empty2);

            var baggage = Baggage.SetBaggage(new Dictionary<string, string>
            {
                [K1] = V1,
                [K2] = V2,
                [K3] = V3,
            });

            var baggage2 = Baggage.RemoveBaggage(K1, baggage);

            Assert.Equal(3, baggage.Count);
            Assert.Equal(2, baggage2.Count);

            Assert.DoesNotContain(new KeyValuePair<string, string>(K1, V1), baggage2.GetBaggage());
        }

        [Fact]
        public void ClearTest()
        {
            var baggage = Baggage.SetBaggage(new Dictionary<string, string>
            {
                [K1] = V1,
                [K2] = V2,
                [K3] = V3,
            });

            Assert.Equal(3, baggage.Count);

            Baggage.ClearBaggage();

            Assert.Equal(0, Baggage.Current.Count);
        }

        [Fact]
        public void ContextFlowTest()
        {
            var baggage = Baggage.SetBaggage(K1, V1);
            var baggage2 = Baggage.Current.SetBaggage(K2, V2);
            Baggage.Current = baggage2;
            var baggage3 = Baggage.SetBaggage(K3, V3);

            Assert.Equal(1, baggage.Count);
            Assert.Equal(2, baggage2.Count);
            Assert.Equal(3, baggage3.Count);

            Baggage.Current = baggage;

            var baggage4 = Baggage.SetBaggage(K3, V3);

            Assert.Equal(2, baggage4.Count);
            Assert.DoesNotContain(new KeyValuePair<string, string>(K2, V2), baggage4.GetBaggage());
        }

        [Fact]
        public void EnumeratorTest()
        {
            var list = new List<KeyValuePair<string, string>>(2)
            {
                new KeyValuePair<string, string>(K1, V1),
                new KeyValuePair<string, string>(K2, V2),
            };

            var baggage = Baggage.SetBaggage(K1, V1);
            baggage = Baggage.SetBaggage(K2, V2, baggage);

            var enumerator = Baggage.GetEnumerator(baggage);

            Assert.True(enumerator.MoveNext());
            var tag1 = enumerator.Current;
            Assert.True(enumerator.MoveNext());
            var tag2 = enumerator.Current;
            Assert.False(enumerator.MoveNext());

            Assert.Equal(list, new List<KeyValuePair<string, string>> { tag1, tag2 });

            Baggage.ClearBaggage();

            enumerator = Baggage.GetEnumerator();

            Assert.False(enumerator.MoveNext());
        }

        [Fact]
        public void EqualsTest()
        {
            var bc1 = new Baggage(new Dictionary<string, string>() { [K1] = V1, [K2] = V2 });
            var bc2 = new Baggage(new Dictionary<string, string>() { [K1] = V1, [K2] = V2 });
            var bc3 = new Baggage(new Dictionary<string, string>() { [K2] = V2, [K1] = V1 });
            var bc4 = new Baggage(new Dictionary<string, string>() { [K1] = V1, [K2] = V1 });
            var bc5 = new Baggage(new Dictionary<string, string>() { [K1] = V2, [K2] = V1 });

            Assert.True(bc1.Equals(bc2));

            Assert.False(bc1.Equals(bc3));
            Assert.False(bc1.Equals(bc4));
            Assert.False(bc2.Equals(bc4));
            Assert.False(bc3.Equals(bc4));
            Assert.False(bc5.Equals(bc4));
            Assert.False(bc4.Equals(bc5));
        }

        [Fact]
        public void CreateBaggageTest()
        {
            var baggage = Baggage.Create(null);

            Assert.Equal(default, baggage);

            baggage = Baggage.Create(new Dictionary<string, string>
            {
                [K1] = V1,
                ["key2"] = "value2",
                ["KEY2"] = "VALUE2",
                ["KEY3"] = "VALUE3",
                ["Key3"] = null,
            });

            Assert.Equal(2, baggage.Count);
            Assert.Contains(baggage.GetBaggage(), kvp => kvp.Key == K1);
            Assert.Equal("VALUE2", Baggage.GetBaggage("key2", baggage));
        }

        [Fact]
        public void EqualityTests()
        {
            var emptyBaggage = Baggage.Create(null);

            var baggage = Baggage.SetBaggage(K1, V1);

            Assert.NotEqual(emptyBaggage, baggage);

            Assert.True(emptyBaggage != baggage);

            baggage = Baggage.ClearBaggage(baggage);

            Assert.Equal(emptyBaggage, baggage);

            baggage = Baggage.SetBaggage(K1, V1);

            var baggage2 = Baggage.SetBaggage(null);

            Assert.Equal(baggage, baggage2);

            Assert.False(baggage.Equals(this));
            Assert.True(baggage.Equals((object)baggage2));
        }

        [Fact]
        public void GetHashCodeTests()
        {
            var baggage = Baggage.Current;
            var emptyBaggage = Baggage.Create(null);

            Assert.Equal(emptyBaggage.GetHashCode(), baggage.GetHashCode());

            baggage = Baggage.SetBaggage(K1, V1, baggage);

            Assert.NotEqual(emptyBaggage.GetHashCode(), baggage.GetHashCode());

            var expectedBaggage = Baggage.Create(new Dictionary<string, string> { [K1] = V1 });

            Assert.Equal(expectedBaggage.GetHashCode(), baggage.GetHashCode());
        }

        [Fact]
        public async Task AsyncLocalTests()
        {
            Baggage.SetBaggage("key1", "value1");

            await InnerTask().ConfigureAwait(false);

            Baggage.SetBaggage("key4", "value4");

            Assert.Equal(4, Baggage.Current.Count);
            Assert.Equal("value1", Baggage.GetBaggage("key1"));
            Assert.Equal("value2", Baggage.GetBaggage("key2"));
            Assert.Equal("value3", Baggage.GetBaggage("key3"));
            Assert.Equal("value4", Baggage.GetBaggage("key4"));

            static async Task InnerTask()
            {
                Baggage.SetBaggage("key2", "value2");

                await Task.Yield();

                Baggage.SetBaggage("key3", "value3");

                // key2 & key3 changes don't flow backward automatically
            }
        }

        [Fact]
        public void ThreadSafetyTest()
        {
            Baggage.SetBaggage("rootKey", "rootValue"); // Note: Required to establish a root ExecutionContext containing the BaggageHolder we use as a lock

            Parallel.For(0, 100, (i) =>
            {
                Baggage.SetBaggage($"key{i}", $"value{i}");
            });

            Assert.Equal(101, Baggage.Current.Count);
        }
    }
}
