{-# 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 request loop and server state.
--
-- A wayland server provides a "display" that can be controlled using a UNIX socket. The server runs an event loop
-- handling requests. These requests may be events, requests to create objects, or requests for a particular object
-- instance.

module Wayland.Server.Display
  ( -- * The Display Monad
    DisplayM
  , withDisplay
  , bindSocket
  ,  -- * Processing Requests
    RequestHandler
  , processRequests
  , requestTermination
  , -- * Managing Objects
    Global
  , registerGlobal
  , unregisterGlobal
  , createObject
  , destroyObject
  , sendEvent
  , flushEvents
  , -- * Event Sources
    EventSource
  , registerSignalEvent
  ) where

import Data.Maybe (isJust)
import Data.Functor (($>))
import Data.Typeable (Proxy(..))
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Foreign (Ptr, FunPtr, Word32, nullPtr)
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)
import Control.Monad (when)
import Wayland.Server.System
import Wayland.Server.Object (ObjectCache, Request(..), newObjectCache, destroyObjectCache, createResource, RequestClass, Object(..), unwrapResource, resolveTypeInterface, sendResourceEvent, NewObject, createNewObject)
import Wayland.Server.Client (wrapClient)

-- | The display monad.
--
-- All server-related logic runs in this monad, including the initial set-up and all request 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 request.
--
-- This handler is executed as part of the event loop.
type RequestHandler s = Request -> DisplayM s ()

-- | Read-only data related to the display.
data Display
  = Display { display         :: !WLDisplay                 -- ^ The underlying wl_display instance.
            , activeEventLoop :: !(IORef (Maybe EventLoop)) -- ^ Mutable reference to the active request handler.
            }

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

-- | A function handling a request in the IO monad.
--
-- The event loop callback needs to retrieve the current state, run the request handler in a Display monad, and update
-- the state.
newtype EventLoop = EventLoop { handleEvent :: Request -> IO () }

-- | Handle to an event source that dispatches requests to the event loop.
--
-- This may be a signal event source, timer event source, etc. The event source handle remains valid even after it has
-- expired or after it was unregistered.
data EventSource = EventSource (FunPtr WLEventLoopSignalFunc) (IORef WLEventSource)
  deriving (Eq)

-- | 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 was
-- unregistered.
data Global = Global (FunPtr WLGlobalBindFunc) (IORef WLGlobal)
  deriving (Eq)

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)

-- | 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 display

-- | Get a function which can be used to handle a request.
--
-- This function can be used even before the event loop is started as it will use the mutable "activeEventLoop"
-- reference which gets updated when the event loop is started.
getRequestHandler :: DisplayM s (Request -> IO ())
getRequestHandler = do
  loopRef <- liftDisplay $ asks activeEventLoop
  pure $ \req -> do
    loop <- readIORef loopRef
    case loop of
      Just handler -> handleEvent handler req
      Nothing      -> putStrLn $ "Lost request because no event loop is active: " ++ show req

-- | 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
  eventLoop <- newIORef Nothing
  let disp = Display { display         = sysDisp
                     , activeEventLoop = eventLoop
                     }
  objectCache <- newObjectCache
  let dispState = DisplayState { _stateObjectCache = objectCache }
  (res, dispState') <- runStateT (runReaderT (runStateT (runDisplayM f) initialState) disp) dispState
  destroyObjectCache (_stateObjectCache dispState')
  pure res

-- | 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 event loop for the display.
--
-- The display will then handle all requests using the provided handler until 'requestTermination' is called. The state
-- may be modified by the request handler.
processRequests :: RequestHandler s -> DisplayM s ()
processRequests h = do
  disp <- liftDisplay ask
  stateBefore <- get
  displayStateBefore <- liftDisplayState get
  stateRef <- liftIO $ newIORef (stateBefore, displayStateBefore)
  let handler :: Request -> IO ()
      handler r = do
        (s, ds) <- readIORef stateRef
        ((_, s'), ds') <- runStateT (runReaderT (runStateT (runDisplayM (h r)) s) disp) ds
        writeIORef stateRef (s', ds')
  loop <- liftIO $ readIORef (activeEventLoop disp)
  when (isJust loop) (error "Cannot start an event loop from within an event handler.")
  liftIO $ writeIORef (activeEventLoop disp) (Just (EventLoop handler))
  liftIO $ wlDisplayRun (display disp)
  (stateAfter, displayStateAfter) <- liftIO $ readIORef stateRef
  put stateAfter
  liftDisplayState $ put displayStateAfter

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

-- | Create a new object instance.
--
-- If the NewObject instance has already been used to create an instance, the same instance is returned again.
createObject
  :: Object o
  => RequestClass destroyReq
  => NewObject o        -- ^ Handle describing the object to be created
  -> (o -> destroyReq)  -- ^ Function creating a request dispatched after the object was destroyed.
  -> DisplayM s o       -- ^ Display monad wrapping the newly created object.
createObject o destroyHook = do
  handle <- getRequestHandler
  liftDisplayState $ zoom stateObjectCache $ createResource handle 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 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" request will be dispatched in the request loop.
registerGlobal
  :: forall o req s. Object o
  => RequestClass req
  => (NewObject o -> req) -- ^ Function used to create a request to bind the global object for a client.
  -> Int                  -- ^ Version of the global published.
  -> DisplayM s Global    -- ^ Display monad wrapping a handle to the global which can be used to unregister it.
registerGlobal mkReq version = do
  disp <- getDisplay
  handle <- getRequestHandler
  iface <- liftDisplayState $ zoom stateObjectCache $ resolveTypeInterface (Proxy :: Proxy o)
  let bind :: WLClient -> Ptr () -> Word32 -> Word32 -> IO ()
      bind client _ ver objID = do wrappedClient <- wrapClient client
                                   obj <- createNewObject wrappedClient objID (fromIntegral ver)
                                   handle $ Request (mkReq obj)
  liftIO $ do
    bindFunc <- wrapWLGlobalBindFunc bind
    global <- wlGlobalCreate disp iface (fromIntegral version) nullPtr bindFunc
    globalRef <- newIORef global
    -- TODO: Free func
    pure $ Global bindFunc globalRef

-- | Unregister a global after it is no longer available.
unregisterGlobal :: Global -> DisplayM s ()
unregisterGlobal _ = pure ()

-- | Register an event to be dispatched in the request loop when a UNIX signal is sent to the server process.
registerSignalEvent
  :: RequestClass req
  => req                    -- ^ The request dispatched when the signal is received.
  -> CInt                   -- ^ The code of the UNIX signal to listen for.
  -> DisplayM s EventSource -- ^ Display monad wrapping a handle to the event source which can be used to unregister it.
registerSignalEvent r sig = do
  disp <- getDisplay
  handle <- getRequestHandler
  liftIO $ do
    loop <- wlDisplayGetEventLoop disp
    func <- wrapWLEventLoopSignalFunc $ \_ _ -> handle (Request r) $> 0
    src <- wlEventLoopAddSignal loop sig func nullPtr
    srcRef <- newIORef src
    -- TODO: Free func
    pure $ EventSource func srcRef