{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ExistentialQuantification #-}

-- |
-- Module     : Wayland.Server.Object
-- Description: Low-level API for wrapping Wayland object/resource handles.
--
-- Objects are represented as wl_resources in libwayland. This module wraps a resource, providing an opaque handle to it
-- that is valid even after the underlying resource is destroyed. This is useful for storing the resource handle in a
-- map and associating other data with it, cleaning it up only after its destruction.

module Wayland.Server.Object
  ( -- * Handles
    Object(..)
  , AnyObject
  , NewObject
  , createNewObject
  , RequestClass
  , EventClass
  , ArgumentMarshaller(..)
  , ArgumentDemarshaller(..)
  , -- * State
    ObjectCache
  , ObjectCacheState
  , newObjectCache
  , destroyObjectCache
  , resolveTypeInterface
  , -- * Lifecycle
    unwrapResource
  , wrapResource
  , createResource
  , sendResourceEvent
  ) where

import System.Posix.Types (Fd(..))
import Data.Kind (Type)
import Data.Hashable (Hashable(..))
import Data.IORef (readIORef, newIORef, IORef, writeIORef)
import Data.Typeable (Proxy(..), TypeRep, typeRep, Typeable)
import Foreign
import Foreign.C (withCString, peekCString, CInt)
import Lens.Micro.Mtl (use, (%=), zoom)
import Lens.Micro.TH (makeLenses)
import Control.Monad.State.Strict (StateT(..), MonadState(..))
import Control.Monad.Cont (ContT(..))
import Control.Monad.Trans (MonadIO(..), MonadTrans(..))
import Control.Monad.Reader (ReaderT (..), MonadReader(..))
import Wayland.Server.System
import Wayland.Server.Interface (MessageDesc(..), InterfaceDesc(..), InterfaceCache, newInterfaceCache, destroyInterfaceCache, resolveInterface)
import Wayland.Server.Client (Client, unwrapClient, wrapClient)
import Wayland.Server.Types (Handle, invalidateHandle, resolveHandle, createHandle, MessageClass, Message(..))
import qualified Data.Map as M

-- | The primary data type for typed object handles.
--
-- An instance of the Object class wraps an opaque object handle ('AnyObject') and additionally provides information
-- about the object's interface.
class (Typeable o, RequestClass (ObjectRequest o), EventClass (ObjectEvent o)) => Object o where
  -- | Associated type for requests sent to object instances.
  type ObjectRequest o :: Type
  -- | Associated type for events emitted by object instances.
  type ObjectEvent o :: Type

  -- | Unwrap this typed object into an opaque object handle.
  toAny :: o -> AnyObject
  -- | Create a typed object from an opaque object handle.
  fromAny :: AnyObject -> o

  -- | Provide an interface description for objects of this type.
  describeInterface :: Proxy o -> InterfaceDesc
  -- | Provide a message description for an event instance.
  eventSignature :: Proxy o -> ObjectEvent o -> MessageDesc
  -- | Provide a message description for a request instance.
  requestSignature :: Proxy o -> ObjectRequest o -> MessageDesc

  -- | Marshal the arguments provided in the event instance into an argument marshaller.
  --
  -- This should marshal exactly @length . msgArguments . eventSignature@ arguments for the passed event and match the
  -- types described in the signature. After that, the opcode for the event should be returned.
  marshalEvent
    :: ArgumentMarshaller m -- ^ Monad used to write marshalled arguments.
    => Proxy o              -- ^ Proxy for the object type this event is referring to.
    -> ObjectEvent o        -- ^ The event that should be marshalled.
    -> m Int                -- ^ Argument marshaller returning the event's opcode.

  -- | Demarshal the arguments provided into a request instance using an argument demarshaller.
  --
  -- This should demarshal exactly @length . msgArguments . requestSignature@ arguments using the demarshaller interface
  -- and match the types described in the signature.
  demarshalRequest
    :: ArgumentDemarshaller m -- ^ Monad used to extract marshalled arguments.
    => o                      -- ^ The instance this request is referring to.
    -> Int                    -- ^ The opcode of the request.
    -> m (ObjectRequest o)    -- ^ Argument demarshaller returning the demarshalled request.

-- | A request for an object sent by a client.
--
-- The request usually includes the object it is referring to, as well as arguments provided by the client. It is
-- processed as part of the server's message loop.
class MessageClass r => RequestClass r where

-- | An event that can be emitted on an object by the server.
--
-- The event stores the arguments passed to the client.
class MessageClass e => EventClass e where

-- | A monad providing an environment for marshalling a fixed number of arguments of various types.
--
-- For libwayland, the arguments are stored into wl_arguments. It is considered a programmer error to write more
-- arguments than expected by an event or different argument types than the ones specified.
class Monad m => ArgumentMarshaller m where
  marshalInt    :: Int                     -> m ()
  marshalUInt   :: Word                    -> m ()
  marshalFixed  :: Float                   -> m ()
  marshalString :: String                  -> m ()
  marshalObject :: Object o => o           -> m ()
  marshalNewID  :: Object o => NewObject o -> m ()
  marshalArray  :: Storable a => [a]       -> m ()
  marshalFD     :: Fd                      -> m ()

-- | A monad providing an environment for reading a fixed number of arguments of various types.
--
-- For libwayland, the arguments are read from wl_arguments. It is considered a programmer error to read more arguments
-- than expected by a request or different argument types than the ones specified.
class Monad m => ArgumentDemarshaller m where
  demarshalInt    :: m Int
  demarshalUInt   :: m Word
  demarshalFixed  :: m Float
  demarshalString :: m String
  demarshalObject :: Object o => m o
  demarshalNewID  :: Object o => m (NewObject o)
  demarshalArray  :: Storable a => m [a]
  demarshalFD     :: m Fd

-- | An opaque handle to a Wayland object.
--
-- The handle remains valid and can be used for equality checks even after the underlying libwayland resource was
-- destroyed. However, actions using the object (like sending events to it) will fail.
--
-- In addition, a custom type-erased cleanup action is stored which is invoked once the resource is destroyed.
data AnyObject
  = AnyObject { objHandle      :: Handle WLResource  -- ^ The underlying resource.
              , objDestroyHook :: AnyObject -> IO () -- ^ Function executed after the object was destroyed.
              }

-- | Handle to an object which has not yet been created but has an ID assigned by the client.
--
-- This handle can be used to create an object. It is a programmer error to use it more than once.
data NewObject o
  = NewObject { newClient  :: !Client            -- ^ The client this object is associated with.
              , newID      :: !Word32            -- ^ Object ID of the new object.
              , newObject  :: !(IORef (Maybe o)) -- ^ Reference to the object, assigned when it is created.
              }

-- | Data related to an object type.
data TypeInfo
  = TypeInfo { typeInterface    :: !(Ptr WLInterface)         -- ^ Reference to the libwayland interface for the type.
             , typeDispatchFunc :: !(FunPtr WLDispatcherFunc) -- ^ The dispatcher function handling requests for the type.
             }

-- | State used to cache object-related data.
data ObjectCache
  = ObjectCache { _cacheTypeInfo    :: !(M.Map TypeRep TypeInfo)       -- ^ Type-specific information.
                , _cacheDestroyFunc :: !(FunPtr WLResourceDestroyFunc) -- ^ Global destroy function for all objects.
                , _cacheInterfaces  :: !InterfaceCache                 -- ^ Cache for native interfaces.
                }

-- | A state monad providing an instance of 'ObjectCache'.
type ObjectCacheState = StateT ObjectCache IO

-- | State used when iterating over a buffer of wl_arguments.
newtype ArgumentIterator = ArgumentIterator (Ptr WLArgument)

-- | Monad used for writing arguments into a wl_argument buffer.
--
-- The monad wraps a ContT monad to allow using alloca when wrapping arguments.
newtype ArgumentWriter r a = ArgumentWriter { runArgumentWriter :: StateT ArgumentIterator (ContT r IO) a }
  deriving (Functor, Applicative, Monad, MonadState ArgumentIterator, MonadIO)

-- | Monad used for reading arguments from a wl_argument buffer.
newtype ArgumentReader a = ArgumentReader { runArgumentReader :: StateT ArgumentIterator (ReaderT Client IO) a }
  deriving (Functor, Applicative, Monad, MonadState ArgumentIterator, MonadIO, MonadReader Client)

makeLenses ''ObjectCache

instance Eq AnyObject where
  a == b = objHandle a == objHandle b

instance Ord AnyObject where
  a <= b = objHandle a <= objHandle b

instance Hashable AnyObject where
  hashWithSalt s o = hashWithSalt s (objHandle o)

instance Show AnyObject where
  show o = show (objHandle o)

instance Show (NewObject o) where
  show o = "NewObject " ++ show (newClient o) ++ " " ++ show (newID o)

-- | Create a handle to an object that has yet to be created.
createNewObject :: Client -> Word32 -> IO (NewObject o)
createNewObject client objID = do
  obj <- newIORef Nothing
  pure $ NewObject { newClient  = client
                   , newID      = objID
                   , newObject  = obj
                   }

-- | Initialize an object cache.
newObjectCache :: IO ObjectCache
newObjectCache = do
  destroyFunc <- wrapWLResourceDestroyFunc destroyResource
  interfaces  <- newInterfaceCache
  pure $ ObjectCache { _cacheTypeInfo    = M.empty
                     , _cacheDestroyFunc = destroyFunc
                     , _cacheInterfaces  = interfaces
                     }

-- | Destroy an object cache, freeing all acquired resources.
--
-- This should only be done after no wl_resource instances are alive as it will free memory used by them.
destroyObjectCache :: ObjectCache -> IO ()
destroyObjectCache c = do
  freeHaskellFunPtr (_cacheDestroyFunc c)
  mapM_ (freeHaskellFunPtr . typeDispatchFunc) (_cacheTypeInfo c)
  destroyInterfaceCache (_cacheInterfaces c)

-- | Get the underlying Wayland resource wrapped by a typed object.
--
-- This function may raise an error if the underlying resource is no longer alive.
unwrapResource :: Object o => o -> IO WLResource
unwrapResource o = resolveHandle (objHandle (toAny o))

-- | Get the typed object wrapping a Wayland resource.
--
-- The object wrapper is stored as a StablePtr in the resource's userdata.
wrapResource :: Object o => WLResource -> IO o
wrapResource o = do
  objPtr <- castPtrToStablePtr <$> wlResourceGetUserData o
  deRefStablePtr objPtr

-- | Handle the destruction of a Wayland resource.
--
-- This will set the object handle's reference to a nullPtr and call the destroy hook.
destroyResource :: WLResource -> IO ()
destroyResource res = do
  objPtr <- castPtrToStablePtr <$> wlResourceGetUserData res
  obj <- deRefStablePtr objPtr
  invalidateHandle (objHandle obj)
  freeStablePtr objPtr
  objDestroyHook obj obj

-- | Create a new Wayland resource and associate it with a typed object wrapper.
--
-- This not only creates the object but also attaches a dispatch function to it and updates the destroy callback.
createResource
  :: forall o destroyMsg. Object o
  => MessageClass destroyMsg
  => (Message -> IO ()) -- ^ Message handler used for dispatching object requests and the destroy hook.
  -> NewObject o        -- ^ Handle describing the object to be created
  -> (o -> destroyMsg)  -- ^ Function creating a message dispatched after the object was destroyed.
  -> ObjectCacheState o -- ^ State monad returning the newly created object.
createResource handler newObj destroyHook = do
  existing <- liftIO $ readIORef (newObject newObj)
  case existing of
    Just _  -> error "The NewObject instance was already used to create an object."
    Nothing -> do typeInfo <- resolveTypeInfo (Proxy :: Proxy o) handler
                  destroyFunc <- use cacheDestroyFunc
                  liftIO $ do
                    let version = Wayland.Server.Interface.ifVersion (describeInterface (Proxy :: Proxy o))
                    client <- unwrapClient (newClient newObj)
                    res <- wlResourceCreate client (typeInterface typeInfo) (fromIntegral version) (newID newObj)
                    handle <- createHandle res
                    let obj = fromAny $ AnyObject { objHandle      = handle
                                                  , objDestroyHook = handler . Message . destroyHook . fromAny
                                                  }
                    objPtr <- newStablePtr obj
                    wlResourceSetDispatcher res (typeDispatchFunc typeInfo) nullPtr (castStablePtrToPtr objPtr) destroyFunc
                    writeIORef (newObject newObj) (Just obj)
                    pure obj

-- | Send an event for a typed object.
--
-- This function may raise an error if the underlying resource is no longer alive.
sendResourceEvent :: forall o. Object o => o -> ObjectEvent o -> IO ()
sendResourceEvent o e = do
  let paramCount = length $ msgArguments (eventSignature (Proxy :: Proxy o) e)
  res <- unwrapResource o
  allocaArray0 paramCount $ \buffer ->
    runContT (runStateT (runArgumentWriter (marshalEvent (Proxy :: Proxy o) e)) (ArgumentIterator buffer)) $ \(opcode, _) ->
      wlResourcePostEventArray res (fromIntegral opcode) buffer

-- | Request handler for a typed object.
--
-- This provides a WLDispatcherFunc for an object type.
dispatchResourceRequest :: forall o. Object o => Proxy o -> (Message -> IO ()) -> Ptr () -> WLResource -> Word32 -> Ptr WLMessage -> Ptr WLArgument -> IO CInt
dispatchResourceRequest _ handler _ res opcode _ args = do
  obj :: o <- wrapResource res
  client <- wlResourceGetClient res >>= wrapClient
  (req, _) <- runReaderT (runStateT (runArgumentReader (demarshalRequest obj (fromIntegral opcode))) (ArgumentIterator args)) client
  handler $ Message req
  pure 0

-- | Resolve a type information for an object type.
--
-- This will first attempt to use an existing instance from the 'ObjectCache'. If there is none, it will create a new
-- structure storing the wl_interface pointer and allocating a function pointer to the typed dispatcher function.
resolveTypeInfo
  :: Object o
  => Proxy o                   -- ^ Proxy for the object type to resolve.
  -> (Message -> IO ())        -- ^ Message handler used to dispatch requests in the type's dispatch func.
  -> ObjectCacheState TypeInfo -- ^ State monad wrapping the type information.
resolveTypeInfo proxy handler = do
  existing <- use cacheTypeInfo
  case M.lookup (typeRep proxy) existing of
    Just info -> pure info
    Nothing   -> do let desc = describeInterface proxy
                    interface    <- zoom cacheInterfaces $ resolveInterface desc
                    dispatchFunc <- liftIO $ wrapWLDispatcherFunc (dispatchResourceRequest proxy handler)
                    let info = TypeInfo { typeInterface    = interface
                                        , typeDispatchFunc = dispatchFunc
                                        }
                    cacheTypeInfo %= M.insert (typeRep proxy) info
                    pure info

-- | Resolve the native interface for an object type.
resolveTypeInterface
  :: Object o
  => Proxy o                            -- ^ Proxy for the object type to resolve.
  -> (Message -> IO ())                 -- ^ Message handler used to dispatch requests in the type's dispatch func.
  -> ObjectCacheState (Ptr WLInterface) -- ^ State monad wrapping the native interface.
resolveTypeInterface proxy handler = typeInterface <$> resolveTypeInfo proxy handler

-- | Get the next argument in an argument iterator.
--
-- This function may raise an error if there is no more argument remaining.
advanceIterator :: MonadState ArgumentIterator m => m (Ptr WLArgument)
advanceIterator = do
  ArgumentIterator arg <- get
  put $ ArgumentIterator $ plusPtr arg (sizeOf (undefined :: WLArgument))
  pure arg

-- | Create an 'ArgumentWriter' wrapping a continuation (like alloca).
liftCont :: ((a -> IO r) -> IO r) -> ArgumentWriter r a
liftCont = ArgumentWriter . lift . ContT

instance ArgumentMarshaller (ArgumentWriter r) where
  marshalInt v = do
    arg <- advanceIterator
    liftIO $ wlArgumentPokeInt arg (fromIntegral v)
  marshalUInt v = do
    arg <- advanceIterator
    liftIO $ wlArgumentPokeUInt arg (fromIntegral v)
  marshalFixed v = do
    arg <- advanceIterator
    liftIO $ wlArgumentPokeFixed arg (WLFixed (round (v * 256)))
  marshalFD (Fd v) = do
    arg <- advanceIterator
    liftIO $ wlArgumentPokeFD arg (fromIntegral v)
  marshalString v = do
    arg <- advanceIterator
    str <- liftCont $ withCString v
    liftIO $ wlArgumentPokeString arg str
  marshalObject v = do
    arg <- advanceIterator
    obj <- liftIO $ unwrapResource v
    liftIO $ wlArgumentPokeObject arg obj
  marshalNewID v = do
    arg <- advanceIterator
    liftIO $ wlArgumentPokeNewID arg (newID v)
  marshalArray v = do
    arg <- advanceIterator
    let len = fromIntegral $ length v * sizeOf (head v)
    buf <- liftCont $ withArray v
    arr <- liftCont $ with $ WLArray { arrSize  = len
                                     , arrAlloc = len
                                     , arrData  = castPtr buf
                                     }
    liftIO $ wlArgumentPokeArray arg arr

instance ArgumentDemarshaller ArgumentReader where
  demarshalInt    = fromIntegral   <$> (advanceIterator >>= liftIO . wlArgumentPeekInt)
  demarshalUInt   = fromIntegral   <$> (advanceIterator >>= liftIO . wlArgumentPeekUInt)
  demarshalFixed  = wlFixedToFloat <$> (advanceIterator >>= liftIO . wlArgumentPeekFixed)
  demarshalString = do
    ptr <- advanceIterator >>= liftIO . wlArgumentPeekString
    liftIO $ peekCString ptr
  demarshalObject = do
    ptr <- advanceIterator >>= liftIO . wlArgumentPeekObject
    liftIO $ wrapResource ptr
  demarshalNewID = do
    objID <- advanceIterator >>= liftIO . wlArgumentPeekNewID
    client <- ask
    liftIO $ createNewObject client objID
  demarshalArray = do
    ptr <- advanceIterator >>= liftIO . wlArgumentPeekArray
    arr <- liftIO $ peek ptr
    liftIO $ peekArray (fromIntegral (arrSize arr)) (castPtr (arrData arr))
  demarshalFD = Fd . fromIntegral <$> (advanceIterator >>= liftIO . wlArgumentPeekFD)