using System.Text.Json;
using JetBrains.Common.CallTreeStorage.Dfs;
using JetBrains.Common.Util.Metadata;
using JetBrains.DotTrace.Dal.Performance.RemoteTree.CoreLogic;
using JetBrains.DotTrace.Dal.Performance.SnapshotComponents.BothSides;
using JetBrains.DotTrace.Dal.Performance.SnapshotComponents.Remote;
using JetBrains.DotTrace.Dal.Performance.SnapshotStorage;
using JetBrains.DotTrace.Dal.Timeline.SnapshotComponents.Remote;
using JetBrains.DotTrace.DalInterface.Performance.CallTree;
using JetBrains.DotTrace.DalInterface.Performance.Fuids;
using JetBrains.DotTrace.DataStructures.CallTree;
using JetBrains.DotTrace.DataStructures.Payloads;
using JetBrains.Lifetimes;
using JetBrains.Metadata.Access;
using JetBrains.Profiler.Snapshot.Converters;
using JetBrains.Profiler.Snapshot.Interface.Section;
using JetBrains.Profiler.Snapshot.Interface.Section.Metadata;
using JetBrains.Profiler.Snapshot.Interface.Section.Metadata.Helpers;
using JetBrains.Profiler.Snapshot.Section.Metadata;
using JetBrains.Util;
using JetBrains.Util.Logging;

return await ProgramMain(args);

static async Task<int> ProgramMain(string[] args)
{
    try
    {
        var options = ParseArguments(args);
        if (options is null)
        {
            Console.Error.WriteLine("Usage: DtpSnapshotExtractor <snapshot.dtp> [--top N] [--filter text] [--flat] [--include-idle]");
            return 2;
        }

        if (!File.Exists(options.SnapshotPath))
        {
            Console.Error.WriteLine($"Snapshot not found: {options.SnapshotPath}");
            return 2;
        }

        RegisterAssemblyResolver(Environment.GetEnvironmentVariable("DOTTRACE_APP_DIR"));

        using var lifetimeDef = Lifetime.Define();
        var lifetime = lifetimeDef.Lifetime;
        var logger = Logger.GetLogger("DtpSnapshotExtractor");
        var masks = new SnapshotMasksComponent();
        var snapshotFile = FileSystemPath.Parse(options.SnapshotPath);
        var container = new SnapshotStorageContainer(lifetime, logger, masks, snapshotFile);
        var callTreeSections = new CallTreeSections(container, masks);
        var headerSections = new HeaderSections(container);
        var headerIndexSections = new HeaderIndexSections(container, masks);
        var headerCombinedSections = new HeaderCombinedSections(headerIndexSections, headerSections, masks);
        var singles = new HeaderIndexOptSections(container, masks);
        var accessor = new CallTreeAndIndexesSnapshotAccessor(callTreeSections, headerCombinedSections, singles);
        var dfsReaders = new DotTraceDfsReaders(accessor);
        var totalPayloadReader = new NodePayloadBatchReaderTotal(accessor);
        var ownPayloadReader = new NodePayloadBatchReaderOwn(accessor);

        var metadataSection = (MetadataSection)new MetadataSectionContainer(container).Get(lifetime);
        var fuidConverter = new FuidToMetadataSectionContainer(container).Get();
        metadataSection.SetFuidConverter(fuidConverter);
        var assemblyProvider = metadataSection.GetMetadataAssemblyProvider(lifetime, loadReferencedAssemblies: false, guardMultiThreading: true);
        var resolver = new SignatureResolver(masks, fuidConverter, assemblyProvider);

        var headers = callTreeSections.AllHeaders().ToArray();
        var expectedNodeCount = headers.Sum(header =>
            checked((int)((header.HeaderFull.SectionSize - header.HeaderFull.SectionHeaderSize) / header.HeaderFull.RecordSize())));

        var nodes = new DfsNode<CallTreeSectionOffset, FunctionUID>[expectedNodeCount];
        var totals = new DotTracePayload[expectedNodeCount];
        var owns = new DotTracePayload[expectedNodeCount];

        var readNodes = dfsReaders.GetNodesReaders(lifetime).ReadForwardOffsetsAscending(accessor.MinOffset, expectedNodeCount, nodes, 0);
        var readTotals = totalPayloadReader.ReadForwardOffsetsAscending(accessor.MinOffset, expectedNodeCount, totals, 0);
        var readOwns = ownPayloadReader.ReadForwardOffsetsAscending(accessor.MinOffset, expectedNodeCount, owns, 0);

        if (readNodes != expectedNodeCount || readTotals != expectedNodeCount || readOwns != expectedNodeCount)
        {
            throw new InvalidOperationException(
                $"Snapshot read mismatch. nodes={readNodes}, totals={readTotals}, owns={readOwns}, expected={expectedNodeCount}.");
        }

        var nodeMap = new Dictionary<CallTreeSectionOffset, MutableNode>(expectedNodeCount);
        foreach (var index in Enumerable.Range(0, expectedNodeCount))
        {
            var node = nodes[index];
            var signature = resolver.Resolve(node.Key);
            nodeMap[node.Offset] = new MutableNode
            {
                Offset = node.Offset,
                ParentOffset = node.ParentOffset,
                Id = node.Offset.ToString(),
                Name = signature.Name,
                Kind = signature.Kind,
                InclusiveTime = GetTotalTime(totals[index]),
                ExclusiveTime = GetTotalTime(owns[index]),
                CallCount = GetCallCount(totals[index]),
                Children = []
            };
        }

        foreach (var index in Enumerable.Range(0, expectedNodeCount))
        {
            var node = nodes[index];
            if (!node.ParentOffset.IsValid)
                continue;

            if (nodeMap.TryGetValue(node.ParentOffset, out var parent))
                parent.Children.Add(nodeMap[node.Offset]);
        }

        var threadNodes = new List<SerializableNode>(headers.Length);
        foreach (var header in headers)
        {
            if (!nodeMap.TryGetValue(header.Root, out var rootNode))
                continue;

            var threadName = string.IsNullOrWhiteSpace(header.ThreadName)
                ? (string.IsNullOrWhiteSpace(header.GroupName) ? $"Thread {header.HeaderFull.GroupId}" : header.GroupName)
                : header.ThreadName;

            threadNodes.Add(new SerializableNode
            {
                Id = $"thread:{header.HeaderFull.GroupId}:{rootNode.Id}",
                Name = threadName,
                Kind = "thread",
                ThreadName = threadName,
                InclusiveTime = rootNode.InclusiveTime,
                ExclusiveTime = 0,
                CallCount = rootNode.CallCount,
                Children = [CloneTree(rootNode, [])]
            });
        }

        var syntheticRoot = new SerializableNode
        {
            Id = "root",
            Name = "<root>",
            Kind = "root",
            ThreadName = null,
            InclusiveTime = threadNodes.Sum(node => node.InclusiveTime),
            ExclusiveTime = 0,
            CallCount = threadNodes.Sum(node => node.CallCount),
            Children = threadNodes
        };

        var filteredTree = string.IsNullOrWhiteSpace(options.NameFilter)
            ? syntheticRoot
            : PruneTree(syntheticRoot, options.NameFilter) ?? CreateEmptyRoot();

        var hotspotCandidates = FilterHotspotCandidates(nodeMap.Values, options);
        var hotspots = BuildHotspots(hotspotCandidates, options.Top);
        var summary = BuildSummary(syntheticRoot, hotspotCandidates, hotspots.Exclusive.FirstOrDefault());
        var payload = new OutputDocument
        {
            Snapshot = new SnapshotInfo
            {
                Path = options.SnapshotPath,
                PayloadType = "time",
                TimeUnit = "nanoseconds",
                ThreadCount = threadNodes.Count,
                NodeCount = nodeMap.Count,
                Diagnostics = new SnapshotDiagnostics
                {
                    HeaderCount = headers.Length,
                    ExpectedNodeCount = expectedNodeCount,
                    ReadNodeCount = readNodes,
                    ReadTotalPayloadCount = readTotals,
                    ReadOwnPayloadCount = readOwns
                }
            },
            Summary = summary,
            ThreadRoots = filteredTree.Children.Select(node => new ThreadRootInfo
            {
                Id = node.Id,
                Name = node.Name,
                InclusiveTime = node.InclusiveTime
            }).ToList(),
            Hotspots = hotspots,
            HotPaths = options.FlatPaths ? BuildHotPaths(filteredTree, options) : null,
            CallTree = filteredTree
        };

        await JsonSerializer.SerializeAsync(Console.OpenStandardOutput(), payload, CreateJsonOptions());
        await Console.Out.WriteLineAsync();
        return 0;
    }
    catch (Exception exception)
    {
        Console.Error.WriteLine(exception);
        return 1;
    }
}

static ExtractorOptions? ParseArguments(string[] args)
{
    if (args.Length == 0)
        return null;

    var snapshotPath = Path.GetFullPath(args[0]);
    var top = 200;
    string? nameFilter = null;
    var flatPaths = false;
    var excludeIdle = true;

    for (var index = 1; index < args.Length; index++)
    {
        switch (args[index])
        {
            case "--top":
                if (++index >= args.Length || !int.TryParse(args[index], out top) || top <= 0)
                    throw new ArgumentException("--top requires a positive integer value.");
                break;
            case "--filter":
                if (++index >= args.Length)
                    throw new ArgumentException("--filter requires a value.");
                nameFilter = args[index];
                break;
            case "--flat":
            case "--paths":
                flatPaths = true;
                break;
            case "--include-idle":
                excludeIdle = false;
                break;
            case "--exclude-idle":
                excludeIdle = true;
                break;
            default:
                throw new ArgumentException($"Unknown argument: {args[index]}");
        }
    }

    return new ExtractorOptions(snapshotPath, top, nameFilter, flatPaths, excludeIdle);
}

static JsonSerializerOptions CreateJsonOptions() => new()
{
    MaxDepth = 4096,
    PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
    WriteIndented = true,
    DefaultIgnoreCondition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingNull
};

static void RegisterAssemblyResolver(string? dotTraceAppDir)
{
    var probeDir = string.IsNullOrWhiteSpace(dotTraceAppDir)
        ? Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), "Applications", "dotTrace.app", "Contents", "DotFiles")
        : dotTraceAppDir;

    AppDomain.CurrentDomain.AssemblyResolve += (_, eventArgs) =>
    {
        var assemblyName = new System.Reflection.AssemblyName(eventArgs.Name).Name;
        if (string.IsNullOrWhiteSpace(assemblyName))
            return null;

        var candidatePath = Path.Combine(probeDir, $"{assemblyName}.dll");
        return File.Exists(candidatePath) ? System.Reflection.Assembly.LoadFrom(candidatePath) : null;
    };
}

static long GetTotalTime(DotTracePayload payload) => payload.PlusTime - payload.MinusTime;

static long GetCallCount(DotTracePayload payload) => payload.PlusCallCount - payload.MinusCallCount;

static IEnumerable<MutableNode> FilterHotspotCandidates(IEnumerable<MutableNode> nodes, ExtractorOptions options) =>
    nodes.Where(node =>
        node.Kind == "method" &&
        (!options.ExcludeIdle || !IsIdleMethod(node.Name)) &&
        MatchesFilter(node.Name, options.NameFilter));

static HotspotLists BuildHotspots(IEnumerable<MutableNode> nodes, int top)
{
    var candidates = nodes
        .Select(node => new HotspotEntry
        {
            Id = node.Id,
            Name = node.Name,
            Kind = node.Kind,
            InclusiveTime = node.InclusiveTime,
            ExclusiveTime = node.ExclusiveTime,
            CallCount = node.CallCount
        })
        .ToArray();

    return new HotspotLists
    {
        Inclusive = candidates
            .OrderByDescending(node => node.InclusiveTime)
            .ThenBy(node => node.Name, StringComparer.Ordinal)
            .Take(top)
            .ToList(),
        Exclusive = candidates
            .OrderByDescending(node => node.ExclusiveTime)
            .ThenBy(node => node.Name, StringComparer.Ordinal)
            .Take(top)
            .ToList()
    };
}

static SummaryInfo BuildSummary(
    SerializableNode callTree,
    IEnumerable<MutableNode> hotspotCandidates,
    HotspotEntry? topExclusive)
{
    var candidateArray = hotspotCandidates.ToArray();
    return new SummaryInfo
    {
        WallTimeMs = NanosecondsToMilliseconds(callTree.Children.Count == 0 ? 0 : callTree.Children.Max(node => node.InclusiveTime)),
        ActiveTimeMs = NanosecondsToMilliseconds(candidateArray.Sum(node => node.ExclusiveTime)),
        TotalSamples = candidateArray.Sum(node => node.CallCount),
        TopExclusiveMethod = topExclusive?.Name,
        TopExclusiveMs = NanosecondsToMilliseconds(topExclusive?.ExclusiveTime ?? 0)
    };
}

static List<HotPathEntry> BuildHotPaths(SerializableNode root, ExtractorOptions options)
{
    var collector = new List<HotPathAccumulator>();
    foreach (var threadNode in root.Children)
    {
        CollectHotPaths(threadNode, [threadNode.Name], collector, options);
    }

    return collector
        .OrderByDescending(path => path.InclusiveTime)
        .ThenBy(path => path.Path, StringComparer.Ordinal)
        .Take(options.Top)
        .Select(path => new HotPathEntry
        {
            Path = path.Path,
            InclusiveTime = path.InclusiveTime,
            InclusiveMs = NanosecondsToMilliseconds(path.InclusiveTime),
            LeafExclusiveTime = path.LeafExclusiveTime,
            LeafExclusiveMs = NanosecondsToMilliseconds(path.LeafExclusiveTime)
        })
        .ToList();
}

static void CollectHotPaths(
    SerializableNode node,
    List<string> path,
    List<HotPathAccumulator> collector,
    ExtractorOptions options)
{
    foreach (var child in node.Children)
    {
        if (child.Kind == "method")
        {
            path.Add(child.Name);

            if ((!options.ExcludeIdle || !IsIdleMethod(child.Name)) && MatchesFilter(child.Name, options.NameFilter))
            {
                collector.Add(new HotPathAccumulator(
                    string.Join(" > ", path),
                    child.InclusiveTime,
                    child.ExclusiveTime));
            }

            CollectHotPaths(child, path, collector, options);
            path.RemoveAt(path.Count - 1);
            continue;
        }

        var appended = child.Kind is not ("root" or "special");
        if (appended)
            path.Add(child.Name);

        CollectHotPaths(child, path, collector, options);
        if (appended)
            path.RemoveAt(path.Count - 1);
    }
}

static SerializableNode CloneTree(MutableNode source, HashSet<CallTreeSectionOffset> ancestry)
{
    var nextAncestry = new HashSet<CallTreeSectionOffset>(ancestry) { source.Offset };
    var children = new List<SerializableNode>(source.Children.Count);

    foreach (var child in source.Children)
    {
        if (nextAncestry.Contains(child.Offset))
            continue;

        children.Add(CloneTree(child, nextAncestry));
    }

    return new SerializableNode
    {
        Id = source.Id,
        Name = source.Name,
        Kind = source.Kind,
        ThreadName = null,
        InclusiveTime = source.InclusiveTime,
        ExclusiveTime = source.ExclusiveTime,
        CallCount = source.CallCount,
        Children = children
    };
}

static SerializableNode? PruneTree(SerializableNode node, string? filter)
{
    if (string.IsNullOrWhiteSpace(filter))
        return node;

    if (MatchesFilter(node.Name, filter))
        return node;

    var prunedChildren = new List<SerializableNode>();
    foreach (var child in node.Children)
    {
        var prunedChild = PruneTree(child, filter);
        if (prunedChild is not null)
            prunedChildren.Add(prunedChild);
    }

    if (prunedChildren.Count == 0)
        return null;

    return node with { Children = prunedChildren };
}

static SerializableNode CreateEmptyRoot() => new()
{
    Id = "root",
    Name = "<root>",
    Kind = "root",
    ThreadName = null,
    InclusiveTime = 0,
    ExclusiveTime = 0,
    CallCount = 0,
    Children = []
};

static bool MatchesFilter(string value, string? filter) =>
    string.IsNullOrWhiteSpace(filter) ||
    value.Contains(filter, StringComparison.OrdinalIgnoreCase);

static bool IsIdleMethod(string name) =>
    GetIdleMethodPatterns().Any(pattern => name.Contains(pattern, StringComparison.Ordinal));

static double NanosecondsToMilliseconds(long value) => Math.Round(value / 1_000_000d, 3);

static string[] GetIdleMethodPatterns()
{
    return
    [
        "WaitHandle.Wait",
        "WaitOneNoCheck",
        "SemaphoreSlim.Wait",
        "LowLevelLifoSemaphore.Wait",
        "LowLevelLifoSemaphore.WaitForSignal",
        "LowLevelSpinWaiter.Wait",
        "Monitor.Wait",
        "SocketAsyncEngine.EventLoop",
        "PollGC",
        "Interop+Sys.WaitForSocketEvents",
        "ManualResetEventSlim.Wait",
        "Task.SpinThenBlockingWait",
        "Thread.SleepInternal"
    ];
}

sealed record ExtractorOptions(
    string SnapshotPath,
    int Top,
    string? NameFilter,
    bool FlatPaths,
    bool ExcludeIdle);

readonly record struct SignatureResult(string Name, string Kind);

readonly record struct HotPathAccumulator(string Path, long InclusiveTime, long LeafExclusiveTime);

sealed class SignatureResolver(
    SnapshotMasksComponent masks,
    FuidToMetadataIdConverter fuidConverter,
    IMetadataSectionAssemblyProvider assemblyProvider)
{
    private readonly Dictionary<FunctionUID, SignatureResult> _cache = [];

    public SignatureResult Resolve(FunctionUID fuid)
    {
        if (_cache.TryGetValue(fuid, out var cached))
            return cached;

        var result = ResolveCore(fuid);
        _cache[fuid] = result;
        return result;
    }

    private SignatureResult ResolveCore(FunctionUID fuid)
    {
        var fid = fuid.ToDotTraceFid(masks);
        if (fid == KnownFunctionId.RootId)
            return new SignatureResult("<thread-root>", "special");
        if (fid == KnownFunctionId.NativeCallId)
            return new SignatureResult("<native>", "special");
        if (fid == KnownFunctionId.ThreadStoppedByUnknownReasonCallId)
            return new SignatureResult("<thread-stopped>", "special");
        if (fid == KnownFunctionId.GcCallId)
            return new SignatureResult("<gc>", "special");
        if (fid == KnownFunctionId.AppDomainShutdownCallId)
            return new SignatureResult("<appdomain-shutdown>", "special");
        if (fid == KnownFunctionId.CodePitchingCallId)
            return new SignatureResult("<code-pitching>", "special");
        if (fid == KnownFunctionId.UnmanagedStackFrameCallId)
            return new SignatureResult("<unmanaged-stack-frame>", "special");
        if (fid == KnownFunctionId.UnsafeStackFrameId)
            return new SignatureResult("<unsafe-stack-frame>", "special");
        if (fid == KnownFunctionId.ContinuationsId)
            return new SignatureResult("<continuations>", "special");
        if (fid == KnownFunctionId.AwaitsId)
            return new SignatureResult("<awaits>", "special");
        if (fid == KnownFunctionId.TaskScheduledId)
            return new SignatureResult("<task-scheduled>", "special");
        if (fid == KnownFunctionId.TaskExecutionId)
            return new SignatureResult("<task-execution>", "special");
        if (fid == KnownFunctionId.TaskRecursionId)
            return new SignatureResult("<task-recursion>", "special");
        if (!KnownFunctionId.IsProcessable(fid))
            return new SignatureResult($"[special:0x{(uint)fid:X8}]", "special");

        try
        {
            var metadataId = fuidConverter.Convert(fid);
            if (metadataId.IsSynthetic)
                return new SignatureResult($"[synthetic:{metadataId}]", "special");

            var function = assemblyProvider.GetFunctionItemsLight(metadataId);
            var className = function.ClassToken != MetadataToken.Nil
                ? assemblyProvider.GetClassItems(MetadataId.Create(metadataId.MetadataIndex, function.ClassToken)).FullName
                : null;
            var qualifiedMethod = string.IsNullOrWhiteSpace(className)
                ? function.FunctionName
                : $"{className}.{function.FunctionName}";
            var methodName = string.IsNullOrEmpty(function.GenericParameters)
                ? qualifiedMethod
                : $"{qualifiedMethod}<{function.GenericParameters}>";
            return new SignatureResult(methodName, "method");
        }
        catch
        {
            return new SignatureResult($"[unresolved:0x{(uint)fid:X8}]", "special");
        }
    }
}

sealed class OutputDocument
{
    public required SnapshotInfo Snapshot { get; init; }

    public required SummaryInfo Summary { get; init; }

    public required List<ThreadRootInfo> ThreadRoots { get; init; }

    public required HotspotLists Hotspots { get; init; }

    public List<HotPathEntry>? HotPaths { get; init; }

    public required SerializableNode CallTree { get; init; }
}

sealed class SnapshotInfo
{
    public required string Path { get; init; }

    public required string PayloadType { get; init; }

    public required string TimeUnit { get; init; }

    public required int ThreadCount { get; init; }

    public required int NodeCount { get; init; }

    public required SnapshotDiagnostics Diagnostics { get; init; }
}

sealed class SnapshotDiagnostics
{
    public required int HeaderCount { get; init; }

    public required int ExpectedNodeCount { get; init; }

    public required int ReadNodeCount { get; init; }

    public required int ReadTotalPayloadCount { get; init; }

    public required int ReadOwnPayloadCount { get; init; }
}

sealed class SummaryInfo
{
    public required double WallTimeMs { get; init; }

    public required double ActiveTimeMs { get; init; }

    public required long TotalSamples { get; init; }

    public string? TopExclusiveMethod { get; init; }

    public required double TopExclusiveMs { get; init; }
}

sealed class ThreadRootInfo
{
    public required string Id { get; init; }

    public required string Name { get; init; }

    public required long InclusiveTime { get; init; }
}

sealed class HotspotLists
{
    public required List<HotspotEntry> Inclusive { get; init; }

    public required List<HotspotEntry> Exclusive { get; init; }
}

sealed class HotspotEntry
{
    public required string Id { get; init; }

    public required string Name { get; init; }

    public required string Kind { get; init; }

    public required long InclusiveTime { get; init; }

    public required long ExclusiveTime { get; init; }

    public required long CallCount { get; init; }
}

sealed class HotPathEntry
{
    public required string Path { get; init; }

    public required long InclusiveTime { get; init; }

    public required double InclusiveMs { get; init; }

    public required long LeafExclusiveTime { get; init; }

    public required double LeafExclusiveMs { get; init; }
}

sealed class MutableNode
{
    public required CallTreeSectionOffset Offset { get; init; }

    public required CallTreeSectionOffset ParentOffset { get; init; }

    public required string Id { get; init; }

    public required string Name { get; init; }

    public required string Kind { get; init; }

    public required long InclusiveTime { get; init; }

    public required long ExclusiveTime { get; init; }

    public required long CallCount { get; init; }

    public required List<MutableNode> Children { get; init; }
}

sealed record SerializableNode
{
    public required string Id { get; init; }

    public required string Name { get; init; }

    public required string Kind { get; init; }

    public string? ThreadName { get; init; }

    public required long InclusiveTime { get; init; }

    public required long ExclusiveTime { get; init; }

    public required long CallCount { get; init; }

    public required List<SerializableNode> Children { get; init; }
}