// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.

using System;
using System.Diagnostics;
using System.IO;
using System.Text;
using System.Text.Encodings.Web;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Html;
using Microsoft.AspNetCore.Mvc.TagHelpers.Cache;
using Microsoft.AspNetCore.Mvc.ViewFeatures.Buffers;
using Microsoft.AspNetCore.Razor.TagHelpers;
using Microsoft.Extensions.Caching.Memory;
using Microsoft.Extensions.Primitives;

namespace Microsoft.AspNetCore.Mvc.TagHelpers
{
    /// <summary>
    /// <see cref="TagHelper"/> implementation targeting &lt;cache&gt; elements.
    /// </summary>
    public class CacheTagHelper : CacheTagHelperBase
    {
        /// <summary>
        /// Prefix used by <see cref="CacheTagHelper"/> instances when creating entries in <see cref="MemoryCache"/>.
        /// </summary>
        public static readonly string CacheKeyPrefix = nameof(CacheTagHelper);

        private const string CachePriorityAttributeName = "priority";

        // We need to come up with a value for the size of entries when storing a gating Task on the cache. Any value
        // greater than 0 will suffice. We choose 56 bytes as an approximation of the size of the task that we store
        // in the cache. This size got calculated as an upper bound for the size of an actual task on an x64 architecture
        // and corresponds to 24 bytes for the object header block plus the 40 bytes added by the members of the task
        // object.
        private const int PlaceholderSize = 64;

        /// <summary>
        /// Creates a new <see cref="CacheTagHelper"/>.
        /// </summary>
        /// <param name="factory">The factory containing the private <see cref="IMemoryCache"/> instance
        /// used by the <see cref="CacheTagHelper"/>.</param>
        /// <param name="htmlEncoder">The <see cref="HtmlEncoder"/> to use.</param>
        public CacheTagHelper(
            CacheTagHelperMemoryCacheFactory factory,
            HtmlEncoder htmlEncoder)
            : base(htmlEncoder)
        {
            MemoryCache = factory.Cache;
        }

        /// <summary>
        /// Gets the <see cref="IMemoryCache"/> instance used to cache entries.
        /// </summary>
        protected IMemoryCache MemoryCache { get; }

        /// <summary>
        /// Gets or sets the <see cref="CacheItemPriority"/> policy for the cache entry.
        /// </summary>
        [HtmlAttributeName(CachePriorityAttributeName)]
        public CacheItemPriority? Priority { get; set; }

        /// <inheritdoc />
        public override async Task ProcessAsync(TagHelperContext context, TagHelperOutput output)
        {
            if (context == null)
            {
                throw new ArgumentNullException(nameof(context));
            }

            if (output == null)
            {
                throw new ArgumentNullException(nameof(output));
            }

            IHtmlContent content;
            if (Enabled)
            {
                var cacheKey = new CacheTagKey(this, context);
                if (MemoryCache.TryGetValue(cacheKey, out Task<IHtmlContent> cachedResult))
                {
                    // There is either some value already cached (as a Task) or a worker processing the output.
                    content = await cachedResult;
                }
                else
                {
                    content = await CreateCacheEntry(cacheKey, output);
                }
            }
            else
            {
                content = await output.GetChildContentAsync();
            }

            // Clear the contents of the "cache" element since we don't want to render it.
            output.SuppressOutput();
            output.Content.SetHtmlContent(content);
        }

        private async Task<IHtmlContent> CreateCacheEntry(CacheTagKey cacheKey, TagHelperOutput output)
        {
            var tokenSource = new CancellationTokenSource();

            var options = GetMemoryCacheEntryOptions();
            options.AddExpirationToken(new CancellationChangeToken(tokenSource.Token));
            options.SetSize(PlaceholderSize);
            var tcs = new TaskCompletionSource<IHtmlContent>(creationOptions: TaskCreationOptions.RunContinuationsAsynchronously);

            // The returned value is ignored, we only do this so that
            // the compiler doesn't complain about the returned task
            // not being awaited
            _ = MemoryCache.Set(cacheKey, tcs.Task, options);

            IHtmlContent content;
            try
            {
                // The entry is set instead of assigning a value to the
                // task so that the expiration options are not impacted
                // by the time it took to compute it.

                // Use the CreateEntry to ensure a cache scope is created that will copy expiration tokens from
                // cache entries created from the GetChildContentAsync call to the current entry.
                var entry = MemoryCache.CreateEntry(cacheKey);

                // The result is processed inside an entry
                // such that the tokens are inherited.

                var result = ProcessContentAsync(output);
                content = await result;
                options.SetSize(GetSize(content));
                entry.SetOptions(options);

                entry.Value = result;

                // An entry gets committed to the cache when disposed gets called. We only want to do this when
                // the content has been correctly generated (didn't throw an exception). For that reason the entry
                // can't be put inside a using block.
                entry.Dispose();

                // Set the result on the TCS once we've committed the entry to the cache since commiting to the cache
                // may throw.
                tcs.SetResult(content);
                return content;
            }
            catch (Exception ex)
            {
                // Remove the worker task from the cache in case it can't complete.
                tokenSource.Cancel();

                // Fail the TCS so other awaiters see the exception.
                tcs.TrySetException(ex);
                throw;
            }
            finally
            {
                // The tokenSource needs to be disposed as the MemoryCache
                // will register a callback on the Token.
                tokenSource.Dispose();
            }
        }

        private long GetSize(IHtmlContent content)
        {
            if (content is CharBufferHtmlContent charBuffer)
            {
                // We need to multiply the size of the buffer
                // by a factor of two due to the fact that
                // characters in .NET are UTF-16 which means
                // every character uses two bytes (surrogates
                // are represented as two characters)
                return charBuffer.Buffer.Length * sizeof(char);
            }

            Debug.Fail($"{nameof(content)} should be an {nameof(CharBufferHtmlContent)}.");
            return -1;
        }

        // Internal for unit testing
        internal MemoryCacheEntryOptions GetMemoryCacheEntryOptions()
        {
            var hasEvictionCriteria = false;
            var options = new MemoryCacheEntryOptions();
            if (ExpiresOn != null)
            {
                hasEvictionCriteria = true;
                options.SetAbsoluteExpiration(ExpiresOn.Value);
            }

            if (ExpiresAfter != null)
            {
                hasEvictionCriteria = true;
                options.SetAbsoluteExpiration(ExpiresAfter.Value);
            }

            if (ExpiresSliding != null)
            {
                hasEvictionCriteria = true;
                options.SetSlidingExpiration(ExpiresSliding.Value);
            }

            if (Priority != null)
            {
                options.SetPriority(Priority.Value);
            }

            if (!hasEvictionCriteria)
            {
                options.SetSlidingExpiration(DefaultExpiration);
            }

            return options;
        }

        private async Task<IHtmlContent> ProcessContentAsync(TagHelperOutput output)
        {
            var content = await output.GetChildContentAsync();

            using (var writer = new CharBufferTextWriter())
            {
                content.WriteTo(writer, HtmlEncoder);
                return new CharBufferHtmlContent(writer.Buffer);
            }
        }

        private class CharBufferTextWriter : TextWriter
        {
            public CharBufferTextWriter()
            {
                Buffer = new PagedCharBuffer(CharArrayBufferSource.Instance);
            }

            public override Encoding Encoding => Null.Encoding;

            public PagedCharBuffer Buffer { get; }

            public override void Write(char value)
            {
                Buffer.Append(value);
            }

            public override void Write(char[] buffer, int index, int count)
            {
                Buffer.Append(buffer, index, count);
            }

            public override void Write(string value)
            {
                Buffer.Append(value);
            }
        }

        private class CharBufferHtmlContent : IHtmlContent
        {
            private readonly PagedCharBuffer _buffer;

            public CharBufferHtmlContent(PagedCharBuffer buffer)
            {
                _buffer = buffer;
            }

            public PagedCharBuffer Buffer => _buffer;

            public void WriteTo(TextWriter writer, HtmlEncoder encoder)
            {
                var length = Buffer.Length;
                if (length == 0)
                {
                    return;
                }

                for (var i = 0; i < Buffer.Pages.Count; i++)
                {
                    var page = Buffer.Pages[i];
                    var pageLength = Math.Min(length, page.Length);
                    writer.Write(page, index: 0, count: pageLength);
                    length -= pageLength;
                }

                Debug.Assert(length == 0);
            }
        }
    }
}
