using Insight.Agent.Interfaces;
using Insight.Agent.Messages;
using System.Runtime.InteropServices;
using System.Runtime.Versioning;

namespace Insight.Agent.Network.Handlers
{
    [SupportedOSPlatform("windows")]
    public class SessionHandler : IAgentMessageHandler<AgentSession>
    {
        public async ValueTask HandleAsync<TMessage>(AgentSession sender, TMessage message, CancellationToken cancellationToken) where TMessage : IAgentMessage
        {
            if (message is GetInventory)
            {
                var result = new SessionList();
                result.AddRange(GetSessions());

                await sender.SendAsync(result, cancellationToken);
            }
        }

        private static List<Session> GetSessions()
        {
            var query = NativeMethods.GetSessions();

            var sessions = new List<Session>();

            foreach (var s in NativeMethods.GetSessions())
            {
                sessions.Add(new Session
                {
                    Sid = s.SessionId.ToString(),
                    User = s.Username,
                    Type = s.Workstation,
                    Status = s.State.ToString(),
                    Remote = s.IPAddress
                });
            }

            return sessions;
        }

        private static partial class NativeMethods
        {
            //public const int WTS_CURRENT_SESSION = -1;

            [DllImport("wtsapi32.dll")]
            static extern int WTSEnumerateSessions(
            nint pServer,
            [MarshalAs(UnmanagedType.U4)] int iReserved,
            [MarshalAs(UnmanagedType.U4)] int iVersion,
            ref nint pSessionInfo,
            [MarshalAs(UnmanagedType.U4)] ref int iCount);

            [DllImport("Wtsapi32.dll")]
            private static extern bool WTSQuerySessionInformation(
                nint pServer,
                int iSessionID,
                WTS_INFO_CLASS oInfoClass,
                out nint pBuffer,
                out uint iBytesReturned);

            [DllImport("wtsapi32.dll")]
            static extern void WTSFreeMemory(
                nint pMemory);

            [StructLayout(LayoutKind.Sequential)]
            private struct WTS_CLIENT_ADDRESS
            {
                public int iAddressFamily;
                [MarshalAs(UnmanagedType.ByValArray, SizeConst = 20)]
                public byte[] bAddress;
            }

            [StructLayout(LayoutKind.Sequential)]
            private struct WTS_SESSION_INFO
            {
                public int iSessionID;
                [MarshalAs(UnmanagedType.LPStr)]
                public string sWinsWorkstationName;
                public WTS_CONNECTSTATE_CLASS oState;
            }

            [StructLayout(LayoutKind.Sequential)]
            private struct WTS_CLIENT_DISPLAY
            {
                public int iHorizontalResolution;
                public int iVerticalResolution;
                //1 = The display uses 4 bits per pixel for a maximum of 16 colors.
                //2 = The display uses 8 bits per pixel for a maximum of 256 colors.
                //4 = The display uses 16 bits per pixel for a maximum of 2^16 colors.
                //8 = The display uses 3-byte RGB values for a maximum of 2^24 colors.
                //16 = The display uses 15 bits per pixel for a maximum of 2^15 colors.
                public int iColorDepth;
            }

            public enum WTS_CONNECTSTATE_CLASS
            {
                WTSActive,
                WTSConnected,
                WTSConnectQuery,
                WTSShadow,
                WTSDisconnected,
                WTSIdle,
                WTSListen,
                WTSReset,
                WTSDown,
                WTSInit
            }

            public enum WTS_INFO_CLASS
            {
                WTSInitialProgram,
                WTSApplicationName,
                WTSWorkingDirectory,
                WTSOEMId,
                WTSSessionId,
                WTSUserName,
                WTSWinStationName,
                WTSDomainName,
                WTSConnectState,
                WTSClientBuildNumber,
                WTSClientName,
                WTSClientDirectory,
                WTSClientProductId,
                WTSClientHardwareId,
                WTSClientAddress,
                WTSClientDisplay,
                WTSClientProtocolType,
                WTSIdleTime,
                WTSLogonTime,
                WTSIncomingBytes,
                WTSOutgoingBytes,
                WTSIncomingFrames,
                WTSOutgoingFrames,
                WTSClientInfo,
                WTSSessionInfo,
                WTSConfigInfo,
                WTSValidationInfo,
                WTSSessionAddressV4,
                WTSIsRemoteSession
            }

            public class WTSSession
            {
                public int SessionId { get; set; }
                public WTS_CONNECTSTATE_CLASS State { get; set; }
                public string? Workstation { get; set; }
                public string? IPAddress { get; set; }
                public string? Username { get; set; }
                public int HorizontalResolution { get; set; }
                public int VerticalResolution { get; set; }
                public int ColorDepth { get; set; }
                public string? ClientApplicationDirectory { get; set; }
            }

            public static IEnumerable<WTSSession> GetSessions()
            {
                var sessions = new List<WTSSession>();

                var pServer = nint.Zero;
                var pSessionInfo = nint.Zero;

                try
                {
                    var count = 0;
                    var sessionCount = WTSEnumerateSessions(pServer, 0, 1, ref pSessionInfo, ref count);
                    var dataSize = (long)Marshal.SizeOf(typeof(WTS_SESSION_INFO));
                    var current = (long)pSessionInfo;

                    if (sessionCount <= 0)
                        return sessions;

                    for (int i = 0; i < count; i++)
                    {
                        if (Marshal.PtrToStructure((nint)current, typeof(WTS_SESSION_INFO)) is not object sessionStructure)
                            continue;

                        var sessionInfo = (WTS_SESSION_INFO)sessionStructure;
                        current += dataSize;

                        var session = new WTSSession
                        {
                            SessionId = sessionInfo.iSessionID,
                            State = sessionInfo.oState,
                            Workstation = sessionInfo.sWinsWorkstationName,
                        };

                        var returned = (uint)0;

                        // get terminal user address
                        var address = nint.Zero;
                        var clientAddress = new WTS_CLIENT_ADDRESS();

                        if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSClientAddress, out address, out returned) == true)
                        {
                            if (Marshal.PtrToStructure(address, clientAddress.GetType()) is not object addressStructure)
                                break;

                            clientAddress = (WTS_CLIENT_ADDRESS)addressStructure;
                            session.IPAddress = clientAddress.bAddress[2] + "." + clientAddress.bAddress[3] + "." + clientAddress.bAddress[4] + "." + clientAddress.bAddress[5];
                        }

                        // get terminal user name
                        if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSUserName, out address, out returned) == true)
                        {
                            session.Username = Marshal.PtrToStringAnsi(address);
                        }

                        // get terminal user domain name
                        if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSDomainName, out address, out returned) == true)
                        {
                            session.Username = Marshal.PtrToStringAnsi(address) + @"\" + session.Username;
                        }

                        // get terminal user display informations
                        var clientDisplay = new WTS_CLIENT_DISPLAY();

                        if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSClientDisplay, out address, out returned) == true)
                        {
                            if (Marshal.PtrToStructure(address, clientDisplay.GetType()) is not object displayStructure)
                                break;

                            clientDisplay = (WTS_CLIENT_DISPLAY)displayStructure;
                            session.HorizontalResolution = clientDisplay.iHorizontalResolution;
                            session.VerticalResolution = clientDisplay.iVerticalResolution;
                            session.ColorDepth = clientDisplay.iColorDepth;
                        }

                        // get terminal user application directory
                        if (WTSQuerySessionInformation(pServer, sessionInfo.iSessionID, WTS_INFO_CLASS.WTSClientDirectory, out address, out returned) == true)
                        {
                            session.ClientApplicationDirectory = Marshal.PtrToStringAnsi(address);
                        }

                        sessions.Add(session);
                    }
                }
                catch (Exception)
                {
                    throw;
                }
                finally
                {
                    WTSFreeMemory(pSessionInfo);
                }

                return sessions;
            }
        }
    }
}