/* This software is licensed by the MIT License, see LICENSE file */
/*                           Copyright © 2024-2025 Gregory Lirent */

#if WINIPC_UA_SERVER
using System.Runtime.InteropServices;
using WinIPC.Models;

namespace WinIPC.Utils
{
    public class InputHandler
    {
        private Dictionary<IntPtr, Dictionary<Button, bool>> _state = new();

        [DllImport("user32.dll", SetLastError = true)]
        private static extern uint SendInput(uint nInputs, IntPtr pInputs, int cbSize);

        [DllImport("user32.dll", SetLastError = true)]
        private static extern IntPtr SendMessage(IntPtr hwnd, int msg, IntPtr wParam, IntPtr lParam);

        [DllImport("user32.dll")]
        private static extern short GetKeyState(int nVirtKey);

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

        private bool IsKeyPressed(int keyCode)
        {
            return (GetKeyState(keyCode) & 0x8000) != 0;
        }

        private ushort GetMouseKeyStates()
        {
            ushort keyStates = 0;
            if (IsKeyPressed(0x01)) keyStates |= 0x0001; // lbutton
            if (IsKeyPressed(0x02)) keyStates |= 0x0002; // rbutton
            if (IsKeyPressed(0x04)) keyStates |= 0x0010; // mbutton
            if (IsKeyPressed(0x05)) keyStates |= 0x0020; // xbutton1
            if (IsKeyPressed(0x06)) keyStates |= 0x0040; // xbutton2
            if (IsKeyPressed(0x10)) keyStates |= 0x0004; // Shift
            if (IsKeyPressed(0x11)) keyStates |= 0x0008; // Ctrl
            return keyStates;
        }

        [StructLayout(LayoutKind.Sequential)]
        private struct MOUSEINPUT
        {
            public int dx;
            public int dy;
            public int mouseData;
            public uint dwFlags;
            public uint time;
            public IntPtr dwExtraInfo;
        }

        [StructLayout(LayoutKind.Sequential)]
        private struct KEYBDINPUT
        {
            public ushort wVk;
            public ushort wScan;
            public uint dwFlags;
            public uint time;
            public IntPtr dwExtraInfo;
        }

        [StructLayout(LayoutKind.Explicit)]
        private struct INPUT
        {
            [FieldOffset(0)]
            public int type;
            [FieldOffset(8)]
            public MOUSEINPUT mi;
            [FieldOffset(8)]
            public KEYBDINPUT ki;
        }
        private static IntPtr MakeParam(int low, int high = 0)
        {
            return (high << 16) | (low & 0xffff);
        }

        private static IntPtr GetCursorPos(IntPtr hwnd)
        {
            WindowScanner.GetCursorPos(out var pos);

            if (hwnd != IntPtr.Zero)
            {
                WindowScanner.ScreenToClient(hwnd, ref pos);
            }

            return MakeParam(pos.X, pos.Y);
        }

        private static readonly int INPUT_SIZE = Marshal.SizeOf(typeof(INPUT));

        private Dictionary<Button, bool> GetStateDict(IntPtr hwnd)
        {
            Dictionary<Button, bool> state;
            if (_state.TryGetValue(hwnd, out var s))
            {
                state = s;
            }
            else { _state.Add(hwnd, state = new()); }

            return state;
        }

        private async Task<bool> ProcessMouseInput(IntPtr hwnd, Button button, bool press)
        {
            bool success = false;
            var state = GetStateDict(hwnd);
            if (state.TryGetValue(button, out var s) && s == press) return true;

            return await Task.Run(async () =>
            {
                if (hwnd == IntPtr.Zero)
                {
                    var buffer = Marshal.AllocHGlobal(INPUT_SIZE);
                    var input = new INPUT
                    {
                        type = 0x00,
                        mi   = new()
                    };

                    if (button.HasFlag(Button.MouseExtraFlag))
                    {
                        input.mi.dwFlags = (uint)Button.MouseExtraFlag << (press ? 0 : 1);
                        input.mi.mouseData = (int)(button & ~Button.MouseExtraFlag);
                    }
                    else
                    {
                        input.mi.dwFlags = (uint)button << (press ? 0 : 1);
                    }

                    Marshal.StructureToPtr(input, buffer, true);
                    success = await Task.Run(() => { return 1 != SendInput(1, buffer, INPUT_SIZE); });
                    Marshal.FreeHGlobal(buffer);
                }
                else
                {
                    int msg = 0;
                    int wFlags = 0;

                    switch(button)
                    {
                        case Button.MouseLeft:
                            msg = press ? 0x0201 : 0x0202;
                            break;
                        case Button.MouseRight:
                            msg = press ? 0x0204 : 0x0205;
                            break;
                        case Button.MouseMiddle:
                            msg = press ? 0x0207 : 0x0208;
                            break;
                        case Button.MouseExtra1:
                        case Button.MouseExtra2:
                            msg = press ? 0x020b : 0x020c;
                            wFlags = (int)button & 0x7f;
                            break;
                        default:
                            // TODO Log
                            break;
                    }

                    SendMessage(hwnd, msg, MakeParam(GetMouseKeyStates(), wFlags), GetCursorPos(hwnd));
                    success = true;
                }

                if (success) state[button] = press;
                return success;
            });
        }

        private async Task<bool> ProcessKeyboardInput(IntPtr hwnd, Button button, bool press)
        {
            bool success = false;
            var state = GetStateDict(hwnd);
            if (state.TryGetValue(button, out var s) && s == press) return true;

            return await Task.Run<bool>(async () =>
            {
                if (hwnd == IntPtr.Zero)
                {
                    var buffer = Marshal.AllocHGlobal(INPUT_SIZE);
                    var flags  = (uint)(press ? 0x00 : 0x02);
                    var input  = new INPUT
                    {
                        type = 0x01,
                        ki = new()
                    };

                    if (button.HasFlag(Button.ExtendedKeyFlag))
                    {
                        flags |= 0x01;
                        button &= ~Button.ExtendedKeyFlag;
                    }

                    input.ki.wVk = (ushort)button;
                    input.ki.dwFlags = flags;

                    Marshal.StructureToPtr(input, buffer, true);
                    success = await Task.Run(() => { return 1 != SendInput(1, buffer, INPUT_SIZE); });
                    Marshal.FreeHGlobal(buffer);
                }
                else
                {
                    var btn = (int)(button & ~Button.ExtendedKeyFlag);
                    var msg = press ? 0x0100 : 0x0101;
                    
                    if (button == Button.LAlt || button == Button.RAlt) msg += 4;

                    int flags = MapVirtualKey(btn, 0) | (button.HasFlag(Button.ExtendedKeyFlag) ? 0x0100 : 0) |
                        (IsKeyPressed(0x12) ? 0x2000 : 0) | (s ? 0x4000 : 0) | (press ? 0 : 0x8000);

                    SendMessage(hwnd, msg, MakeParam(btn), MakeParam(1, flags));
                    success = true;
                }

                if (success) state[button] = press;
                return success;
            });
        }

        public async Task<bool> KeyUp(IntPtr hwnd, Button button, int delay)
        {
            Task<bool> pTask;
            bool success;

            if (button.HasFlag(Button.MouseKeyFlag))
            {
                pTask = ProcessMouseInput(hwnd, button & ~Button.MouseKeyFlag, false);
            }
            else
            {
                pTask = ProcessKeyboardInput(hwnd, button, false);
            }

            if (success = await pTask)
            {
                await Task.Delay(delay);
            }

            return success;
        }

        public async Task<bool> KeyDown(IntPtr hwnd, Button button, int delay)
        {
            Task<bool> pTask;
            bool success;

            if (button.HasFlag(Button.MouseKeyFlag))
            {
                pTask = ProcessMouseInput(hwnd, button & ~Button.MouseKeyFlag, true);
            }
            else
            {
                pTask = ProcessKeyboardInput(hwnd, button, true);
            }

            if (success = await pTask)
            {
                await Task.Delay(delay);
            }

            return success;
        }

        public async Task<bool> MouseScroll(IntPtr hwnd, int offset, int delay)
        {
            bool success = false;

            if (hwnd == IntPtr.Zero)
            {
                var buffer = Marshal.AllocHGlobal(INPUT_SIZE);
                var input = new INPUT
                {
                    type = 0x00,
                    mi = new MOUSEINPUT
                    {
                        dwFlags = 0x0800,
                        mouseData = offset
                    }
                };

                Marshal.StructureToPtr(input, buffer, true);
                success = await Task.Run(() => { return 1 != SendInput(1, buffer, INPUT_SIZE); });
                Marshal.FreeHGlobal(buffer);
            }
            else
            {
                SendMessage(hwnd, 0x020a, MakeParam(GetMouseKeyStates(), offset), GetCursorPos(IntPtr.Zero));
                success = true;
            }

            await Task.Delay(delay);

            return success;
        }

        public async Task<bool> MouseMove(IntPtr hwnd, int x, int y, int delay)
        {
            bool success = false;

            if (hwnd == IntPtr.Zero)
            {
                var buffer = Marshal.AllocHGlobal(INPUT_SIZE);
                var input = new INPUT
                {
                    type = 0x00,
                    mi = new MOUSEINPUT
                    {
                        dwFlags = 0x8001,
                        dx = x,
                        dy = y
                    }
                };

                Marshal.StructureToPtr(input, buffer, true);
                success = await Task.Run(() => { return 1 != SendInput(1, buffer, INPUT_SIZE); });
                Marshal.FreeHGlobal(buffer);
            } 
            else
            {
                SendMessage(hwnd, 0x0200, MakeParam(GetMouseKeyStates()), MakeParam(x, y));
                success = true;
            }

            await Task.Delay(delay);

            return success;
        }
    }
}
#endif