-- 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 Internal.Intercept where

import XMonad

import Text.Printf
import Data.Monoid (Endo(..))
import Control.Monad.Trans.Class (lift)
import Control.Monad (forM, forM_)
import Data.Default (Default, def)
import Data.Map (Map)
import qualified Data.Map as Map
import qualified XMonad.Util.ExtensibleState as XS
import Data.Monoid (All(..))

import Internal.Logger

type WindowHook = Query ()

newtype NoPersist a = NoPersist a
  deriving (Typeable)

instance Show (NoPersist a) where
  show (NoPersist _) = show ()

instance (Default a) => Read (NoPersist a) where
  readsPrec i s = map (\(_, s) -> (NoPersist def, s)) (readsPrec i s :: [((), String)])

instance (Default a) => Default (NoPersist a) where
  def = NoPersist def

instance (Default a, Typeable a) => ExtensionClass (NoPersist a) where
  initialValue = NoPersist def


data InterceptState = InterceptState (NoPersist (Map (KeyMask, KeySym) (X ())))

data RemapState = RemapState (NoPersist (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 (NoPersist 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)

interceptHook :: Event -> X All
interceptHook event = do
  InterceptState (NoPersist map) <- XS.get
  case event of
    KeyEvent { 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 (m, keysym) map of

        Just xdo -> do
          xdo
          mapM_ (\m' -> io $ ungrabKey dpy code (m .|. m') rootw)
              =<< extraModifiers
          return (All False)

        Nothing -> return (All True)

    _ -> return (All True)

setIntercept :: (KeyMask, KeySym) -> X () -> X ()
setIntercept (keyMask, keysym) action = do
  XS.modify $ \(InterceptState (NoPersist m)) -> InterceptState $ NoPersist $
    Map.insert (keyMask, keysym) action m
  XConf { display = dpy, theRoot = rootw } <- ask

  doGrab dpy rootw (keyMask, keysym)


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 (NoPersist keyMap)) -> RemapState $ NoPersist $
      Map.insert (window, keyFrom) action keyMap

rebindKey :: (KeyMask, KeySym) -> (KeyMask, KeySym) -> WindowHook
rebindKey keyFrom keyTo = do
  window <- ask
  remapKey keyFrom $ do
    XConf { display = disp, theRoot = rootw } <- ask
    codes <- io $ getKeyCodesForKeysym disp (snd keyTo)
    case codes of
      (keyCode:_) -> do
        io $ allocaXEvent $ \xEv -> do
          setEventType xEv keyPress
          setKeyEvent xEv window rootw none (fst keyTo) keyCode True
          sendEvent disp window True keyPressMask xEv
          setEventType xEv keyRelease
          sendEvent disp window True keyPressMask xEv

      _ -> return ()