module Wetterhorn.Core.Keys where

import Control.Monad (forever, void, when)
import Control.Monad.Cont.Class
import Control.Monad.IO.Class
import Control.Monad.State (MonadState (get, put), MonadTrans (lift), StateT, evalStateT, gets, modify)
import Control.Monad.Trans.Cont
import Data.Bits
import Data.Word
import Wetterhorn.Core.ButtonEvent (ButtonEvent)
import Wetterhorn.Core.KeyEvent
import qualified Wetterhorn.Core.KeyEvent as KeyEvent
import qualified Wetterhorn.Core.ButtonEvent as ButtonEvent
import Wetterhorn.Core.W
import Wetterhorn.Foreign.WlRoots (wlrSeatKeyboardNotifyKey, wlrSeatSetKeyboard)

-- | Forwards the given key event to the focused window.
forwardKey :: KeyEvent -> W ()
forwardKey keyEvent = do
  seatPtr <- getSeat
  wio $ do
    wlrSeatSetKeyboard
      seatPtr
      (device keyEvent)

    wlrSeatKeyboardNotifyKey
      seatPtr
      (timeMs keyEvent)
      (keycode keyEvent)
      ( case state keyEvent of
          KeyReleased -> 0
          _ -> 1
      )

-- | Forwards the current key event to the focused window.
forwardEvent :: KeyEvent -> KeysM ()
forwardEvent = liftW . forwardKey

-- | Enumeration of possible modifiers
data Modifier = Shift | Lock | Control | Mod1 | Mod2 | Mod3 | Mod4 | Mod5
  deriving (Eq, Ord, Show, Read, Enum, Bounded)

-- | Converts a modifier to its associated mask.
modifierToMask :: Modifier -> Word32
modifierToMask m =
  1
    `shiftL` case m of
      Shift -> 0
      Lock -> 1
      Control -> 2
      Mod1 -> 3
      Mod2 -> 4
      Mod3 -> 5
      Mod4 -> 6
      Mod5 -> 7

data KeysState = KeysState
  { -- | Reference to the top. Used for a continue statement.
    keysTop :: KeysM (),
    handleContinuation :: KeyContinuation -> W ()
  }

-- | The Keys monad. This monad abstracts away control flow for handling key
-- bindings. This makes it easy to make key-sequence bindings.
-- newtype KeysM a = KeysM ((KeyEvent -> W ()) -> KeyEvent -> W (KeysMR a))
newtype KeysM a = KeysM (ContT () (StateT KeysState W) a)
  deriving (Monad, Functor, Applicative, MonadCont, MonadIO)

-- | KeysM can be lifted from a W action.
instance Wlike KeysM where
  liftW = KeysM . lift . lift

type KeyContinuation = KeyEvent -> W ()

useKeysWithContinuation :: (KeyContinuation -> W ()) -> KeysM () -> W ()
useKeysWithContinuation continuation (forever -> km@(KeysM c)) =
  evalStateT (evalContT c) (KeysState km continuation)

useKeys :: KeysM () -> W ()
useKeys = useKeysWithContinuation putKeyHandler

-- | Returns the next key event.
nextKeyEvent :: KeysM KeyEvent
nextKeyEvent = do
  st <- KeysM $ lift get
  KeysM $
    shiftT
      ( \keyHandler ->
          lift . lift $
            handleContinuation st (\kp -> evalStateT (keyHandler kp) st)
      )

-- | Discards the rest of the continuation and starts again from the top. Useful
-- for keybinds where once the key is handled, there's nothing left to do.
continue :: KeysM ()
continue = do
  st <- KeysM $ lift get
  let (KeysM topCont) = keysTop st

  -- This shift discards the rest of the computation and instead returns to the
  -- top of the handler.
  KeysM $ shiftT (\_ -> resetT topCont)

-- | Returns the "top" continuation.
getTop :: KeysM (KeysM ())
getTop = KeysM (gets keysTop)

putKeyHandler :: KeyContinuation -> W ()
putKeyHandler handler = do
  s@State {currentHooks = hooks} <- get
  put
    s
      { currentHooks =
          hooks
            { keyHook = void <$> handler
            }
      }

nextButtonEvent :: KeysM ButtonEvent
nextButtonEvent = do
  st <- KeysM get
  KeysM $
    shiftT $ \h ->
      lift $ lift $ putButtonHandler (\ev -> evalStateT (h ev) st)
  where
    putButtonHandler h = do
      modify $ \st -> st {currentHooks = (currentHooks st) {buttonHook = h}}

nextButtonOrKeyEvent :: KeysM (Either ButtonEvent KeyEvent)
nextButtonOrKeyEvent = do
  st <- KeysM get
  KeysM $
    shiftT $ \rest -> 
      lift $ lift $ do
        putButtonHandler (\ev -> evalStateT (rest (Left ev)) st)
        handleContinuation st (\ev -> evalStateT (rest (Right ev)) st)

  where
    putButtonHandler h = do
      modify $ \st -> st {currentHooks = (currentHooks st) {buttonHook = h}}

nextButtonOrKeyPress :: KeysM (Either ButtonEvent KeyEvent)
nextButtonOrKeyPress = do
  ev <- nextButtonOrKeyEvent
  case ev of
    Left bev | ButtonEvent.state bev == ButtonEvent.ButtonPressed -> return ev
    Left bev -> forwardButtonEvent bev >> nextButtonOrKeyPress
    Right kev | KeyEvent.state kev == KeyEvent.KeyPressed -> return ev
    Right kev -> forwardEvent kev >> nextButtonOrKeyPress

 where
   forwardButtonEvent _ = return ()
  

-- | Returns the next KeyPressed event. This is likely what 90% of use cases
-- want rather than nextKeyEvent.
nextKeyPress :: KeysM KeyEvent
nextKeyPress = do
  k <- nextKeyEvent
  if KeyEvent.state k /= KeyPressed
    then forwardEvent k >> nextKeyPress
    else return k

--
-- binding EDSL used to expressively create key bindings and subbindings inside
-- a KeysM () context.
--

data KeyMatcher = KeyMatcher Word32 Char
  deriving (Show)

-- | Like a KeyMatcher, but allows additional modifiers to be pressed, not just
-- the exact ones given.
newtype WeakKeyMatcher = WeakKeyMatcher KeyMatcher

-- | Converts a KeyMatcher to a weak key matcher.
weak :: KeyMatcher -> WeakKeyMatcher
weak = WeakKeyMatcher

class KeyMatcherId r where
  toKeyMatcher :: r -> KeyMatcher

instance KeyMatcherId KeyMatcher where
  toKeyMatcher = id

instance KeyMatcherId Char where
  toKeyMatcher = KeyMatcher 0

class KeyMatcherBuilder b where
  (.+) :: (KeyMatcherId i) => b -> i -> KeyMatcher

instance KeyMatcherBuilder Modifier where
  (.+) m (toKeyMatcher -> (KeyMatcher mods ch)) =
    KeyMatcher (mods .|. modifierToMask m) ch

infixr 9 .+

class MatchKey m where
  matchKey :: m -> KeyEvent -> Bool

instance MatchKey (KeyEvent -> Bool) where
  matchKey = ($)

instance MatchKey Bool where
  matchKey = const

instance MatchKey Char where
  matchKey ch ev = ch == KeyEvent.codepoint ev

instance MatchKey KeyMatcher where
  matchKey (KeyMatcher m ch) ev =
    ch == KeyEvent.codepoint ev && m == KeyEvent.modifiers ev

instance MatchKey WeakKeyMatcher where
  matchKey (WeakKeyMatcher (KeyMatcher m ch)) ev =
    ch == KeyEvent.codepoint ev && (m .|. ms) == ms
    where
      ms = KeyEvent.modifiers ev

class IsKeysM m where
  toKeysM :: m a -> KeysM a

instance IsKeysM W where
  toKeysM = liftW

instance IsKeysM KeysM where
  toKeysM = id

bind :: (MatchKey m, IsKeysM k) => KeyEvent -> m -> k () -> KeysM ()
bind ev m act = do
  when (matchKey m ev) $ do
    toKeysM act
    continue

ignoreReleaseEvents :: KeyEvent -> KeysM ()
ignoreReleaseEvents ev = do
  when (KeyEvent.state ev /= KeyEvent.KeyPressed) $ do
    forwardEvent ev
    continue