use std::collections::HashSet;
use std::io::{self, Write};
use std::sync::{Mutex, OnceLock};

use serde::Serialize;
use windows_sys::Win32::Foundation::{LPARAM, LRESULT, WPARAM};
use windows_sys::Win32::UI::WindowsAndMessaging::{
  CallNextHookEx, DispatchMessageW, GetMessageW, KBDLLHOOKSTRUCT, MSG, SetWindowsHookExW,
  TranslateMessage, UnhookWindowsHookEx, VK_CONTROL, VK_LCONTROL, VK_LMENU, VK_LSHIFT, VK_LWIN,
  VK_MENU, VK_RMENU, VK_RSHIFT, VK_RWIN, VK_SHIFT, WH_KEYBOARD_LL, WM_KEYDOWN, WM_KEYUP,
  WM_SYSKEYDOWN, WM_SYSKEYUP,
};

static MODIFIER_KEYS_DOWN: OnceLock<Mutex<HashSet<&'static str>>> = OnceLock::new();

#[derive(Serialize)]
struct HotkeyEvent<'a> {
  event: &'a str,
  key: &'a str,
  state: &'a str,
}

#[derive(Serialize)]
struct CapabilityEvent<'a> {
  event: &'a str,
  platform: &'a str,
  #[serde(rename = "supportsFn")]
  supports_fn: bool,
  #[serde(rename = "supportsModifierOnly")]
  supports_modifier_only: bool,
  #[serde(rename = "requiresAccessibilityPermission")]
  requires_accessibility_permission: bool,
}

#[derive(Serialize)]
struct StatusEvent<'a> {
  event: &'a str,
  state: &'a str,
  message: Option<&'a str>,
}

#[derive(Serialize)]
struct ErrorEvent<'a> {
  event: &'a str,
  code: &'a str,
  message: &'a str,
}

fn modifier_keys_down() -> &'static Mutex<HashSet<&'static str>> {
  MODIFIER_KEYS_DOWN.get_or_init(|| Mutex::new(HashSet::new()))
}

fn emit_event(key: &'static str, state: &'static str) {
  let payload = HotkeyEvent {
    event: "modifier",
    key,
    state,
  };

  if let Ok(json) = serde_json::to_string(&payload) {
    let _ = writeln!(io::stdout(), "{json}");
    let _ = io::stdout().flush();
  }
}

fn emit_json<T: Serialize>(payload: &T) {
  if let Ok(json) = serde_json::to_string(payload) {
    let _ = writeln!(io::stdout(), "{json}");
    let _ = io::stdout().flush();
  }
}

fn map_vk_code(vk_code: u32) -> Option<&'static str> {
  match vk_code {
    VK_LCONTROL | VK_CONTROL => Some("left_control"),
    VK_LMENU | VK_MENU => Some("left_alt"),
    VK_LSHIFT | VK_SHIFT => Some("left_shift"),
    VK_LWIN => Some("left_command"),
    VK_RMENU => Some("right_alt"),
    VK_RSHIFT => Some("right_shift"),
    VK_RWIN => Some("right_command"),
    _ => None,
  }
}

unsafe extern "system" fn keyboard_proc(code: i32, wparam: WPARAM, lparam: LPARAM) -> LRESULT {
  if code >= 0 {
    let hook = &*(lparam as *const KBDLLHOOKSTRUCT);

    if let Some(key) = map_vk_code(hook.vkCode) {
      let key_set = modifier_keys_down();
      match wparam as u32 {
        WM_KEYDOWN | WM_SYSKEYDOWN => {
          if let Ok(mut down) = key_set.lock() {
            if down.insert(key) {
              emit_event(key, "down");
            }
          }
        }
        WM_KEYUP | WM_SYSKEYUP => {
          if let Ok(mut down) = key_set.lock() {
            if down.remove(key) {
              emit_event(key, "up");
            }
          }
        }
        _ => {}
      }
    }
  }

  CallNextHookEx(std::ptr::null_mut(), code, wparam, lparam)
}

fn main() {
  emit_json(&StatusEvent {
    event: "status",
    state: "starting",
    message: None,
  });
  emit_json(&CapabilityEvent {
    event: "capability",
    platform: "windows",
    supports_fn: false,
    supports_modifier_only: true,
    requires_accessibility_permission: false,
  });

  let hook = unsafe { SetWindowsHookExW(WH_KEYBOARD_LL, Some(keyboard_proc), 0, 0) };
  if hook.is_null() {
    emit_json(&ErrorEvent {
      event: "error",
      code: "failed_to_install_low_level_keyboard_hook",
      message: "Unable to install WH_KEYBOARD_LL hook.",
    });
    emit_json(&StatusEvent {
      event: "status",
      state: "degraded",
      message: Some("Unable to install WH_KEYBOARD_LL hook."),
    });
    return;
  }

  emit_json(&StatusEvent {
    event: "status",
    state: "ready",
    message: None,
  });

  let mut message = MSG::default();
  while unsafe { GetMessageW(&mut message, std::ptr::null_mut(), 0, 0) } > 0 {
    unsafe {
      TranslateMessage(&message);
      DispatchMessageW(&message);
    }
  }

  unsafe {
    UnhookWindowsHookEx(hook);
  }
}
