-- Module for intercepting key presses not explicity mapped in the key bindings.
-- This uses some deep magic with grabKey and windows and everything else, but
-- it makes window-specific key bindings awesome!
module Rahm.Desktop.RebindKeys
  ( remapHook,
    getKeyCodesForKeysym,
    doGrab,
    disableKey,
    remapKey,
    sendKeyQ,
    sendKey,
    rebindKey,
    WindowHook (..),
  )
where

import Control.Monad (forM, forM_)
import Control.Monad.Trans.Class (lift)
import Data.Default (def)
import Data.Map (Map)
import qualified Data.Map as Map
  ( delete,
    findWithDefault,
    fromListWith,
    insert,
    lookup,
  )
import Data.Monoid (All (..))
import XMonad
  ( Default (def),
    Display,
    Event (KeyEvent, ev_event_type, ev_keycode, ev_state, ev_window),
    ExtensionClass (initialValue),
    KeyCode,
    KeyMask,
    KeySym,
    MonadReader (ask),
    Query (..),
    Window,
    X,
    XConf (XConf, display, theRoot),
    allocaXEvent,
    displayKeycodes,
    extraModifiers,
    grabKey,
    grabModeAsync,
    io,
    keyPress,
    keyPressMask,
    keyRelease,
    keyReleaseMask,
    keycodeToKeysym,
    liftX,
    noSymbol,
    none,
    sendEvent,
    setEventType,
    setKeyEvent,
    (.|.),
  )
import qualified XMonad.Util.ExtensibleState as XS (get, modify)

type WindowHook = Query ()

newtype InterceptState
  = InterceptState (Map (KeyMask, KeySym) (X ()))

newtype RemapState
  = RemapState (Map (Window, (KeyMask, KeySym)) (X ()))

instance ExtensionClass InterceptState where
  initialValue = InterceptState def

instance ExtensionClass RemapState where
  initialValue = RemapState def

remapHook :: Event -> X All
remapHook event = do
  RemapState map <- XS.get

  case event of
    KeyEvent {ev_window = win, ev_event_type = typ, ev_keycode = code, ev_state = m}
      | typ == keyPress -> do
        XConf {display = dpy, theRoot = rootw} <- ask
        keysym <- io $ keycodeToKeysym dpy code 0

        case Map.lookup (win, (m, keysym)) map of
          Just xdo -> do
            xdo
            return (All False)
          Nothing -> return (All True)
    _ -> return (All True)

getKeyCodesForKeysym :: Display -> KeySym -> IO [KeyCode]
getKeyCodesForKeysym dpy keysym = do
  let (minCode, maxCode) = displayKeycodes dpy
      allCodes = [fromIntegral minCode .. fromIntegral maxCode]

  syms <- forM allCodes $ \code -> keycodeToKeysym dpy code 0
  let keysymMap' = Map.fromListWith (++) (zip syms [[code] | code <- allCodes])

  -- keycodeToKeysym returns noSymbol for all unbound keycodes, and we don't
  -- want to grab those whenever someone accidentally uses def :: KeySym
  let keysymMap = Map.delete noSymbol keysymMap'
  let keysymToKeycodes sym = Map.findWithDefault [] keysym keysymMap

  return $ keysymToKeycodes keysym

doGrab :: Display -> Window -> (KeyMask, KeySym) -> X ()
doGrab dpy win (keyMask, keysym) = do
  let grab kc m = io $ grabKey dpy kc m win True grabModeAsync grabModeAsync

  codes <- io $ getKeyCodesForKeysym dpy keysym

  forM_ codes $ \kc ->
    mapM_ (grab kc . (keyMask .|.)) =<< extraModifiers

disableKey :: (KeyMask, KeySym) -> WindowHook
disableKey key = remapKey key (return ())

remapKey :: (KeyMask, KeySym) -> X () -> WindowHook
remapKey keyFrom action = do
  window <- ask
  Query $
    lift $ do
      XConf {display = disp, theRoot = rootw} <- ask
      doGrab disp window keyFrom

      XS.modify $ \(RemapState keyMap) ->
        RemapState $
          Map.insert (window, keyFrom) action keyMap

-- sendKey, but as a query.
sendKeyQ :: (KeyMask, KeySym) -> Query ()
sendKeyQ key = do
  win <- ask
  liftX (sendKey key win)

sendKey :: (KeyMask, KeySym) -> Window -> X ()
sendKey (keymask, keysym) w = do
  XConf {display = disp, theRoot = rootw} <- ask

  codes <- io $ getKeyCodesForKeysym disp keysym

  case codes of
    (keycode : _) ->
      io $
        allocaXEvent $ \xEv -> do
          setEventType xEv keyPress
          setKeyEvent xEv w rootw none keymask keycode True
          sendEvent disp w True keyPressMask xEv

          setEventType xEv keyRelease
          sendEvent disp w True keyReleaseMask xEv
    _ -> return ()

rebindKey :: (KeyMask, KeySym) -> (KeyMask, KeySym) -> WindowHook
rebindKey keyFrom keyTo =
  (remapKey keyFrom . sendKey keyTo) =<< ask