{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module     : Wayland.Server.Display
-- Description: Entrypoint for the server, managing the message loop and server state.
--
-- A wayland server provides a "display" that can be controlled using a UNIX socket. The server runs a message loop
-- handling messages. These messages may be timers, UNIX signals, requests to create objects, or requests for a
-- particular object instance.

module Wayland.Server.Display
  ( -- * The Display Monad
    DisplayM
  , withDisplay
  , bindSocket
  ,  -- * Processing Messages
    MessageHandler
  , processMessages
  , requestTermination
  , -- * Managing Clients
    destroyClient
  , flushClientEvents
  , -- * Managing Objects
    Global
  , registerGlobal
  , hideGlobal
  , unregisterGlobal
  , createObject
  , destroyObject
  , sendEvent
  , flushEvents
  , -- * Message Sources
    SignalSource
  , TimerSource
  , FileSource
  , FileEventMask(..)
  , registerSignalSource
  , unregisterSignalSource
  , registerTimerSource
  , setTimerSource
  , unregisterTimerSource
  , registerFileSource
  , updateFileSourceMask
  , unregisterFileSource
  ) where

import Data.Maybe (isJust)
import Data.Functor (($>))
import Data.Typeable (Proxy(..))
import Data.IORef (IORef, newIORef, readIORef, writeIORef, modifyIORef)
import Data.Hashable (Hashable(..))
import Data.Time.Clock (DiffTime, diffTimeToPicoseconds)
import Data.Bits ((.&.), (.|.))
import Foreign (Ptr, FunPtr, Word32, nullPtr, freeHaskellFunPtr, free, new, castPtr)
import Foreign.C (peekCString, CInt)
import Lens.Micro.TH (makeLenses)
import Lens.Micro.Mtl.Internal (Zoom(..), Zoomed)
import Control.Monad.State.Strict (StateT(..), MonadState(..))
import Control.Monad.Reader (ReaderT(..), ask, asks)
import Control.Monad.Trans (MonadIO(..), MonadTrans(..))
import Control.Exception (bracket, bracket_)
import Control.Monad (when, void)
import Wayland.Server.System
import Wayland.Server.Object
import Wayland.Server.Client (wrapClient, ClientData, newClientData, destroyClientData, ClientMessage(..), createClient, Client, unwrapClient)
import Wayland.Server.Types (Message(..), MessageClass, Handle, createHandle, resolveHandle, invalidateHandle)
import Wayland.Server.Interface (InterfaceDesc(..))
import System.Posix.Types (Fd (..))

-- | The display monad.
--
-- All server-related logic runs in this monad, including the initial set-up and all message handlers. The monad
-- provides custom state and support for IO operations.
newtype DisplayM s a = DisplayM { runDisplayM :: StateT s (ReaderT Display (StateT DisplayState IO)) a }
  deriving (Functor, Applicative, Monad, MonadIO, MonadState s)

-- | A function handling a type-erased message.
--
-- This handler is executed as part of the message loop.
type MessageHandler s = Message -> DisplayM s ()

-- | Read-only data related to the display.
data Display
  = Display { dispDisplay           :: !WLDisplay                   -- ^ The underlying wl_display instance.
            , dispActiveMessageLoop :: !(IORef (Maybe MessageLoop)) -- ^ Mutable reference to the active message handler.
            , dispClientData        :: !ClientData
            }

-- | Mutable state for the display.
newtype DisplayState = DisplayState { _stateObjectCache :: ObjectCache }

-- | A function handling a message in the IO monad.
--
-- The message loop callback needs to retrieve the current state, run the message handler in a Display monad, and update
-- the state.
newtype MessageLoop = MessageLoop (Message -> IO ())

-- | Handle to a message source that dispatches messages when a UNIX signal was raised.
--
-- The message source handle remains valid even after it is unregistered. However, actions using the source (like
-- unregistering it) will fail.
data SignalSource
  = SignalSource { signalHandle :: !(Handle WLEventSource)
                 , signalFunc   :: !(FunPtr WLEventLoopSignalFunc)
                 }

-- | Handle to a message source that dispatches messages after a timeout has passed.
--
-- The timer can be reused and/or reset multiple times before it is unregistered.
--
-- The message source handle remains valid even after it is unregistered. However, actions using the source (like
-- unregistering it) will fail.
data TimerSource
  = TimerSource { timerHandle :: !(Handle WLEventSource)
                , timerFunc   :: !(FunPtr WLEventLoopTimerFunc)
                }

-- | Handle to a message source that dispatches messages for file descriptor events.
--
-- The message source handle remains valid even after it is unregistered. However, actions using the source (like
-- unregistering it) will fail.
data FileSource
  = FileSource { fileHandle :: !(Handle WLEventSource)
               , fileFunc   :: !(FunPtr WLEventLoopFdFunc)
               }

-- | The file events listened for by a 'FileSource' on a UNIX file descriptor.
data FileEventMask
  = FileEventMask { fileEventReadable :: Bool
                  , fileEventWritable :: Bool
                  , fileEventHangup   :: Bool
                  , fileEventError    :: Bool
                  }
  deriving (Show, Eq, Ord)

-- | Handle to a global object shared between all clients.
--
-- For each client, the global can be "bound" to a specific object.
--
-- The global handle remains valid even after it is destroyed. However, actions using the global (like unregistering it)
-- will fail.
data Global
  = Global { glbHandle   :: Handle WLGlobal
           , glbBindFunc :: FunPtr WLGlobalBindFunc
           }

makeLenses ''DisplayState

type instance Zoomed (DisplayM s) = Zoomed (StateT s (ReaderT Display (StateT DisplayState IO)))

instance Zoom (DisplayM s) (DisplayM t) s t where
  zoom l (DisplayM m) = DisplayM (zoom l m)

instance Eq SignalSource where
  a == b = signalHandle a == signalHandle b

instance Ord SignalSource where
  a <= b = signalHandle a <= signalHandle b

instance Hashable SignalSource where
  hashWithSalt s o = hashWithSalt s (signalHandle o)

instance Show SignalSource where
  show o = "SignalSource " ++ show (signalHandle o)

instance Eq TimerSource where
  a == b = timerHandle a == timerHandle b

instance Ord TimerSource where
  a <= b = timerHandle a <= timerHandle b

instance Hashable TimerSource where
  hashWithSalt s o = hashWithSalt s (timerHandle o)

instance Show TimerSource where
  show o = "TimerSource " ++ show (timerHandle o)

instance Eq FileSource where
  a == b = fileHandle a == fileHandle b

instance Ord FileSource where
  a <= b = fileHandle a <= fileHandle b

instance Hashable FileSource where
  hashWithSalt s o = hashWithSalt s (fileHandle o)

instance Show FileSource where
  show o = "FileSource " ++ show (fileHandle o)

instance Eq Global where
  a == b = glbHandle a == glbHandle b

instance Ord Global where
  a <= b = glbHandle a <= glbHandle b

instance Hashable Global where
  hashWithSalt s o = hashWithSalt s (glbHandle o)

instance Show Global where
  show o = "Global " ++ show (glbHandle o)

-- | Lift a 'Display' reader operation into the display monad.
liftDisplay :: ReaderT Display (StateT DisplayState IO) a -> DisplayM s a
liftDisplay = DisplayM . lift

-- | Lift a 'DisplayState' state operation into the display monad.
liftDisplayState :: StateT DisplayState IO a -> DisplayM s a
liftDisplayState = DisplayM . lift . lift

-- | Get a reference to the underlying wl_display instance.
getDisplay :: DisplayM s WLDisplay
getDisplay = liftDisplay $ asks dispDisplay

-- | Get a function which can be used to handle a message.
--
-- This function can be used even before the message loop is started as it will use the mutable "activeMessageLoop"
-- reference which gets updated when the message loop is started.
getMessageHandler :: DisplayM s (Message -> IO ())
getMessageHandler = liftDisplay $ asks (handleMessage . dispActiveMessageLoop)

-- | Handle a message within a display's message loop.
--
-- The message loop supports reentrant calls. In this case, dispatching the message is delayed until after the current
-- message has been processed.
handleMessage :: IORef (Maybe MessageLoop) -> Message -> IO ()
handleMessage loopRef msg = do
  loop <- readIORef loopRef
  case loop of
    Just (MessageLoop handler) -> handler msg
    Nothing                    -> putStrLn $ "Lost message because no message loop is active: " ++ show msg

-- | Create a new display and execute the given action on it.
--
-- After executing the action, all resources associated with the display are cleaned up.
withDisplay
  :: s            -- ^ Initial state.
  -> DisplayM s a -- ^ The display action to execute.
  -> IO (a, s)    -- ^ IO monad returning the result and the resulting state.
withDisplay initialState f = bracket wlDisplayCreate wlDisplayDestroy $ \sysDisp -> do
  messageLoop <- newIORef Nothing
  clientData  <- newClientData (handleMessage messageLoop)
  objectCache <- newObjectCache
  let disp = Display { dispDisplay           = sysDisp
                     , dispActiveMessageLoop = messageLoop
                     , dispClientData        = clientData
                     }
  let dispState = DisplayState { _stateObjectCache = objectCache }
  (res, dispState') <- handleClientCreated disp $
    runStateT (runReaderT (runStateT (runDisplayM f) initialState) disp) dispState
  destroyObjectCache (_stateObjectCache dispState')
  destroyClientData clientData
  pure res

-- | Handle the "client created" event, creating a wrapper and dispatching a message.
handleClientCreated :: Display -> IO a -> IO a
handleClientCreated disp action =
  let clientCreated :: Ptr WLListener -> Ptr () -> IO ()
      clientCreated _ c = do
        client <- runReaderT (createClient (WLClient (castPtr c))) (dispClientData disp)
        handleMessage (dispActiveMessageLoop disp) $ Message $ ClientCreated client
  in bracket (wrapWLNotifyFunc clientCreated) freeHaskellFunPtr $ \clientCreatedFunc ->
    bracket (new $ WLListener clientCreatedFunc) free $ \clientListener ->
      bracket_ (wlDisplayAddClientCreatedListener (dispDisplay disp) clientListener) (wlListenerRemove clientListener) action

-- | Bind the display to a UNIX socket so it is reachable from clients.
--
-- This will return the name of the Wayland display.
bindSocket :: DisplayM s String
bindSocket = do
  disp <- getDisplay
  name <- liftIO $ wlDisplayAddSocketAuto disp
  liftIO $ peekCString name

-- | Start the message loop for the display.
--
-- The display will then handle all messages using the provided handler until 'requestTermination' is called. The state
-- may be modified by the message handler.
processMessages :: MessageHandler s -> DisplayM s ()
processMessages h = do
  disp <- liftDisplay ask
  stateBefore <- get
  displayStateBefore <- liftDisplayState get
  stateRef <- liftIO $ newIORef $ Just (stateBefore, displayStateBefore)
  queued   <- liftIO $ newIORef []
  let snoc r (x:xs) = x : snoc r xs
      snoc r []     = [r]
  let runQueued :: MessageHandler s -> DisplayM s ()
      runQueued h' = do
        x <- liftIO $ readIORef queued
        case x of
          (r:rs) -> do liftIO $ writeIORef queued rs
                       h' r
                       runQueued h'
          []    -> pure ()
  let handler :: Message -> IO ()
      handler r = do
        stateVal <- readIORef stateRef
        modifyIORef queued (snoc r) -- In most cases, queued is empty, so we don't really care about snoc's runtime.
        case stateVal of
          Just (s, ds) -> do writeIORef stateRef Nothing
                             ((_, s'), ds') <- runStateT (runReaderT (runStateT (runDisplayM (runQueued h)) s) disp) ds
                             writeIORef stateRef $ Just (s', ds')
          Nothing      -> pure () -- Reentrant call, dispatch after completing the current message.
  loop <- liftIO $ readIORef (dispActiveMessageLoop disp)
  when (isJust loop) (error "Cannot start a message loop from within a message handler.")
  liftIO $ writeIORef (dispActiveMessageLoop disp) (Just (MessageLoop handler))
  liftIO $ wlDisplayRun (dispDisplay disp)
  stateVal <- liftIO $ readIORef stateRef
  case stateVal of
    Just (stateAfter, displayStateAfter) -> do put stateAfter
                                               liftDisplayState $ put displayStateAfter
    Nothing                              -> pure ()

-- | Request the message loop to stop.
--
-- The 'processMessages' will return following this action.
requestTermination :: DisplayM s ()
requestTermination = do
  disp <- getDisplay
  liftIO $ wlDisplayTerminate disp

-- | Destroy a client instance.
destroyClient :: Client -> DisplayM s ()
destroyClient c = liftIO (unwrapClient c >>= wlClientDestroy)

-- | Ensure that all queued events are sent to a particular client.
flushClientEvents :: Client -> DisplayM s ()
flushClientEvents c = liftIO (unwrapClient c >>= wlClientFlush)

-- | Register a global object which can be bound by clients.
--
-- The object is published through the global wl_registry instance. When a client wants to instantiate the global
-- object, the given "bind" message will be dispatched in the message loop.
registerGlobal
  :: forall o bindMsg s. Object o
  => MessageClass bindMsg
  => (Global -> NewObject o -> bindMsg) -- ^ Function used to create a message to bind the global object for a client.
  -> DisplayM s Global                  -- ^ Display monad wrapping a handle to the global which can be used to unregister it.
registerGlobal mkBindMsg = do
  disp    <- getDisplay
  handler <- getMessageHandler
  iface   <- liftDisplayState $ zoom stateObjectCache $ resolveTypeInterface (Proxy :: Proxy o) handler
  liftIO $ do
    ref <- newIORef Nothing
    let version = Wayland.Server.Interface.ifVersion (describeInterface (Proxy :: Proxy o))
    let bind :: WLClient -> Ptr () -> Word32 -> Word32 -> IO ()
        bind client _ _ objID = do wrappedClient <- wrapClient client
                                   Just global   <- readIORef ref
                                   obj           <- createNewObject wrappedClient objID
                                   handler $ Message (mkBindMsg global obj)
    bindFunc <- wrapWLGlobalBindFunc bind
    handle   <- wlGlobalCreate disp iface (fromIntegral version) nullPtr bindFunc >>= createHandle
    let res = Global handle bindFunc
    writeIORef ref (Just res)
    pure res

-- | Stop advertizing a global object to client.
--
-- The global will continue to exist until it is destroyed using 'destroyGlobal'. This is useful for waiting for a
-- timeout to elapse before actually destroying the global.
hideGlobal :: Global -> DisplayM s ()
hideGlobal global = liftIO (resolveHandle (glbHandle global) >>= wlGlobalRemove)

-- | Unregister a global once it is no longer available.
--
-- The global and all resources associated with it will be destroyed.
unregisterGlobal :: Global -> DisplayM s ()
unregisterGlobal global = liftIO $ do
  resolveHandle (glbHandle global) >>= wlGlobalDestroy
  invalidateHandle (glbHandle global)
  freeHaskellFunPtr (glbBindFunc global)

-- | Create a new object instance.
--
-- It is a programmer error to use a NewObject instance more than once.
createObject
  :: Object o
  => MessageClass destroyMsg
  => NewObject o        -- ^ Handle describing the object to be created
  -> (o -> destroyMsg)  -- ^ Function creating a message dispatched after the object was destroyed.
  -> DisplayM s o       -- ^ Display monad wrapping the newly created object.
createObject o destroyHook = do
  handler <- getMessageHandler
  liftDisplayState $ zoom stateObjectCache $ createResource handler o destroyHook

-- | Destroy an object instance.
--
-- This will invoke the object's "destroy" hook.
destroyObject :: Object o => o -> DisplayM s ()
destroyObject o = liftIO $ do
  ptr <- unwrapResource o
  wlResourceDestroy ptr

-- | Send an event for an object to the client it is associated with.
sendEvent :: Object o => o -> ObjectEvent o -> DisplayM s ()
sendEvent o e = liftIO $ sendResourceEvent o e

-- | Ensure that all queued events are sent out to clients.
flushEvents :: DisplayM s ()
flushEvents = do
  disp <- getDisplay
  liftIO $ wlDisplayFlushClients disp

-- | Register a message to be dispatched in the message loop when a UNIX signal is sent to the server process.
registerSignalSource
  :: MessageClass msg
  => (SignalSource -> msg)   -- ^ Function used to create the message when the signal is received.
  -> CInt                    -- ^ The code of the UNIX signal to listen for.
  -> DisplayM s SignalSource -- ^ Display monad wrapping a handle to the newly registered message source.
registerSignalSource mkMsg sig = do
  disp    <- getDisplay
  handler <- getMessageHandler
  liftIO $ do
    ref    <- newIORef Nothing
    loop   <- wlDisplayGetEventLoop disp
    func   <- wrapWLEventLoopSignalFunc $ \_ _ -> do
      Just src <- readIORef ref
      handler (Message (mkMsg src)) $> 0
    handle <- wlEventLoopAddSignal loop sig func nullPtr >>= createHandle
    let res = SignalSource handle func
    writeIORef ref (Just res)
    pure res

-- | Unregister a signal message source and stop it from dispatching any more messages.
unregisterSignalSource :: SignalSource -> DisplayM s ()
unregisterSignalSource signal = liftIO $ do
  _ <- resolveHandle (signalHandle signal) >>= wlEventSourceRemove
  invalidateHandle (signalHandle signal)
  freeHaskellFunPtr (signalFunc signal)

-- | Register a message to be dispatched using a timer.
--
-- Initially, the timer is not set. Use 'setTimerSource' to arm/disarm the timer later.
registerTimerSource
  :: MessageClass msg
  => (TimerSource -> msg)   -- ^ Function used to create a message on timeout.
  -> DisplayM s TimerSource -- ^ Display monad wrapping a handle to the newly registered message source.
registerTimerSource mkMsg = do
  disp    <- getDisplay
  handler <- getMessageHandler
  liftIO $ do
    ref    <- newIORef Nothing
    loop   <- wlDisplayGetEventLoop disp
    func   <- wrapWLEventLoopTimerFunc $ \_ -> do
      Just src <- readIORef ref
      handler (Message (mkMsg src)) $> 0
    handle <- wlEventLoopAddTimer loop func nullPtr >>= createHandle
    let res = TimerSource handle func
    writeIORef ref (Just res)
    pure res

-- | Arm/disarm a timer message source.
--
-- Any active timeout is stopped. If the specified timeout is not Nothing, a new timeout is started.
setTimerSource :: TimerSource -> Maybe DiffTime -> DisplayM s ()
setTimerSource timer timeout = liftIO $ do
  src <- resolveHandle (timerHandle timer)
  let msDelay = case timeout of
                  Just diff -> fromIntegral $ max 1 (diffTimeToPicoseconds diff `div` 1000000000)
                  Nothing   -> 0
  void $ wlEventSourceTimerUpdate src msDelay

-- | Unregister a timer message source and stop it from dispatching any more messages.
unregisterTimerSource :: TimerSource -> DisplayM s ()
unregisterTimerSource timer = liftIO $ do
  _ <- resolveHandle (timerHandle timer) >>= wlEventSourceRemove
  invalidateHandle (timerHandle timer)
  freeHaskellFunPtr (timerFunc timer)

-- | Register a message to be dispatched for file events on a UNIX file descriptor.
--
-- Initially, file descriptors are usually writable so the event should only be listened for after a write fails.
registerFileSource
  :: MessageClass msg
  => (FileSource -> FileEventMask -> msg) -- ^ Function used to create a message following one or more file events (as specified).
  -> Fd                                   -- ^ The UNIX file descriptor to watch.
  -> FileEventMask                        -- ^ Events to watch for.
  -> DisplayM s FileSource                -- ^ Display monad wrapping a handle to the newly registered message source.
registerFileSource mkMsg (Fd fd) events = do
  disp    <- getDisplay
  handler <- getMessageHandler
  liftIO $ do
    ref    <- newIORef Nothing
    loop   <- wlDisplayGetEventLoop disp
    func   <- wrapWLEventLoopFdFunc $ \_ e _ -> do
      Just src <- readIORef ref
      handler (Message (mkMsg src (toFileEventMask e))) $> 0
    handle <- wlEventLoopAddFd loop fd (fromFileEventMask events) func nullPtr >>= createHandle
    let res = FileSource handle func
    writeIORef ref (Just res)
    pure res

-- | Update the file events to watch for using an existing file message source.
updateFileSourceMask :: FileSource -> FileEventMask -> DisplayM s ()
updateFileSourceMask file events = liftIO $ do
  src <- resolveHandle (fileHandle file)
  void $ wlEventSourceFdUpdate src (fromFileEventMask events)

-- | Unregister a file message source and stop it from dispatching any more messages.
unregisterFileSource :: FileSource -> DisplayM s ()
unregisterFileSource file = liftIO $ do
  _ <- resolveHandle (fileHandle file) >>= wlEventSourceRemove
  invalidateHandle (fileHandle file)
  freeHaskellFunPtr (fileFunc file)

fromFileEventMask :: FileEventMask -> Word32
fromFileEventMask events = conditionalMask wlEventReadable (fileEventReadable events)
                         $ conditionalMask wlEventWritable (fileEventWritable events)
                         $ conditionalMask wlEventHangup   (fileEventHangup   events)
                         $ conditionalMask wlEventError    (fileEventError    events) 0
  where conditionalMask mask enabled x = if enabled then x .|. mask else x

toFileEventMask :: Word32 -> FileEventMask
toFileEventMask bits = FileEventMask { fileEventReadable = testMask wlEventReadable
                                     , fileEventWritable = testMask wlEventWritable
                                     , fileEventHangup   = testMask wlEventHangup
                                     , fileEventError    = testMask wlEventError
                                     }
  where testMask mask = (bits .&. mask) /= 0