{-# LANGUAGE FunctionalDependencies #-}

module Rahm.Desktop.Keys.Dsl
  ( doc,
    (-|-),
    ButtonBinding (..),
    ButtonBindings,
    Documented (..),
    HasConfig,
    KeyBinding (..),
    KeyBindings,
    altAltgrMod,
    altHyperAltgrMod,
    altHyperMod,
    altMask,
    altMod,
    altSuperAltgrMod,
    altSuperHyperAltgrMod,
    altSuperHyperMod,
    altSuperMod,
    altgrMask,
    altgrMod,
    bind,
    buttonDocumentation,
    controlAltAltgrMod,
    controlAltHyperAltgrMod,
    controlAltHyperMod,
    controlAltMod,
    controlAltSuperAltgrMod,
    controlAltSuperHyperAltgrMod,
    controlAltSuperHyperMod,
    controlAltSuperMod,
    controlAltgrMod,
    controlHyperAltgrMod,
    controlHyperMod,
    controlMod,
    controlSuperAltgrMod,
    controlSuperHyperAltgrMod,
    controlSuperHyperMod,
    controlSuperMod,
    documentation,
    getConfig,
    hyperAltgrMod,
    hyperMask,
    hyperMod,
    justMod,
    maskMod,
    noMod,
    rawMask,
    runButtons,
    runKeys,
    shiftAltAltgrMod,
    shiftAltHyperAltgrMod,
    shiftAltHyperMod,
    shiftAltMod,
    shiftAltSuperAltgrMod,
    shiftAltSuperHyperAltgrMod,
    shiftAltSuperHyperMod,
    shiftAltSuperMod,
    shiftAltgrMod,
    shiftControlAltAltgrMod,
    shiftControlAltHyperAltgrMod,
    shiftControlAltHyperMod,
    shiftControlAltMod,
    shiftControlAltSuperAltgrMod,
    shiftControlAltSuperHyperAltgrMod,
    shiftControlAltSuperHyperMod,
    shiftControlAltSuperMod,
    shiftControlAltgrMod,
    shiftControlHyperAltgrMod,
    shiftControlHyperMod,
    shiftControlMod,
    shiftControlSuperAltgrMod,
    shiftControlSuperHyperAltgrMod,
    shiftControlSuperHyperMod,
    shiftControlSuperMod,
    shiftHyperAltgrMod,
    shiftHyperMod,
    shiftMod,
    shiftSuperAltgrMod,
    shiftSuperHyperAltgrMod,
    shiftSuperHyperMod,
    shiftSuperMod,
    superAltgrMod,
    superHyperAltgrMod,
    superHyperMod,
    superMask,
    superMod,
  )
where

import Control.Arrow (first, second)
import Control.Monad.State (State, execState, modify')
import Control.Monad.Writer
  ( MonadWriter (tell),
    execWriter,
    forM_,
    when,
  )
import Data.Bits ((.&.))
import Data.List (intercalate, sortOn)
import Data.Map (Map)
import qualified Data.Map as Map
  ( empty,
    fromList,
    fromListWith,
    toList,
  )
import Text.Printf (printf)
import XMonad
  ( Button,
    ButtonMask,
    KeyMask,
    KeySym,
    MonadState (get),
    Window,
    X,
    XConfig (modMask),
    controlMask,
    keysymToString,
    mod1Mask,
    mod3Mask,
    mod4Mask,
    shiftMask,
    (.|.),
  )

data Documented t = Documented String t

data KeyBinding
  = Action (X ())
  | Submap KeyBindings
  | Repeat KeyBindings

type KeyBindings = Map (KeyMask, KeySym) (Documented KeyBinding)

data ButtonBinding
  = ButtonAction (Window -> X ())
  | ButtonSubmap ButtonBindings
  | ButtonContinuous ButtonBindings

-- Window -> X ()

type ButtonBindings = Map (KeyMask, Button) (Documented ButtonBinding)

{- Module that defines a DSL for binding keys. -}
newtype KeysM l a = KeysM (State (XConfig l, KeyBindings) a)
  deriving (Functor, Applicative, Monad)

newtype ButtonsM l a = ButtonsM (State (XConfig l, ButtonBindings) a)
  deriving (Functor, Applicative, Monad)

newtype BindingBuilder b a = BindingBuilder (State (KeyMask, [(KeyMask, b)]) a)
  deriving (Functor, Applicative, Monad)

class HasConfig m where
  getConfig :: m l (XConfig l)

class Bindable k where
  type BindableValue k :: *
  type BindableMonad k :: (* -> *) -> * -> *

  bind :: k -> BindingBuilder (BindableValue k) a -> BindableMonad k l ()

-- section :: String -> BindableMonad k l () -> BindableMonad k l ()

class Binding k b where
  toB :: k -> b

  rawMask :: KeyMask -> k -> BindingBuilder b ()
  rawMask m x = BindingBuilder $ modify' (second ((m, toB x) :))

instance Binding (Window -> X ()) (Documented ButtonBinding) where
  toB = Documented "" . ButtonAction

instance Binding (X ()) (Documented KeyBinding) where
  toB = Documented "" . Action

instance Binding KeyBindings (Documented KeyBinding) where
  toB = Documented "" . Submap

instance Binding a (Documented a) where
  toB = Documented ""

instance Binding a a where
  toB = id

-- Relationships to witness which types can be used with the "doc" function,
-- which is used to document actions in a safe and programmable way..
class Relation k b | k -> b

instance Relation (X ()) KeyBinding

instance Relation KeyBinding KeyBinding

instance Relation ButtonBinding ButtonBinding

instance Relation (Window -> X ()) ButtonBinding

doc :: (Relation k b, Binding k (Documented b)) => String -> k -> Documented b
doc str k = let (Documented _ t) = toB k in Documented str t

runKeys :: KeysM l a -> XConfig l -> KeyBindings
runKeys (KeysM stateM) config =
  snd $ execState stateM (config, Map.empty)

runButtons :: ButtonsM l a -> XConfig l -> ButtonBindings
runButtons (ButtonsM stateM) config =
  snd $ execState stateM (config, Map.empty)

instance HasConfig KeysM where
  getConfig = fst <$> KeysM get

instance HasConfig ButtonsM where
  getConfig = fst <$> ButtonsM get

{- Generally it is assumed that the mod key shoud be pressed, but not always. -}
noMod :: (Binding k b) => k -> BindingBuilder b ()
noMod = rawMask 0

maskMod :: (Binding k b) => KeyMask -> k -> BindingBuilder b ()
maskMod mask action = do
  modMask <- fst <$> BindingBuilder get
  rawMask (modMask .|. mask) action

altMask :: KeyMask
altMask = mod1Mask

hyperMask :: KeyMask
hyperMask = mod3Mask

altgrMask :: KeyMask
altgrMask = 0x80

superMask :: KeyMask
superMask = mod4Mask

justMod :: (Binding k b) => k -> BindingBuilder b ()
justMod = maskMod 0

instance Bindable KeySym where
  type BindableValue KeySym = Documented KeyBinding
  type BindableMonad KeySym = KeysM

  -- bind :: KeySym -> BindingBuilder (X x) a -> KeysM l ()
  bind key (BindingBuilder stM) = do
    m <- modMask <$> getConfig
    let (_, values) = execState stM (m, [])

    KeysM $
      modify' $
        second $
          flip (<>) (Map.fromList (map (\(m, v) -> ((m, key), v)) values))

instance Bindable Button where
  type BindableValue Button = Documented ButtonBinding
  type BindableMonad Button = ButtonsM

  -- bind :: KeySym -> BindingBuilder (Window -> X ()) a -> ButtonsM l ()
  bind button (BindingBuilder stM) = do
    m <- modMask <$> getConfig
    let (_, values) = execState stM (m, [])

    ButtonsM $
      modify' $
        second $
          flip (<>) (Map.fromList (map (\(m, v) -> ((m, button), v)) values))

shiftControlAltSuperHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltSuperHyperAltgrMod =
  maskMod (shiftMask .|. controlMask .|. altMask .|. superMask .|. hyperMask .|. altgrMask)

shiftControlAltSuperHyperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltSuperHyperMod =
  maskMod (shiftMask .|. controlMask .|. altMask .|. superMask .|. hyperMask)

shiftControlAltSuperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltSuperAltgrMod =
  maskMod (shiftMask .|. controlMask .|. altMask .|. superMask .|. altgrMask)

shiftControlAltSuperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltSuperMod =
  maskMod (shiftMask .|. controlMask .|. altMask .|. superMask)

shiftControlAltHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltHyperAltgrMod =
  maskMod (shiftMask .|. controlMask .|. altMask .|. hyperMask .|. altgrMask)

shiftControlAltHyperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltHyperMod =
  maskMod (shiftMask .|. controlMask .|. altMask .|. hyperMask)

shiftControlAltAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltAltgrMod =
  maskMod (shiftMask .|. controlMask .|. altMask .|. altgrMask)

shiftControlAltMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltMod =
  maskMod (shiftMask .|. controlMask .|. altMask)

shiftControlSuperHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlSuperHyperAltgrMod =
  maskMod (shiftMask .|. controlMask .|. superMask .|. hyperMask .|. altgrMask)

shiftControlSuperHyperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlSuperHyperMod =
  maskMod (shiftMask .|. controlMask .|. superMask .|. hyperMask)

shiftControlSuperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlSuperAltgrMod =
  maskMod (shiftMask .|. controlMask .|. superMask .|. altgrMask)

shiftControlSuperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlSuperMod =
  maskMod (shiftMask .|. controlMask .|. superMask)

shiftControlHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlHyperAltgrMod =
  maskMod (shiftMask .|. controlMask .|. hyperMask .|. altgrMask)

shiftControlHyperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlHyperMod =
  maskMod (shiftMask .|. controlMask .|. hyperMask)

shiftControlAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlAltgrMod =
  maskMod (shiftMask .|. controlMask .|. altgrMask)

shiftControlMod :: (Binding k b) => k -> BindingBuilder b ()
shiftControlMod =
  maskMod (shiftMask .|. controlMask)

shiftAltSuperHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltSuperHyperAltgrMod =
  maskMod (shiftMask .|. altMask .|. superMask .|. hyperMask .|. altgrMask)

shiftAltSuperHyperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltSuperHyperMod =
  maskMod (shiftMask .|. altMask .|. superMask .|. hyperMask)

shiftAltSuperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltSuperAltgrMod =
  maskMod (shiftMask .|. altMask .|. superMask .|. altgrMask)

shiftAltSuperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltSuperMod =
  maskMod (shiftMask .|. altMask .|. superMask)

shiftAltHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltHyperAltgrMod =
  maskMod (shiftMask .|. altMask .|. hyperMask .|. altgrMask)

shiftAltHyperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltHyperMod =
  maskMod (shiftMask .|. altMask .|. hyperMask)

shiftAltAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltAltgrMod =
  maskMod (shiftMask .|. altMask .|. altgrMask)

shiftAltMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltMod =
  maskMod (shiftMask .|. altMask)

shiftSuperHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftSuperHyperAltgrMod =
  maskMod (shiftMask .|. superMask .|. hyperMask .|. altgrMask)

shiftSuperHyperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftSuperHyperMod =
  maskMod (shiftMask .|. superMask .|. hyperMask)

shiftSuperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftSuperAltgrMod =
  maskMod (shiftMask .|. superMask .|. altgrMask)

shiftSuperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftSuperMod =
  maskMod (shiftMask .|. superMask)

shiftHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftHyperAltgrMod =
  maskMod (shiftMask .|. hyperMask .|. altgrMask)

shiftHyperMod :: (Binding k b) => k -> BindingBuilder b ()
shiftHyperMod =
  maskMod (shiftMask .|. hyperMask)

shiftAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
shiftAltgrMod =
  maskMod (shiftMask .|. altgrMask)

shiftMod :: (Binding k b) => k -> BindingBuilder b ()
shiftMod = maskMod shiftMask

controlAltSuperHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltSuperHyperAltgrMod =
  maskMod (controlMask .|. altMask .|. superMask .|. hyperMask .|. altgrMask)

controlAltSuperHyperMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltSuperHyperMod =
  maskMod (controlMask .|. altMask .|. superMask .|. hyperMask)

controlAltSuperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltSuperAltgrMod =
  maskMod (controlMask .|. altMask .|. superMask .|. altgrMask)

controlAltSuperMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltSuperMod =
  maskMod (controlMask .|. altMask .|. superMask)

controlAltHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltHyperAltgrMod =
  maskMod (controlMask .|. altMask .|. hyperMask .|. altgrMask)

controlAltHyperMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltHyperMod =
  maskMod (controlMask .|. altMask .|. hyperMask)

controlAltAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltAltgrMod =
  maskMod (controlMask .|. altMask .|. altgrMask)

controlAltMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltMod =
  maskMod (controlMask .|. altMask)

controlSuperHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
controlSuperHyperAltgrMod =
  maskMod (controlMask .|. superMask .|. hyperMask .|. altgrMask)

controlSuperHyperMod :: (Binding k b) => k -> BindingBuilder b ()
controlSuperHyperMod =
  maskMod (controlMask .|. superMask .|. hyperMask)

controlSuperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
controlSuperAltgrMod =
  maskMod (controlMask .|. superMask .|. altgrMask)

controlSuperMod :: (Binding k b) => k -> BindingBuilder b ()
controlSuperMod =
  maskMod (controlMask .|. superMask)

controlHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
controlHyperAltgrMod =
  maskMod (controlMask .|. hyperMask .|. altgrMask)

controlHyperMod :: (Binding k b) => k -> BindingBuilder b ()
controlHyperMod =
  maskMod (controlMask .|. hyperMask)

controlAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
controlAltgrMod =
  maskMod (controlMask .|. altgrMask)

controlMod :: (Binding k b) => k -> BindingBuilder b ()
controlMod = maskMod controlMask

altSuperHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
altSuperHyperAltgrMod =
  maskMod (altMask .|. superMask .|. hyperMask .|. altgrMask)

altSuperHyperMod :: (Binding k b) => k -> BindingBuilder b ()
altSuperHyperMod =
  maskMod (altMask .|. superMask .|. hyperMask)

altSuperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
altSuperAltgrMod =
  maskMod (altMask .|. superMask .|. altgrMask)

altSuperMod :: (Binding k b) => k -> BindingBuilder b ()
altSuperMod =
  maskMod (altMask .|. superMask)

altHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
altHyperAltgrMod =
  maskMod (altMask .|. hyperMask .|. altgrMask)

altHyperMod :: (Binding k b) => k -> BindingBuilder b ()
altHyperMod =
  maskMod (altMask .|. hyperMask)

altAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
altAltgrMod =
  maskMod (altMask .|. altgrMask)

altMod :: (Binding k b) => k -> BindingBuilder b ()
altMod = maskMod altMask

superHyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
superHyperAltgrMod =
  maskMod (superMask .|. hyperMask .|. altgrMask)

superHyperMod :: (Binding k b) => k -> BindingBuilder b ()
superHyperMod =
  maskMod (superMask .|. hyperMask)

superAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
superAltgrMod =
  maskMod (superMask .|. altgrMask)

superMod :: (Binding k b) => k -> BindingBuilder b ()
superMod = maskMod superMask

hyperAltgrMod :: (Binding k b) => k -> BindingBuilder b ()
hyperAltgrMod =
  maskMod (hyperMask .|. altgrMask)

hyperMod :: (Binding k b) => k -> BindingBuilder b ()
hyperMod = maskMod hyperMask

altgrMod :: (Binding k b) => k -> BindingBuilder b ()
altgrMod = maskMod altgrMask

{- Can combine two or more of the functions above to apply the same action to
 - multiple masks. -}
(-|-) ::
  (Binding k b) =>
  (k -> BindingBuilder b ()) ->
  (k -> BindingBuilder b ()) ->
  k ->
  BindingBuilder b ()
(-|-) fn1 fn2 f = fn1 f >> fn2 f

buttonDocumentation :: ButtonBindings -> String
buttonDocumentation = execWriter . document' ""
  where
    document' pref keybindings =
      forM_ (sortOn (map (\(a, b) -> (b, a)) . snd . snd) $ Map.toList (keyBindingsToList keybindings)) $ \(doc, (thing, keys)) -> do
        when (not (null doc) || hasSubmap thing) $
          tell $ printf "%s%s: %s\n" pref (intercalate " or " $ map prettyShow keys) doc
        case thing of
          ButtonAction _ -> return ()
          ButtonSubmap submap -> document' (pref ++ "  ") submap
          ButtonContinuous submap -> do
            tell pref
            tell "  (repeatable):\n"
            document' (pref ++ "    ") submap

    keyBindingsToList :: ButtonBindings -> Map String (ButtonBinding, [(ButtonMask, Button)])
    keyBindingsToList b =
      (\list -> ((\(_, Documented _ t) -> t) (head list), map fst list))
        <$> group (\(_, Documented doc _) -> doc) (sortOn (snd . fst) $ Map.toList b)

    prettyShow :: (ButtonMask, Button) -> String
    prettyShow (mask, button) = printf "%s%s" (showMask mask) (buttonToString button)

    buttonToString = \case
      1 -> "Left Click"
      2 -> "Middle Click"
      3 -> "Right Click"
      4 -> "Wheel Up"
      5 -> "Wheel Down"
      6 -> "Wheel Left"
      7 -> "Wheel Right"
      8 -> "Browser Back"
      9 -> "Browser Forward"
      13 -> "Thumb Target"
      14 -> "Index Forward"
      15 -> "Index Back"
      b -> "Button " ++ show b

    hasSubmap b = case b of
      ButtonAction _ -> False
      _ -> True

documentation :: KeyBindings -> String
documentation = execWriter . document' ""
  where
    document' pref keybindings =
      forM_ (sortOn (map (\(a, b) -> (b, a)) . snd . snd) $ Map.toList (keyBindingsToList keybindings)) $ \(doc, (thing, keys)) -> do
        when (not (null doc) || hasSubmap thing) $
          tell $ printf "%s%s: %s\n" pref (intercalate " or " $ map prettyShow keys) doc
        case thing of
          Action _ -> return ()
          Submap submap -> document' (pref ++ "  ") submap
          Repeat submap -> do
            tell pref
            tell "  (repeatable):\n"
            document' (pref ++ "    ") submap

    keyBindingsToList :: KeyBindings -> Map String (KeyBinding, [(KeyMask, KeySym)])
    keyBindingsToList b =
      (\list -> ((\(_, Documented _ t) -> t) (head list), map fst list))
        <$> group (\(_, Documented doc _) -> doc) (sortOn (snd . fst) $ Map.toList b)

    prettyShow :: (KeyMask, KeySym) -> String
    prettyShow (mask, key) = printf "%s%s" (showMask mask) (keysymToString key)

    hasSubmap b = case b of
      Action _ -> False
      _ -> True

showMask :: KeyMask -> String
showMask mask =
  let masks =
        [ (shiftMask, "S"),
          (altMask, "A"),
          (mod3Mask, "H"),
          (mod4Mask, "M"),
          (altgrMask, "AGr"),
          (controlMask, "C")
        ]
   in concatMap ((++ "-") . snd) $ filter ((/= 0) . (.&. mask) . fst) masks

group :: (Ord b) => (a -> b) -> [a] -> Map b [a]
group fn = Map.fromListWith (++) . map (first fn . (\a -> (a, [a])))