//
// Copyright (c) 2022-present, Trail of Bits, Inc.
// All rights reserved.
//
// This source code is licensed in accordance with the terms specified in
// the LICENSE file found in the root directory of this source tree.
//
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using RpcInvestigator.Util;
using System.Windows.Forms.Integration;
using GraphX.Common.Enums;
using GraphX.Logic.Algorithms.OverlapRemoval;
using GraphX.Logic.Models;
using GraphX.Controls;
using GraphX.Controls.Models;
using QuickGraph;
using GraphX.Common.Models;
using System.Windows;
using MessageBox = System.Windows.Forms.MessageBox;
using GraphX.Logic.Algorithms.LayoutAlgorithms;
using Brushes = System.Windows.Media.Brushes;
using Style = System.Windows.Style;
using Binding = System.Windows.Data.Binding;
using Control = System.Windows.Controls.Control;
using GraphX.Common.Interfaces;

namespace RpcInvestigator.Windows.Controls
{
    public class SnifferGraph
    {
        private readonly RpcLibrary m_Library;
        private readonly SnifferCallbackTable m_Callbacks;
        private ZoomControl m_Zoom;
        private SnifferGraphArea m_GraphArea;
        private Dictionary<long, string> m_PidLookupCache;

        public SnifferGraph(
            ElementHost Host,
            RpcLibrary Library,
            SnifferCallbackTable Callbacks
            )
        {
            m_PidLookupCache = new Dictionary<long, string>();
            m_Callbacks = Callbacks;
            m_Library = Library;

            m_Zoom = new ZoomControl();
            Host.Child = m_Zoom;
            m_Zoom.ZoomToFill();
            ZoomControl.SetViewFinderVisibility(m_Zoom, System.Windows.Visibility.Visible);
            var logic = new GXLogicCore<SnifferNode, SnifferEdge, BidirectionalSnifferGraph>();
            m_GraphArea = new SnifferGraphArea
            {
                LogicCore = logic,
                EdgeLabelFactory = new DefaultEdgelabelFactory()
            };
            m_GraphArea.LogicCore.Graph = new BidirectionalSnifferGraph();
            logic.DefaultLayoutAlgorithm = LayoutAlgorithmTypeEnum.LinLog;
            var layoutParams = (LinLogLayoutParameters)
                logic.AlgorithmFactory.CreateLayoutParameters(LayoutAlgorithmTypeEnum.LinLog);
            logic.DefaultLayoutAlgorithmParams = layoutParams;
            logic.DefaultOverlapRemovalAlgorithm = OverlapRemovalAlgorithmTypeEnum.FSA;
            logic.DefaultOverlapRemovalAlgorithmParams = logic.AlgorithmFactory.CreateOverlapRemovalParameters(OverlapRemovalAlgorithmTypeEnum.FSA);
            ((OverlapRemovalParameters)logic.DefaultOverlapRemovalAlgorithmParams).HorizontalGap = 20;
            ((OverlapRemovalParameters)logic.DefaultOverlapRemovalAlgorithmParams).VerticalGap = 20;
            logic.DefaultEdgeRoutingAlgorithm = EdgeRoutingAlgorithmTypeEnum.None;
            logic.AsyncAlgorithmCompute = false;
            m_Zoom.Content = m_GraphArea;
            OverrideTemplateStyles();
            m_GraphArea.RelayoutFinished += SnifferGraphArea_RelayoutFinished;
            m_GraphArea.VertexClicked += NodeClicked;
            ToggleVisibility(false);
        }

        void SnifferGraphArea_RelayoutFinished(object sender, EventArgs e)
        {
            m_Zoom.ZoomToFill();
        }

        private void NodeClicked(object sender, VertexClickedEventArgs e)
        {
            var node = e.Control.Vertex as SnifferNode;
            if (node.UserData == null || node.UserData.GetType() != typeof(SnifferNodeUserData))
            {
                return;
            }
            if (!Guid.TryParse(node.UserData.InterfaceId, out Guid parsed))
            {
                MessageBox.Show("The InterfaceUuid has an invalid format.");
                return;
            }
            m_Callbacks.ShowRpcServerDetailsCallback(parsed.ToString());
        }

        private
        void
        OverrideTemplateStyles()
        {
            //
            // Ellipsis node type
            //
            var ellipsisNodeTrigger = new DataTrigger()
            {
                Binding = new Binding("NodeType"),
                Value = SnifferNodeType.Ellipsis
            };
            ellipsisNodeTrigger.Setters.Add(new Setter()
            {
                Property = Control.BackgroundProperty,
                Value = Brushes.LightBlue
            });
            //
            // RPC server node type
            //
            var rpcServerNodeTrigger = new DataTrigger()
            {
                Binding = new Binding("NodeType"),
                Value = SnifferNodeType.RpcServer
            };
            rpcServerNodeTrigger.Setters.Add(new Setter()
            {
                Property = Control.BackgroundProperty,
                Value = Brushes.Red
            });
            rpcServerNodeTrigger.Setters.Add(new Setter()
            {
                Property = VertexControl.VertexShapeProperty,
                Value = VertexShape.Rectangle
            });
            //
            // Process node type
            //
            var processNodeTrigger = new DataTrigger()
            {
                Binding = new Binding("NodeType"),
                Value = SnifferNodeType.Process
            };
            processNodeTrigger.Setters.Add(new Setter()
            {
                Property = Control.BackgroundProperty,
                Value = Brushes.Blue
            });
            processNodeTrigger.Setters.Add(new Setter()
            {
                Property = Control.ForegroundProperty,
                Value = Brushes.White
            });

            //
            // There's probably a better way, but this works.
            //
            m_Zoom.Resources.MergedDictionaries.Add(new ResourceDictionary
            {
                Source = new Uri("Windows/Controls/GraphTemplate.xaml", UriKind.Relative)
            });
            var style = (Style)m_Zoom.Resources.MergedDictionaries[0][typeof(VertexControl)];
            style.Setters.Clear();
            style.Triggers.Clear();
            style.Triggers.Add(ellipsisNodeTrigger);
            style.Triggers.Add(rpcServerNodeTrigger);
            style.Triggers.Add(processNodeTrigger);
        }

        public
        void
        ToggleVisibility(bool Visible)
        {
            m_GraphArea.Visibility = Visible ? Visibility.Visible : Visibility.Hidden;
        }

        public void Update(List<ParsedEtwEvent> Events)
        {
            foreach (var evt in Events)
            {
                //
                // Ignore events that don't contain an interface ID for the RPC server
                // involved or activity ID.  This greatly reduces noise and the
                // back-and-forth parts of a typical RPC call which show up here as
                // multiple ETW events but represent one logical call.
                // Also, ignore this tool as well.
                //
                if (!evt.UserDataProperties.ContainsKey("InterfaceUuid") ||
                    evt.ActivityId == null || evt.ActivityId == Guid.Empty ||
                    evt.ProcessId == Process.GetCurrentProcess().Id)
                {
                    continue;
                }

                //
                // Only draw nodes/edge for an activity event group one time.
                //
                var activityId = evt.ActivityId.ToString();
                var edgeId = activityId;
                var existingActivityEdge = m_GraphArea.EdgesList.Where(
                    e => e.Key.SnifferId == activityId);
                if (existingActivityEdge.Count() != 0)
                {
                    continue;
                }

                var rpcServerUuid = evt.UserDataProperties["InterfaceUuid"];
                var existing = m_GraphArea.VertexList.Where(
                    v => v.Key.SnifferId == rpcServerUuid);
                SnifferNode rpcServerNode;

                if (existing.Count() == 0)
                {
                    var node = new SnifferNode()
                    {
                        NodeType = SnifferNodeType.RpcServer,
                        SnifferId = rpcServerUuid,
                        UserData = new SnifferNodeUserData()
                        {
                            InterfaceId = rpcServerUuid,
                            EdgeCount = 0,
                        },
                        Text = rpcServerUuid,
                    };
                    var servers = m_Library.Get(new Guid(rpcServerUuid));
                    if (servers != null)
                    {
                        //
                        // TODO: Icon to show this info came from our RPC library.
                        // Also, we should add ALPC port information if it's available.
                        //
                        var friendly = "";
                        if (!string.IsNullOrEmpty(servers[0].ServiceName))
                        {
                            friendly = servers[0].ServiceName + Environment.NewLine;
                        }
                        else if (!string.IsNullOrEmpty(servers[0].FilePath))
                        {
                            friendly = servers[0].FilePath + Environment.NewLine;
                        }
                        node.Text = friendly + rpcServerUuid;
                    }
                    var control = new VertexControl(node);
                    m_GraphArea.AddVertexAndData(node, control);
                    rpcServerNode = node;
                }
                else
                {
                    rpcServerNode = existing.ToList()[0].Key;
                }

                //
                // Avoid edge explosion for chatty RPC servers.
                // TODO: Make this customizable.
                //
                if (rpcServerNode.UserData.EdgeCount > 10)
                {
                    AddOrUpdateEllipsisNode(rpcServerNode);
                    continue;
                }
                ConnectRpcNode(activityId, rpcServerNode, evt);
            }
            m_GraphArea.SetEdgesDrag(false);
            m_GraphArea.SetVerticesDrag(false);
            m_GraphArea.GenerateGraph(true);
            m_GraphArea.ShowAllEdgesLabels(false);
            m_Zoom.ZoomToFill();
        }

        private void ConnectRpcNode(
            string ActivityId,
            SnifferNode RpcServerNode,
            ParsedEtwEvent Event
            )
        {
            //
            // Increment the node's edge count to prevent edge explosion.
            //
            RpcServerNode.UserData.EdgeCount++;

            //
            // TODO: The PID could have been reused. We should probably
            // build and maintain a best-effort mapping.
            //
            string processName = null;
            if (!m_PidLookupCache.ContainsKey(Event.ProcessStartKey))
            {
                try
                {
                    Process p = Process.GetProcessById((int)Event.ProcessId);
                    processName = p.ProcessName + "(" + Event.ProcessId + ")";
                }
                catch (Exception)
                {
                    processName = "[" + Event.ProcessId.ToString() + "]";
                }
                m_PidLookupCache.Add(Event.ProcessStartKey, processName);
            }
            else
            {
                processName = m_PidLookupCache[Event.ProcessStartKey];
            }

            var existingProcessNode = m_GraphArea.VertexList.Where(
                v => v.Key.SnifferId == processName);
            SnifferNode processNode;

            if (existingProcessNode.Count() == 0)
            {
                var node = new SnifferNode()
                {
                    NodeType = SnifferNodeType.Process,
                    SnifferId = processName,
                    Text = processName,
                };
                var control = new VertexControl(node);
                m_GraphArea.AddVertexAndData(node, control);
                processNode = node;
            }
            else
            {
                processNode = existingProcessNode.ToList()[0].Key;
            }

            var processNodeControl = m_GraphArea.VertexList[processNode];

            //
            // RPC etw events have tasks (RpcClientCall, RpcServerCall)
            // and opcodes (Start, Stop) which would result in multiple
            // redundant edges between a process and an RPC server. We
            // have already filtered by ActivityId, so we should never
            // have to check that an edge exists here.
            //
            var rpcServerNodeControl = m_GraphArea.VertexList[RpcServerNode];
            var edgeId = ActivityId;
            var edgeFromRpcServerNodeToProcessNode =
                new SnifferEdge(processNode, RpcServerNode)
            {
                SnifferId = edgeId,
                RpcNodeSnifferId = RpcServerNode.SnifferId,
            };
            var edgeControl = new EdgeControl(processNodeControl,
                    rpcServerNodeControl,
                    edgeFromRpcServerNodeToProcessNode);
            m_GraphArea.AddEdgeAndData(edgeFromRpcServerNodeToProcessNode, edgeControl);
        }

        private void AddOrUpdateEllipsisNode(SnifferNode RpcServerNode)
        {
            if (RpcServerNode.UserData.EllipsisNode != null)
            {
                var count = RpcServerNode.UserData.EdgeCount++;
                RpcServerNode.UserData.EllipsisNode.Text = "+" + count + " more...";
                return;
            }

            //
            // No ellipsis node exists for this RPC server. Create it now.
            //
            var nodeId = RpcServerNode.UserData.InterfaceId + "-ellipsis";
            var node = new SnifferNode()
            {
                NodeType = SnifferNodeType.Ellipsis,
                SnifferId = nodeId,
                UserData = new SnifferNodeUserData(),
                Text = "+1 more..."
            };
            var control = new VertexControl(node);
            m_GraphArea.AddVertexAndData(node, control);
            //
            // Connect it to the RPC server node.
            //
            var rpcServerNodeControl = m_GraphArea.VertexList[RpcServerNode];
            var edgeFromRpcServerNodeToEllipsisNode =
                new SnifferEdge(node, RpcServerNode);
            m_GraphArea.AddEdgeAndData(
                edgeFromRpcServerNodeToEllipsisNode,
                new EdgeControl(
                    control,
                    rpcServerNodeControl,
                    edgeFromRpcServerNodeToEllipsisNode));
            RpcServerNode.UserData.EllipsisNode = node;
        }

        public void Reset()
        {
            m_GraphArea.RemoveAllVertices();
            m_GraphArea.RemoveAllEdges();
        }

        public void SaveAsImage()
        {
            if (m_GraphArea.LogicCore.Graph.IsVerticesEmpty)
            {
                return;
            }
            m_GraphArea.ExportAsImageDialog(ImageType.PNG);
        }
    }

    public class SnifferGraphArea : GraphArea<SnifferNode, SnifferEdge, BidirectionalSnifferGraph> { }

    public class BidirectionalSnifferGraph : BidirectionalGraph<SnifferNode, SnifferEdge> { }

    public enum SnifferNodeType
    {
        Ellipsis,
        Process,
        RpcServer
    }

    public class SnifferNode : VertexBase
    {
        public string SnifferId { get; set; }
        public SnifferNodeUserData UserData { get; set; }
        public string Text { get; set; }
        public SnifferNodeType NodeType { get; set; }

        public SnifferNode() { }

        public override string ToString()
        {
            return Text;
        }
    }

    public class SnifferNodeUserData
    {
        public string InterfaceId;
        public int EdgeCount;
        public SnifferNode EllipsisNode;
    }

    public class SnifferEdge : EdgeBase<SnifferNode>
    {
        public string SnifferId { get; set; }
        public string RpcNodeSnifferId { get; set; }
        public object UserData { get; set; }
        public string Text { get; set; }

        public SnifferEdge(SnifferNode source, SnifferNode target, double weight = 1)
            : base(source, target, weight)
        {
        }

        public SnifferEdge()
            : base(null, null, 1)
        {
        }

        public override string ToString()
        {
            return Text;
        }
    }
}