// Unity C# reference source
// Copyright (c) Unity Technologies. For terms of use, see
// https://unity3d.com/legal/licenses/Unity_Reference_Only_License

using System;
using System.Collections.Generic;
using Unity.Profiling;
using Unity.Profiling.LowLevel.Unsafe;

namespace UnityEditor.Search.Providers
{
    class ProfilerMarkersProvider : BasePerformanceProvider<ProfilerMarkersProvider.ProfilerRecorderInfo>
    {
        internal struct ProfilerRecorderCumulativeData
        {
            public long totalValue;
            public long totalCount;
            public long maxValue;
            public double average;

            public static ProfilerRecorderCumulativeData clear = new();
        }

        internal class ProfilerRecorderInfo
        {
            public ProfilerRecorder recorder;
            public ProfilerRecorderDescription description;
            public ProfilerRecorderCumulativeData data;

            public long GetSampleCount()
            {
                return data.totalCount;
            }

            public long GetMaxValue()
            {
                return data.maxValue;
            }

            public long GetTotalValue()
            {
                return data.totalValue;
            }

            public double GetAverageValue()
            {
                return data.average;
            }

            public void ResetData()
            {
                data = ProfilerRecorderCumulativeData.clear;
            }
        }

        internal const string providerId = "profilermarkers";
        const int k_SamplesCount = 5;
        const double k_NanoSecondsInSeconds = 1e9;
        const double k_KiloByte = 1024 * 1024;
        const double k_MegaByte = k_KiloByte * 1024;

        List<ProfilerRecorderInfo> m_Recorders = new();
        bool m_Enabled;

        static readonly UnitTypeHandle k_SecondUnitTypeHandle = UnitType.GetHandle("s");
        static readonly UnitTypeHandle k_ByteUnitTypeHandle = UnitType.GetHandle("b");
        static readonly UnitTypeHandle k_PercentUnitTypeHandle = UnitType.GetHandle("%");
        static readonly UnitTypeHandle k_HertzUnitTypeHandle = UnitType.GetHandle("hz");

        [SearchItemProvider]
        public static SearchProvider CreateProvider()
        {
            var p = new ProfilerMarkersProvider(providerId, "Profiler Markers");
            p.Initialize();
            p.filterId = "profile:";
            return p;
        }

        protected ProfilerMarkersProvider(string id, string displayName)
            : base(id, displayName)
        {
            var defaultTimeAvgLimit = GetDefaultPerformanceLimit(sampleAvgSelector);
            var defaultTimePeakLimit = GetDefaultPerformanceLimit(samplePeakSelector);

            AddPerformanceLimit(samplePeakSelector, ProfilerMarkerDataUnit.TimeNanoseconds, defaultTimePeakLimit.warningLimit, defaultTimePeakLimit.errorLimit);
            AddPerformanceLimit(sampleAvgSelector, ProfilerMarkerDataUnit.TimeNanoseconds, defaultTimeAvgLimit.warningLimit, defaultTimeAvgLimit.errorLimit);
            AddPerformanceLimit(samplePeakSelector, ProfilerMarkerDataUnit.Bytes, 100 * k_KiloByte, k_MegaByte);
            AddPerformanceLimit(sampleAvgSelector, ProfilerMarkerDataUnit.Bytes, 100 * k_KiloByte, k_MegaByte);

            AddUnitType("s", k_SecondUnitTypeHandle, UnitPowerType.One, UnitPowerType.Milli, UnitPowerType.Micro, UnitPowerType.Nano);
            AddUnitType("b", k_ByteUnitTypeHandle, UnitPowerType.One, UnitPowerType.Kilo, UnitPowerType.Mega, UnitPowerType.Giga, UnitPowerType.Tera, UnitPowerType.Peta);
            AddUnitType("%", k_PercentUnitTypeHandle, UnitPowerType.One);
            AddUnitType("hz", k_HertzUnitTypeHandle, UnitPowerType.One, UnitPowerType.Kilo, UnitPowerType.Mega, UnitPowerType.Giga, UnitPowerType.Tera, UnitPowerType.Peta);
        }

        public override void Initialize()
        {
            base.Initialize();
            onEnable = EnableProvider;
            onDisable = DisableProvider;
        }

        void EnableProvider()
        {
            var availableStatHandles = new List<ProfilerRecorderHandle>();
            ProfilerRecorderHandle.GetAvailable(availableStatHandles);

            foreach (var profilerRecorderHandle in availableStatHandles)
            {
                if (!profilerRecorderHandle.Valid)
                    continue;

                var recorderInfo = new ProfilerRecorderInfo()
                {
                    data = new ProfilerRecorderCumulativeData(),
                    description = ProfilerRecorderHandle.GetDescription(profilerRecorderHandle),
                    recorder = new ProfilerRecorder(profilerRecorderHandle, k_SamplesCount)
                };
                recorderInfo.recorder.Start();
                m_Recorders.Add(recorderInfo);
            }

            EditorApplication.update += OnUpdate;
            m_Enabled = true;
        }

        void DisableProvider()
        {
            m_Enabled = false;
            EditorApplication.update -= OnUpdate;
            foreach (var profilerRecorder in m_Recorders)
            {
                profilerRecorder.recorder.Stop();
                profilerRecorder.recorder.Dispose();
            }
            m_Recorders.Clear();
        }

        void OnUpdate()
        {
            if (!m_Enabled)
                return;
            foreach (var profilerRecorder in m_Recorders)
            {
                profilerRecorder.data = UpdateRecorderData(profilerRecorder.recorder, profilerRecorder.data);
            }
        }

        static ProfilerRecorderCumulativeData UpdateRecorderData(ProfilerRecorder recorder, ProfilerRecorderCumulativeData data)
        {
            var samplesCount = recorder.Capacity;
            if (samplesCount != k_SamplesCount)
                return data;

            if (recorder.Count == 0)
                return data;

            unsafe
            {
                var samples = stackalloc ProfilerRecorderSample[samplesCount];
                recorder.CopyTo(samples, samplesCount);

                var lastSampleIndex = recorder.Count - 1;
                var lastSample = samples[lastSampleIndex];
                if (lastSample is { Value: >= 0, Count: > 0 })
                {
                    data.totalValue = AddWithoutOverflow(data.totalValue, lastSample.Value);
                    data.maxValue = Math.Max(data.maxValue, lastSample.Value);
                    data.totalCount = AddWithoutOverflow(data.totalCount, lastSample.Count);
                }

                // Compute average over samples
                double totalValue = 0;
                double totalCount = 0;
                for (var i = 0; i < samplesCount; ++i)
                {
                    var sampleValue = samples[i].Value;
                    var sampleCount = samples[i].Count;
                    if (sampleValue < 0 || sampleCount <= 0)
                        continue;
                    totalValue += sampleValue;
                    totalCount += sampleCount;
                }

                if (totalCount > 0)
                {
                    data.average = totalValue / totalCount;
                }
            }

            return data;
        }

        static long AddWithoutOverflow(long a, long b)
        {
            if (b > long.MaxValue - a)
                return long.MaxValue;
            return a + b;
        }

        static ProfilerRecorderInfo GetRecorderInfo(SearchItem item)
        {
            return (ProfilerRecorderInfo)item.data;
        }

        protected override string FormatColumnValue(SearchColumnEventArgs args)
        {
            if (!m_Enabled)
                return string.Empty;

            if (args.value == null)
                return string.Empty;

            var valueWithUnit = (ValueWithUnit)args.value;
            if (args.column.selector == sampleCountSelector)
                return valueWithUnit.value.ToString("F0");

            var pri = GetRecorderInfo(args.item);
            if (!pri.recorder.Valid)
                return string.Empty;

            return FormatUnit(valueWithUnit.value, args.column.selector, pri.recorder.UnitType);
        }

        static double ConvertLongWithMaxValue(long value)
        {
            if (value == long.MaxValue)
                return double.PositiveInfinity;
            return value;
        }

        protected override void ResetItems(SearchItem[] items)
        {
            foreach (var item in items)
                GetRecorderInfo(item).ResetData();
        }

        static ValueWithUnit GetSampleCount(SearchItem item)
        {
            var pri = GetRecorderInfo(item);
            return ProfilerRecorderInfoToUnitWithValue(pri, pri.GetSampleCount());
        }

        static ValueWithUnit GetMaxValue(SearchItem item)
        {
            var pri = GetRecorderInfo(item);
            return ProfilerRecorderInfoToUnitWithValue(pri, pri.GetMaxValue());
        }

        static ValueWithUnit GetTotalValue(SearchItem item)
        {
            var pri = GetRecorderInfo(item);
            return ProfilerRecorderInfoToUnitWithValue(pri, ConvertLongWithMaxValue(pri.GetTotalValue()));
        }

        static ValueWithUnit GetAverageValue(SearchItem item)
        {
            var pri = GetRecorderInfo(item);
            return ProfilerRecorderInfoToUnitWithValue(pri, pri.GetAverageValue());
        }

        static ValueWithUnit GetLastTime(SearchItem item)
        {
            var pri = GetRecorderInfo(item);
            return ProfilerRecorderInfoToUnitWithValue(pri, 0);
        }

        [SearchSelector(sampleCountSelector, provider: providerId, cacheable = false)] static object SelectCount(SearchSelectorArgs args) => GetSampleCount(args.current);
        [SearchSelector(samplePeakSelector, provider: providerId, cacheable = false)] static object SelectPeak(SearchSelectorArgs args) => GetMaxValue(args.current);
        [SearchSelector(sampleAvgSelector, provider: providerId, cacheable = false)] static object SelectAvg(SearchSelectorArgs args) => GetAverageValue(args.current);
        [SearchSelector(sampleTotalSelector, provider: providerId, cacheable = false)] static object SelectTotal(SearchSelectorArgs args) => GetTotalValue(args.current);
        [SearchSelector(sampleTotalSelector, provider: providerId, cacheable = false)] static object SelectLastTime(SearchSelectorArgs args) => GetTotalValue(args.current);

        protected override string FetchDescription(SearchItem item, SearchContext context)
        {
            var fullDescription = item.options.HasAny(SearchItemOptions.FullDescription);
            var description = GetTrackerDescription(item, fullDescription ? '\n' : ' ');
            if (item.options.HasAny(SearchItemOptions.Compacted))
                return $"<b>{item.id}</b> {description}";
            return description;
        }

        string GetTrackerDescription(SearchItem item, char splitter)
        {
            if (!m_Enabled)
                return string.Empty;

            var pri = GetRecorderInfo(item);
            var sampleCount = pri.GetSampleCount();
            var peak = pri.GetMaxValue();
            var avg = pri.GetAverageValue();
            var total = pri.GetTotalValue();
            var unitType = pri.recorder.UnitType;
            return $"Sample Count: <b>{sampleCount}</b>{splitter}" +
                $"Peak: {FormatUnit(peak, samplePeakSelector, unitType)}{splitter}" +
                $"Avg: {FormatUnit(avg, sampleAvgSelector, unitType)}{splitter}" +
                $"Total: {FormatUnit(total, sampleTotalSelector, unitType)}{splitter}" +
                $"Category: {pri.description.Category}{splitter}" +
                $"Unit Type: {pri.description.UnitType}{splitter}" +
                $"Data Type: {pri.description.DataType}";
        }

        string FormatUnit(double value, string selector, ProfilerMarkerDataUnit unit)
        {
            var performanceLimitKey = GetPerformanceLimitKey(selector, unit);
            var performanceLimit = GetPerformanceLimit(performanceLimitKey);

            switch (unit)
            {
                case ProfilerMarkerDataUnit.Undefined:
                    return FormatUndefined(value, performanceLimit);
                case ProfilerMarkerDataUnit.TimeNanoseconds:
                    return FormatTime(value, performanceLimit);
                case ProfilerMarkerDataUnit.Bytes:
                    return FormatByte(value, performanceLimit);
                case ProfilerMarkerDataUnit.Count:
                    return FormatCount(value, performanceLimit);
                case ProfilerMarkerDataUnit.Percent:
                    if (selector == sampleTotalSelector)
                        return string.Empty;
                    return FormatPercent(value, performanceLimit);
                case ProfilerMarkerDataUnit.FrequencyHz:
                    if (selector == sampleTotalSelector)
                        return string.Empty;
                    return FormatFrequency(value, performanceLimit);
                default:
                    throw new ArgumentOutOfRangeException(nameof(unit), unit, null);
            }
        }

        static string FormatUndefined(double value, in PerformanceLimit performanceLimit)
        {
            return GetPerformanceLimitLabel(value, performanceLimit, d => $"{d}");
        }

        static string FormatCount(double value, in PerformanceLimit performanceLimit)
        {
            return GetPerformanceLimitLabel(value, performanceLimit, d => $"{Utils.FormatCount((ulong)d)} hit(s)");
        }

        static string FormatTime(double value, in PerformanceLimit performanceLimit)
        {
            // GetTimeLabel expects time in seconds
            return GetTimeLabel(value / k_NanoSecondsInSeconds, performanceLimit);
        }

        static string FormatByte(double value, in PerformanceLimit performanceLimit)
        {
            return GetPerformanceLimitLabel(value, performanceLimit, d => $"{Utils.FormatBytes((long)d)}");
        }

        static string FormatPercent(double value, in PerformanceLimit performanceLimit)
        {
            return GetPerformanceLimitLabel(value, performanceLimit, d => $"{d}%");
        }

        static string FormatFrequency(double value, in PerformanceLimit performanceLimit)
        {
            return GetPerformanceLimitLabel(value, performanceLimit, d => $"{Utils.FormatCount((ulong)d)}Hz");
        }

        protected override IEnumerable<string> YieldPerformanceDataWords(ProfilerRecorderInfo pri)
        {
            yield return pri.description.Name;
            yield return pri.description.Category.Name;
        }

        protected override ValueWithUnit GetPerformanceAverageValue(ProfilerRecorderInfo pri) => ProfilerRecorderInfoToUnitWithValue(pri, pri.GetAverageValue());
        protected override ValueWithUnit GetPerformanceTotalValue(ProfilerRecorderInfo pri) => ProfilerRecorderInfoToUnitWithValue(pri, pri.GetTotalValue());
        protected override ValueWithUnit GetPerformancePeakValue(ProfilerRecorderInfo pri) => ProfilerRecorderInfoToUnitWithValue(pri, pri.GetMaxValue());
        protected override ValueWithUnit GetPerformanceSampleCountValue(ProfilerRecorderInfo pri) => ProfilerRecorderInfoToUnitWithValue(pri, pri.GetSampleCount());
        protected override ValueWithUnit GetPerformanceLastTimeValue(ProfilerRecorderInfo pri) => ProfilerRecorderInfoToUnitWithValue(pri, 0);

        static ValueWithUnit ProfilerRecorderInfoToUnitWithValue(ProfilerRecorderInfo pri, double value)
        {
            switch (pri.description.UnitType)
            {
                case ProfilerMarkerDataUnit.Undefined:
                    return new ValueWithUnit(value, k_UnitlessTypeHandle, UnitPowerType.One);
                case ProfilerMarkerDataUnit.TimeNanoseconds:
                    return new ValueWithUnit(value, k_SecondUnitTypeHandle, UnitPowerType.Nano);
                case ProfilerMarkerDataUnit.Bytes:
                    return new ValueWithUnit(value, k_ByteUnitTypeHandle, UnitPowerType.One);
                case ProfilerMarkerDataUnit.Count:
                    return new ValueWithUnit(value, k_UnitlessTypeHandle, UnitPowerType.One);
                case ProfilerMarkerDataUnit.Percent:
                    return new ValueWithUnit(value, k_PercentUnitTypeHandle, UnitPowerType.One);
                case ProfilerMarkerDataUnit.FrequencyHz:
                    return new ValueWithUnit(value, k_HertzUnitTypeHandle, UnitPowerType.One);
                default:
                    throw new ArgumentOutOfRangeException();
            }
        }

        protected override IEnumerable<SearchItem> FetchItem(SearchContext context, SearchProvider provider)
        {
            var query = m_QueryEngine.ParseQuery(context.searchQuery);
            if (!query.valid)
                yield break;

            foreach (var trackerName in query.Apply(m_Recorders))
                yield return CreateItem(context, provider, trackerName);
        }

        static SearchItem CreateItem(in SearchContext context, in SearchProvider provider, in ProfilerRecorderInfo pri)
        {
            var markerName = pri.description.Name;
            var item = provider.CreateItem(context, markerName, markerName, null, null, pri);
            item.options = SearchItemOptions.AlwaysRefresh;
            return item;
        }

        static int GetPerformanceLimitKey(string selector, ProfilerMarkerDataUnit unit)
        {
            return HashCode.Combine(selector, unit);
        }

        void AddPerformanceLimit(string selector, ProfilerMarkerDataUnit unit, double warningLimit, double errorLimit)
        {
            AddPerformanceLimit(GetPerformanceLimitKey(selector, unit), warningLimit, errorLimit);
        }
    }
}
