import com.jniwrapper.*;
import com.jniwrapper.win32.Handle;
import com.jniwrapper.win32.IntPtr;
import com.jniwrapper.win32.LastErrorException;
import com.jniwrapper.win32.ui.*;

import javax.swing.*;
import java.awt.event.WindowAdapter;
import java.awt.event.WindowEvent;

/**
 * This example demonstrates how to implement key logging
 * using <a href="https://msdn.microsoft.com/en-us/library/windows/desktop/ms645536%28v=vs.85%29.aspx">Raw Input</a> Windows API.
 */
public class RawKeyboardInputSample
{
    /**
     * Wrapper for <a href="https://msdn.microsoft.com/en-us/library/windows/desktop/ms645575%28v=vs.85%29.aspx">RAWKEYBOARD</a> structure.
     */
    static class RAWKEYBOARD extends Structure {
        UShortInt MakeCode = new UShortInt();
        UShortInt Flags = new UShortInt();
        UShortInt Reserved = new UShortInt();
        UShortInt VKey = new UShortInt();
        UInt Message = new UInt();
        ULongInt ExtraInformation = new ULongInt();
        public RAWKEYBOARD() {
            init(new Parameter[] {MakeCode, Flags, Reserved, VKey, Message, ExtraInformation}, (short) 8);
        }
    }

    /**
     * Wrapper for <a href="https://msdn.microsoft.com/en-us/library/windows/desktop/ms645565%28v=vs.85%29.aspx">RAWINPUTDEVICE</a> structure.
     */
    static class RAWINPUTDEVICE extends Structure {
        UShortInt usUsagePage = new UShortInt();
        UShortInt usUsage = new UShortInt();
        UInt32 dwFlags = new UInt32();
        Wnd hwndTarget = new Wnd();

        public RAWINPUTDEVICE() {
            init(new Parameter[] {usUsagePage, usUsage, dwFlags, hwndTarget}, (short) 8);
        }
    }

    /**
     * Wrapper for <a href="https://msdn.microsoft.com/en-us/library/windows/desktop/ms645571%28v=vs.85%29.aspx">RAWINPUTHEADER</a> structure.
     */
    static class RAWINPUTHEADER extends Structure {
        UInt32 dwType = new UInt32();
        UInt32 dwSize = new UInt32();
        Handle hDevice = new Handle();
        IntPtr wParam = new IntPtr();

        public RAWINPUTHEADER() {
            init(new Parameter[] {dwType, dwSize, hDevice, wParam}, (short) 8);
        }
    }

    /**
     * Wrapper for <a href="https://msdn.microsoft.com/en-us/library/windows/desktop/ms645562%28v=vs.85%29.aspx">RAWINPUT</a> structure.
     */
    static class RAWINPUT extends Structure {
        RAWINPUTHEADER rawinputheader = new RAWINPUTHEADER();
        RAWKEYBOARD keyboard = new RAWKEYBOARD();

        public RAWINPUT() {
            init(new Parameter[] {rawinputheader, keyboard}, (short) 8);
        }
    }

    /**
     * <a href="https://msdn.microsoft.com/en-us/library/windows/desktop/ms645590%28v=vs.85%29.aspx">WM_INPUT</a> event.
     */
    public static final int WM_INPUT = 0x00FF;

    public static void main(String[] args) {

        // a sample frame to redirect raw messages to
        JFrame mainFrame = new JFrame("Main Frame");
        mainFrame.setSize(800, 600);
        mainFrame.setLocationRelativeTo(null);
        mainFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        mainFrame.setVisible(true);

        User32 user32 = User32.getInstance();
        final Function getRawInputData = user32.getFunction("GetRawInputData");
        final Function registerRawInputDevices = User32.getInstance().getFunction("RegisterRawInputDevices");

        Wnd wnd = new Wnd(mainFrame);

        // Substitute frame's window procedure to receive WM_INPUT events
        final WindowProc windowProc = new WindowProc(wnd);
        windowProc.addMessageListener(new WindowMessageListener() {
            public boolean canHandle(WindowMessage message, boolean beforeWindowProc) {
                return message.getMsg() == WM_INPUT;
            }

            public int handle(WindowMessage message) {
                final int RID_INPUT = 0x10000003;

                RAWINPUTHEADER rih = new RAWINPUTHEADER();
                UInt32 size = new UInt32();
                UInt result = new UInt();
                getRawInputData.invoke(result, new Parameter[]{
                        new Pointer.Void(message.getLParam()),
                        new UInt(RID_INPUT),
                        new Pointer.Void(),
                        new Pointer(size),
                        new Int(rih.getLength())
                });
                if (result.getValue() == -1) {
                    return 0;
                }

                PrimitiveArray buffer = new PrimitiveArray(UInt8.class, (int) size.getValue());
                Pointer bufferPtr = new Pointer(buffer);
                getRawInputData.invoke(result, new Parameter[]{
                        new Pointer.Void(message.getLParam()),
                        new UInt(RID_INPUT),
                        bufferPtr,
                        new Pointer(size),
                        new Int(rih.getLength())
                });

                if (result.getValue() != size.getValue()) {
                    return 0;
                }

                RAWINPUT rawinput = new RAWINPUT();
                Pointer rawinputPtr = new Pointer(rawinput);
                bufferPtr.castTo(rawinputPtr);

                RAWKEYBOARD keyboard = rawinput.keyboard;
                System.out.println("Keyboard: make = " + keyboard.MakeCode + "; flags = " + keyboard.Flags + "; Reserved = " + keyboard.Reserved +
                        "; ExtraInformation = " + keyboard.ExtraInformation + "; Message = " + keyboard.Message + "; VKey = " + keyboard.VKey);

                return 0;
            }
        });
        windowProc.substitute();

        mainFrame.addWindowListener(new WindowAdapter() {
            public void windowClosing(WindowEvent e) {
                windowProc.restoreNative();
            }
        });


        RAWINPUTDEVICE rid = new RAWINPUTDEVICE();
        rid.usUsagePage.setValue(1); // generic desktop controls
        rid.usUsage.setValue(6); // keyboard
        rid.dwFlags.setValue(0x00000030 | 0x00000100);
        rid.hwndTarget.setValue(wnd.getValue()); // use Frame's wnd (message queue) to receive events

        Bool result = new Bool();
        long lastErrorCode = registerRawInputDevices.invoke(result, new Pointer(rid), new UInt(1), new UInt(rid.getLength()));
        if (!result.getValue()) {
            throw new LastErrorException(lastErrorCode);
        }
    }
}