// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using System.Collections.Immutable;
using System.Text;
using System.Text.RegularExpressions;
using Bicep.Core.SourceGraph;
using Bicep.Core.Syntax;
using Bicep.Core.Text;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;

namespace Bicep.Core.Emit
{
    public record RawSourceMap(
        IList<RawSourceMapFileEntry> Entries);

    public record RawSourceMapFileEntry(
        BicepSourceFile SourceFile,
        IList<SourceMapRawEntry> SourceMap);

    public record SourceMapRawEntry(
        TextSpan SourcePosition,
        IList<TextSpan> TargetPositions);

    public class PositionTrackingJsonTextWriter : JsonTextWriter
    {
        private class PositionTrackingTextWriter : TextWriter
        {
            private readonly TextWriter internalWriter = new StringWriter();

            public int CurrentPosition;

            public List<int> CommaPositions = new();

            public PositionTrackingTextWriter(TextWriter textWriter)
            {
                this.internalWriter = textWriter;
                this.NewLine = textWriter.NewLine;
            }

            public override Encoding Encoding => this.internalWriter.Encoding;

            public override string? ToString() => internalWriter.ToString();

            public override void Write(char value)
            {
                if (value == ',')
                {
                    CommaPositions.Add(CurrentPosition);
                }

                this.internalWriter.Write(value);

                CurrentPosition++;
            }
        }

        private static readonly Regex JsonWhitespaceStrippingRegex = new(@"(""(?:[^""\\]|\\.)*"")|\s+", RegexOptions.Compiled);

        private readonly RawSourceMap rawSourceMap;
        private readonly BicepSourceFile? sourceFile;
        private readonly PositionTrackingTextWriter trackingWriter;

        public PositionTrackingJsonTextWriter(TextWriter textWriter, BicepSourceFile? sourceFile = null, RawSourceMap? rawSourceMap = null)
            : this(new(textWriter), sourceFile, rawSourceMap)
        {
        }

        private PositionTrackingJsonTextWriter(PositionTrackingTextWriter trackingWriter, BicepSourceFile? sourceFile, RawSourceMap? rawSourceMap)
            : base(trackingWriter)
        {
            this.rawSourceMap = rawSourceMap ?? new RawSourceMap(new List<RawSourceMapFileEntry>());
            this.sourceFile = sourceFile;
            this.trackingWriter = trackingWriter;
        }

        public void WriteExpressionWithPosition(IPositionable? sourcePosition, Action expressionFunc)
        {
            var startPos = this.trackingWriter.CurrentPosition;

            expressionFunc();

            AddSourceMapping(sourcePosition, startPos);
        }

        public void WriteObjectWithPosition(IPositionable? sourcePosition, Action propertiesFunc)
        {
            var startPos = this.trackingWriter.CurrentPosition;

            base.WriteStartObject();
            propertiesFunc();
            base.WriteEndObject();

            AddSourceMapping(sourcePosition, startPos);
        }

        public void WriteArrayWithPosition(IPositionable? sourcePosition, Action itemsFunc)
        {
            var startPos = this.trackingWriter.CurrentPosition;

            base.WriteStartArray();
            itemsFunc();
            base.WriteEndArray();

            AddSourceMapping(sourcePosition, startPos);
        }

        public void WritePropertyWithPosition(IPositionable? keyPosition, string name, Action valueFunc)
        {
            var startPos = this.trackingWriter.CurrentPosition;

            base.WritePropertyName(name);
            valueFunc();

            AddSourceMapping(keyPosition, startPos);
        }

        public void AddNestedSourceMap(PositionTrackingJsonTextWriter nestedTrackingJsonWriter)
        {
            // offset all raw mappings by current position
            var offset = this.trackingWriter.CurrentPosition;
            foreach (var fileEntry in nestedTrackingJsonWriter.rawSourceMap.Entries)
            {
                foreach (var sourceMapEntry in fileEntry.SourceMap)
                {
                    foreach (var position in sourceMapEntry.TargetPositions)
                    {
                        var positionWithOffset = new TextSpan(position.Position + offset, position.Length);

                        AddRawMapping(fileEntry.SourceFile, sourceMapEntry.SourcePosition, positionWithOffset);
                    }
                }
            }
        }

        public override string? ToString() => this.trackingWriter.ToString();

        private void AddSourceMapping(IPositionable? bicepSyntax, int jsonStartPosition)
        {
            if (this.sourceFile == null || bicepSyntax == null || bicepSyntax.Span.Length == 0)
            {
                return;
            }

            TextSpan bicepPosition = bicepSyntax.Span;

            // account for leading nodes (decorators)
            if (bicepSyntax is StatementSyntax syntax)
            {
                var lastLeadingNode = syntax.LeadingNodes.LastOrDefault();
                if (lastLeadingNode is not null)
                {
                    bicepPosition = new TextSpan(
                        lastLeadingNode.Span.Position + lastLeadingNode.Span.Length,
                        syntax.Span.Length - syntax.LeadingNodes.Sum(node => node.Span.Length));
                }
            }

            // increment start position if starting on a comma (happens when outputting successive items in objects and arrays)
            if (this.trackingWriter.CommaPositions.Contains(jsonStartPosition))
            {
                jsonStartPosition++;
            }

            var jsonEndPosition = this.trackingWriter.CurrentPosition - 1;
            var jsonPosition = new TextSpan(jsonStartPosition, jsonEndPosition - jsonStartPosition);

            AddRawMapping(this.sourceFile, bicepPosition, jsonPosition);
        }

        private void AddRawMapping(BicepSourceFile bicepFile, TextSpan bicepPosition, TextSpan jsonPosition)
        {
            var fileEntry = this.rawSourceMap.Entries.FirstOrDefault(entry => entry.SourceFile.Equals(bicepFile));
            if (fileEntry is null)
            {
                fileEntry = new RawSourceMapFileEntry(bicepFile, new List<SourceMapRawEntry>());

                this.rawSourceMap.Entries.Add(fileEntry);
            }

            var sourceMapEntry = fileEntry.SourceMap.FirstOrDefault(i => i.SourcePosition == bicepPosition);
            if (sourceMapEntry is null)
            {
                sourceMapEntry = new SourceMapRawEntry(bicepPosition, new List<TextSpan>());

                fileEntry.SourceMap.Add(sourceMapEntry);
            }

            sourceMapEntry.TargetPositions.Add(jsonPosition);
        }

        public SourceMap? ProcessRawSourceMap(JToken templateWithHash)
        {
            if (this.sourceFile == null || this.rawSourceMap == null)
            {
                return null;
            }

            var formattedTemplateLines = templateWithHash
                .ToString(Formatting.Indented)
                .Split(Environment.NewLine, StringSplitOptions.None);

            // get line starts of unformatted JSON by stripping formatting from each line of formatted JSON
            var unformattedLineStarts = formattedTemplateLines
                .Aggregate(
                    new List<int>() { 0 }, // first line starts at position 0
                    (lineStarts, line) =>
                    {
                        var unformattedLine = JsonWhitespaceStrippingRegex.Replace(line, "$1");
                        lineStarts.Add(lineStarts.Last() + unformattedLine.Length);
                        return lineStarts;
                    });

            // get position and length of template hash (relying on the first occurrence)
            (var templateHashStartPosition, var templateHashLength) = formattedTemplateLines
                .Select((value, index) => new { lineNumber = index, lineValue = value })
                .Where(item => item.lineValue.Contains(TemplateWriter.TemplateHashPropertyName))
                .Select(item =>
                {
                    var startPosition = unformattedLineStarts[item.lineNumber];
                    var unformattedLineLength = unformattedLineStarts[item.lineNumber + 1] - startPosition;
                    return (startPosition, unformattedLineLength + 1); // account for comma by adding 1 to length
                })
                .FirstOrDefault();


            var weights = new int[unformattedLineStarts.Count];
            Array.Fill(weights, int.MaxValue);

            var sourceMapFileEntries = new List<SourceMapFileEntry>();
            var entrypointFileName = sourceFile.GetFileName();

            foreach (var bicepFileEntry in this.rawSourceMap.Entries)
            {
                var bicepRelativeFilePath = bicepFileEntry.SourceFile == sourceFile
                    ? entrypointFileName
                    : bicepFileEntry.SourceFile.FileHandle.Uri.GetPathRelativeTo(sourceFile.FileHandle.Uri);
                var sourceMapEntries = new List<SourceMapEntry>();

                foreach (var sourceMapEntry in bicepFileEntry.SourceMap)
                {
                    var bicepLine = TextCoordinateConverter.GetPosition(
                        bicepFileEntry.SourceFile.LineStarts, sourceMapEntry.SourcePosition.Position).line;

                    for (int i = 0; i < sourceMapEntry.TargetPositions.Count; i++)
                    {
                        var jsonPosition = sourceMapEntry.TargetPositions[i];

                        // increment positions by templateHashLength that occur after hash start position
                        if (jsonPosition.Position >= templateHashStartPosition)
                        {
                            jsonPosition = new TextSpan(jsonPosition.Position + templateHashLength, jsonPosition.Length);
                            sourceMapEntry.TargetPositions[i] = jsonPosition; // update RawSourceMap in place
                        }

                        var jsonStartPos = jsonPosition.Position;
                        var jsonEndPos = jsonStartPos + jsonPosition.Length;

                        // transform offsets in rawSourceMap to line numbers for formatted JSON using unformattedLineStarts
                        var jsonStartLine = TextCoordinateConverter.GetPosition(unformattedLineStarts, jsonStartPos).line;
                        var jsonEndLine = TextCoordinateConverter.GetPosition(unformattedLineStarts, jsonEndPos).line;

                        // write most specific mapping available for each json line (less lines => stronger weight)
                        int weight = jsonEndLine - jsonStartLine;
                        for (int jsonLine = jsonStartLine; jsonLine <= jsonEndLine; jsonLine++)
                        {
                            // write new mapping if weight is stronger than existing mapping
                            if (weight < weights[jsonLine])
                            {
                                sourceMapEntries.RemoveAll(i => i.TargetLine == jsonLine);
                                sourceMapEntries.Add(new SourceMapEntry(bicepLine, jsonLine));
                                weights[jsonLine] = weight;
                            }
                        }
                    }
                }

                // ensure ordering is deterministic
                var sortedEntries = sourceMapEntries
                    .OrderBy(x => x.SourceLine).ThenBy(x => x.TargetLine)
                    .ToImmutableArray();

                var fileEntry = new SourceMapFileEntry(
                    bicepRelativeFilePath,
                    sortedEntries);
                sourceMapFileEntries.Add(fileEntry);
            }

            return new SourceMap(
               entrypointFileName,
               [.. sourceMapFileEntries]);
        }
    }
}
