{-# LANGUAGE UndecidableInstances #-}

-- | This module creates a DSL for binding keys in a succinct and expressive
-- way. This DSL follows the pattern:
--
-- bind <key> $ do
--    <mask> $ <doc string $?> binding
--    <mask> $ <doc string $?> binding
--    <mask> $ <doc string $?> binding
--
-- for example:
--
-- bind xK_x $ do
--   justMod $ doc "Kill the current window" (withFocused X.kill)
--   shiftMod $ doc "Restart xmonad" restart
--
-- bind xK_v $ do
--   justMod $
--     continuous $ do
--       bind xK_plus $ doc "increase volume" increaseVolume
--       bind xK_minus $ doc "decrease volume" decreaseVolume
module Rahm.Desktop.Keys.Dsl2 where

import Control.Applicative ((<|>))
import Control.Monad.Fix (fix)
import Control.Monad.RWS (All (All), MonadTrans (lift), MonadWriter, forM, forM_, when)
import Control.Monad.Reader (Reader, ask, runReader)
import Control.Monad.State (MonadTrans, StateT (StateT))
import Control.Monad.Trans.Maybe (MaybeT (..))
import Control.Monad.Trans.Writer (Writer, WriterT, execWriter, execWriterT, runWriter)
import Control.Monad.Writer.Class (tell)
import Data.Bits ((.&.))
import Data.Functor.Identity (Identity)
import Data.List (intercalate)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)
import Rahm.Desktop.Common (pointerWindow, runMaybeT_)
import Rahm.Desktop.Keys.Grab
import Rahm.Desktop.Keys.KeyCodeMapping (setupKeycodeMapping)
import Rahm.Desktop.Logger (LogLevel (Debug, Info), logs)
import Rahm.Desktop.Submap (ButtonOrKeyEvent (ButtonPress, KeyPress, event_keycode, event_mask), getStringForKey, nextButtonOrKeyEvent)
import Rahm.Desktop.XMobarLog (spawnXMobar)
import Rahm.Desktop.XMobarLog.PendingBuffer (pushAddPendingBuffer, pushPendingBuffer, clearPendingBuffer)
import XMonad

-- | A documented "thing." It is essentially an item with a string attached to
-- it. A glorified tuple (String, t)
data Documented t = Documented
  { docString :: String,
    undocument :: t
  }

-- | The Documented type is a functor.
instance Functor Documented where
  fmap fn (Documented s t) = Documented s (fn t)

-- | Type family for an action associated with a type. This type family
-- indicates what type of action a keytype can be bound to.
type family Action t where
  -- KeySyms are bound to contextless actions with type X ()
  Action KeySymOrKeyCode = X ()
  -- Buttons are associated with actions with type Window -> X (). In other
  -- words, actions bound to a button have windows associated with it.
  Action Button = Window -> X ()

class (Bind (Super k)) => LiftBinding k where
  type Super k :: *
  doLift :: k -> Super k

instance LiftBinding KeySymOrKeyCode where
  type Super KeySymOrKeyCode = KeySymOrKeyCode
  doLift = id

instance LiftBinding KeySym where
  type Super KeySym = Super KeySymOrKeyCode
  doLift = doLift . Ks

instance LiftBinding KeyCode where
  type Super KeyCode = Super KeySymOrKeyCode
  doLift = doLift . Kc

instance LiftBinding Button where
  type Super Button = Button
  doLift = id

-- | An GADT for XConfig that hides the 'l' parameter. This keeps type
-- signatures clean by not having to carry around a superfluous type variable.
data XConfigH where
  XConfigH :: forall l. XConfig l -> XConfigH

-- | A type for binding. These are the type of bindings keys and buttons can be
-- bound to The type 't' refers to the key-type of the binding: either `KeySym`
-- or `Button`.
data Binding t
  = -- | Just an action. When this key is pressed, just do some action.
    -- This uses the type-family defined above to determine what type the action
    -- should be.
    Action (Action t)
  | -- | Sub-bindings. These bindings require another key to resolve.
    Submap (forall l. XConfig l -> BindingsMap)
  | -- | Like submap, but these bindings are repeatable. This operates as a
    -- self-loop in the keybinding graph.
    Repeat (Binding t) (forall l. XConfig l -> BindingsMap)
  | -- | No action bound to the key.
    NoBinding

-- | Compiled bindings map. This is the type built up by the Binder monad.
data BindingsMap = BindingsMap
  { -- | Bindings for keys
    key_bindings :: Map (KeyMask, KeySymOrKeyCode) (Documented (Binding KeySymOrKeyCode)),
    -- | Bindings for buttons.
    button_bindings :: Map (KeyMask, Button) (Documented (Binding Button)),
    -- | If no mapping for a key sym exists, this function is called to handle
    -- the no match key.
    no_match_catch_key :: (KeyMask, KeySym, String) -> X (),
    -- | If no mapping for a button exists, this function is called to handle
    -- the unbound button.
    no_match_catch_button :: (KeyMask, Button) -> Window -> X ()
  }

-- | BindingMaps are combined by combining the bindings map and combining the
-- handling functio
instance Semigroup BindingsMap where
  (BindingsMap mk1 mb1 fk1 fb1) <> (BindingsMap mk2 mb2 fk2 fb2) =
    BindingsMap
      (mk1 <> mk2)
      (mb1 <> mb2)
      (\a -> fk1 a >> fk2 a)
      (\a b -> fb1 a b >> fb2 a b)

instance Monoid BindingsMap where
  mempty = BindingsMap mempty mempty (\_ -> return ()) (\_ _ -> return ())

-- | This type is an intermediate monad in the DSL. It binds an action to a key
-- mask. This is all bound to a key or button in the Binder monad.
newtype MaskBinder k a = MaskBinder
  { unMaskBinder :: WriterT (Map KeyMask (Documented (Binding k))) (Reader XConfigH) a
  }
  deriving
    ( Functor,
      Applicative,
      Monad,
      MonadWriter (Map KeyMask (Documented (Binding k))),
      MonadReader XConfigH
    )

-- | The ultimate monad for this DSL. This Binder monad builds up a BindingsMap
newtype Binder a = Binder (WriterT BindingsMap (Reader XConfigH) a)
  deriving (Functor, Applicative, Monad, MonadWriter BindingsMap, MonadReader XConfigH)

-- | This typeclas esentiallly defines what types can be bound in the DSL. I.e.
-- what types are allowed to be arguments to the xMod, rawMask, etc. functions.
class BindingType a where
  -- | What key this action belongs to.
  type BoundTo a :: *

  -- | Convert the action into a Documented binding
  toBinding :: a -> Documented (Binding (BoundTo a))

-- | Bindings are trivially a BindingType.
instance BindingType (Binding t) where
  type BoundTo (Binding t) = t
  toBinding = Documented ""

-- | An X () can be bound to a Binding KeySym
instance BindingType (X ()) where
  type BoundTo (X ()) = KeySymOrKeyCode
  toBinding = Documented "" . Action

-- | A Window -> X () can be bound to a Binding Button.
instance BindingType (Window -> X ()) where
  type BoundTo (Window -> X ()) = Button
  toBinding = Documented "" . Action

-- | Any Documented BindingType is also a BindingType.
instance (BindingType a) => BindingType (Documented a) where
  type BoundTo (Documented a) = BoundTo a
  toBinding (Documented s (toBinding -> (Documented _ a))) = Documented s a

-- | This typeclass is responsible for converting a compiled MaskBinder into a
-- BindingsMap using a key type 'k'.
class Bind k where
  doBinding :: k -> Map KeyMask (Documented (Binding k)) -> BindingsMap
  rawMaskRaw :: KeyMask -> Documented (Binding k) -> MaskBinder k ()

-- | For this, it adds the bindings to the buttonMap
instance Bind Button where
  doBinding but mp = mempty {button_bindings = Map.mapKeys (,but) mp}
  rawMaskRaw mask act = tell (Map.singleton mask act)

-- | For this, it adds the bindings to the keysMap
instance Bind KeySymOrKeyCode where
  doBinding key mp = mempty {key_bindings = Map.mapKeys (,key) mp}
  rawMaskRaw mask act = tell (Map.singleton mask act)

rawMask :: (Bind (BoundTo a), BindingType a) => KeyMask -> a -> MaskBinder (BoundTo a) ()
rawMask mask act = rawMaskRaw mask (toBinding act)

withMod :: (Bind (BoundTo a), BindingType a) => KeyMask -> a -> MaskBinder (BoundTo a) ()
withMod m act = do
  (XConfigH (modMask -> mm)) <- ask
  rawMask (mm .|. m) act

-- | Mask bindings.
noMod, justMod, shiftMod, controlMod, altMod :: (Bind (BoundTo a), BindingType a) => a -> MaskBinder (BoundTo a) ()
justMod = withMod 0
noMod = rawMask 0
shiftMod = withMod shiftMask
controlMod = withMod controlMask
altMod = withMod mod1Mask

-- | allows easy assigning multiple masks to the same action.
(-|-) ::
  (Bind (BoundTo a), BindingType a) =>
  (a -> MaskBinder (BoundTo a) ()) ->
  (a -> MaskBinder (BoundTo a) ()) ->
  a ->
  MaskBinder (BoundTo a) ()
m1 -|- m2 = \act -> m1 act >> m2 act

-- | Bind a key to a maksBinder.
bind :: (LiftBinding k) => k -> MaskBinder (Super k) () -> Binder ()
bind k h =
  tell . doBinding (doLift k) . runReader (execWriterT $ unMaskBinder h) =<< ask

-- | Bind multiple keys to the same mask binder.
bindL :: (LiftBinding k) => [k] -> MaskBinder (Super k) () -> Binder ()
bindL ks h = mapM_ (`bind` h) ks

-- | Convenience function for Documented.
doc :: String -> a -> Documented a
doc = Documented

-- | A concrete-typed version of 'const'
noWindow :: X () -> Window -> X ()
noWindow fn _ = fn

data Bindings where
  Bindings ::
    (forall l. XConfig l -> Map (KeyMask, KeySym) (X ())) ->
    (forall l. XConfig l -> Map (KeyMask, KeyCode) (X ())) ->
    (forall l. XConfig l -> Map (ButtonMask, Button) (Window -> X ())) ->
    Bindings

-- | Turn a BindingsMap into two values usable values for the XMonad config.
resolveBindings ::
  BindingsMap -> Bindings
resolveBindings (BindingsMap keyAndKeyCodeBindings buttonBindings _ _) =
  Bindings
    (\c -> Map.mapWithKey (\k v -> pushK k (bindingToX c) (undocument v)) keyBindings)
    (\c -> Map.mapWithKey (\k v -> bindingToX c (undocument v)) keycodeBindings)
    (\c -> Map.mapWithKey (\k v -> pushB k (bindingToWinX c) (undocument v)) buttonBindings)
  where
    (keyBindings, keycodeBindings) =
      partitionMap
        ( \case
            (m, Kc keyCode) -> Right (m, keyCode)
            (m, Ks keySym) -> Left (m, keySym)
        )
        keyAndKeyCodeBindings

    pushB (_, b) fn binding win = do
      if isRepeatOrSubmap binding
        then pushPendingBuffer ("b" ++ show b ++ " ") $ fn binding win
        else fn binding win
      clearPendingBuffer

    pushK (m, k) fn binding = do
      if isRepeatOrSubmap binding
        then do
          let s = getStringForKey (m, k)
          pushPendingBuffer (s ++ " ") $ fn binding
        else fn binding
      clearPendingBuffer

    bindingToX :: forall l. XConfig l -> Binding KeySymOrKeyCode -> X ()
    bindingToX conf = \case
      NoBinding -> return ()
      Action a -> a
      Submap sm -> doSubmap conf (sm conf) (return ())
      Repeat a sm -> bindingToX conf a >> fix (doSubmap conf (sm conf))

    bindingToWinX :: forall l. XConfig l -> Binding Button -> Window -> X ()
    bindingToWinX conf binding win = case binding of
      NoBinding -> return ()
      Action fn -> fn win
      Submap sm -> doSubmap conf (sm conf) (return ())
      Repeat a sm -> bindingToWinX conf a win >> fix (doSubmap conf (sm conf))

    doSubmap :: forall l. XConfig l -> BindingsMap -> X () -> X ()
    doSubmap conf (BindingsMap kbind bbind catk catb) after = do
      nextPressEvent $ \str ->
        \case
          (ButtonPress m b) -> do
            win <- pointerWindow
            case Map.lookup (m, b) bbind of
              (Just binding) ->
                pushAddPendingBuffer (str ++ " ") $ do
                  bindingToWinX conf (undocument binding) win
                  after
              Nothing -> catb (m, b) win
          (KeyPress m k c s) -> do
            case Map.lookup (m, Kc c) kbind <|> Map.lookup (m, Ks k) kbind of
              (Just binding) ->
                pushAddPendingBuffer (str ++ " ") $ do
                  bindingToX conf (undocument binding)
                  after
              Nothing -> catk (m, k, s)

    isRepeatOrSubmap = \case
      Repeat {} -> True
      Submap {} -> True
      _ -> False

    nextPressEvent fn = runMaybeT_ $ do
      ev <- nextButtonOrKeyEvent
      let str = case ev of
            ButtonPress m b -> "b" ++ show b
            KeyPress _ _ _ s -> s
      lift $
        fn str ev

-- Create a submap in place of an action.
subbind :: Binder () -> Binding t
subbind (Binder b) =
  Submap $ \config ->
    runReader (execWriterT b) (XConfigH config)

repeatable :: Binder () -> Binding t
repeatable (Binder b) =
  Repeat NoBinding $ \config ->
    runReader (execWriterT b) (XConfigH config)

-- Similar to repeatable, but all the keys in the binder start the loop.
continuous :: Binder () -> Binder ()
continuous (Binder b) = do
  conf <- ask
  let bm@(BindingsMap keyBinds mouseBinds _ _) =
        runReader (execWriterT b) conf

  forM_ (Map.toList keyBinds) $ \((m, k), Documented _ b) ->
    bind k $ rawMask m $ Repeat b $ const bm

  forM_ (Map.toList mouseBinds) $ \((m, k), Documented _ b) ->
    bind k $ rawMask m $ Repeat b $ const bm

runBinder :: XConfig l -> Binder a -> BindingsMap
runBinder conf (Binder binder) = runReader (execWriterT binder) (XConfigH conf)

withBindings :: Binder a -> XConfig l -> XConfig l
withBindings b config =
  let (Bindings keyBinds keycodeBinds buttonBinds) =
        resolveBindings $ runBinder config b
   in setupKeycodeMapping keycodeBinds $
        config
          { keys = keyBinds,
            mouseBindings = buttonBinds
          }

documentation :: XConfig l -> Binder () -> String
documentation conf binder =
  documentation' $ runBinder conf binder
  where
    documentation' :: BindingsMap -> String
    documentation' (BindingsMap kmap bmap _ _) = execWriter $ do
      forM_
        ( Map.toList $
            invert $
              Map.union
                (Map.map documentBinding (Map.mapKeys keyToStr kmap))
                (Map.map documentBinding (Map.mapKeys buttonToStr bmap))
        )
        $ \(doc, keys) ->
          when (doc /= "") $ do
            when (length keys > 1) (tell "\n")
            tell (intercalate ",\n" keys)
            tell " -> "
            tell (tindent doc)
            when (length keys > 1) (tell "\n")

    documentBinding :: Documented (Binding r) -> String
    documentBinding = \case
      (Documented s (Action _)) -> s
      (Documented s (Submap mp)) ->
        s ++ "\n" ++ indent (documentation' (mp conf))
      (Documented s (Repeat _ mp)) ->
        s ++ " (repeatable)\n" ++ indent (documentation' (mp conf))
      _ -> ""

    indent = unlines . map ("  " ++) . lines
    tindent (lines -> (h : t)) = unlines (h : map ("  " ++) t)
    tindent x = x

    keyToStr (m, k) = showMask m ++ keysymOrKeyCodeToString k
    keysymOrKeyCodeToString (Kc code) = show code
    keysymOrKeyCodeToString (Ks sym) = keysymToString sym

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

    invert :: (Ord a, Ord b) => Map a b -> Map b [a]
    invert = Map.fromListWith (++) . map (\(a, b) -> (b, [a])) . Map.toList

    buttonToStr (m, b) = showMask m ++ buttonNumToStr b
    buttonNumToStr = \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

bindOtherKeys :: ((KeyMask, KeySym, String) -> X ()) -> Binder ()
bindOtherKeys fn = Binder $ tell (mempty {no_match_catch_key = fn})

bindOtherButtons :: ((KeyMask, Button) -> Window -> X ()) -> Binder ()
bindOtherButtons fn = Binder $ tell (mempty {no_match_catch_button = fn})

partitionMap :: (Ord k, Ord k1, Ord k2) => (k -> Either k1 k2) -> Map k a -> (Map k1 a, Map k2 a)
partitionMap f mp =
  foldl
    ( \(mp1, mp2) (k, v) ->
        case f k of
          Left k1 -> (Map.insert k1 v mp1, mp2)
          Right k2 -> (mp1, Map.insert k2 v mp2)
    )
    (mempty, mempty)
    (Map.toList mp)