module Rahm.Desktop.KeysM where

import Data.List
import Data.Bits ((.&.))
import Control.Monad.Writer
import Text.Printf
import Control.Arrow (second, first)
import Control.Monad (void)
import Control.Monad.State (State(..), modify', get, execState)
import XMonad
import Data.Map (Map)
import qualified Data.Map as Map

data Documented t = Documented String t

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

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

type ButtonBinding = Window -> X ()
type ButtonBindings = Map (KeyMask, Button) 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 (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

doc :: (Binding k (Documented KeyBinding)) => String -> k -> Documented KeyBinding
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 = 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

{- Meant for submapping, binds all alphanumeric charactes to (fn c). -}
mapNumbersAndAlpha :: KeyMask -> (Char -> X ()) -> KeysM l ()
mapNumbersAndAlpha km fn = do
  mapNumbers km fn
  mapAlpha km fn

{- Meant for submapping. This binds all numbers to (fn x) where x is the number
 - pressed and fn is the function provided. -}
mapNumbers :: KeyMask -> (Char -> X ()) -> KeysM l ()
mapNumbers km fn = do
  mapM_ (\(key, ch) -> bind key $ rawMask km (fn ch))
    [ (xK_0, '0')
    , (xK_1, '1')
    , (xK_2, '2')
    , (xK_3, '3')
    , (xK_4, '4')
    , (xK_5, '5')
    , (xK_6, '6')
    , (xK_7, '7')
    , (xK_8, '8')
    , (xK_9, '9')
    -- Programmer Dvorak shifts the numbers so I have to map to their unshifted
    -- form.
    , (xK_bracketright, '6')
    , (xK_exclam, '8')
    , (xK_bracketleft, '7')
    , (xK_braceleft, '5')
    , (xK_braceright, '3')
    , (xK_parenleft, '1')
    , (xK_equal, '9')
    , (xK_asterisk, '0')
    , (xK_parenright, '2')
    , (xK_plus, '4') ]

{- Meant for submapping. This binds all alpha charactes to (fn c) where c is the
 - character pressed and fn is the function provided. -}
mapAlpha :: KeyMask -> (Char -> X ()) -> KeysM l ()
mapAlpha km fn =
  mapM_ (\(key, ch) -> bind key $ rawMask km (fn ch)) [
    (xK_a, 'a')
  , (xK_b, 'b')
  , (xK_c, 'c')
  , (xK_d, 'd')
  , (xK_e, 'e')
  , (xK_f, 'f')
  , (xK_g, 'g')
  , (xK_h, 'h')
  , (xK_i, 'i')
  , (xK_j, 'j')
  , (xK_k, 'k')
  , (xK_l, 'l')
  , (xK_m, 'm')
  , (xK_n, 'n')
  , (xK_o, 'o')
  , (xK_p, 'p')
  , (xK_q, 'q')
  , (xK_r, 'r')
  , (xK_s, 's')
  , (xK_t, 't')
  , (xK_u, 'u')
  , (xK_v, 'v')
  , (xK_w, 'w')
  , (xK_x, 'x')
  , (xK_y, 'y')
  , (xK_z, 'z')
  ]


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 =
      fmap (\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])))