{-# 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
  , Request(..)
  , fromRequest
  , ArgumentMarshaller(..)
  , ArgumentDemarshaller(..)
  , -- * State
    ObjectCache
  , ObjectCacheState
  , newObjectCache
  , destroyObjectCache
  , resolveTypeInterface
  , -- * Lifecycle
    unwrapResource
  , wrapResource
  , createResource
  , sendResourceEvent
  ) where

import System.Mem.StableName (makeStableName, StableName)
import Data.Kind (Type)
import Data.Hashable (Hashable(..))
import Data.IORef (readIORef, newIORef, writeIORef, IORef)
import Data.Typeable (Proxy(..), TypeRep, typeRep, Typeable, cast)
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 (when)
import Wayland.Server.System
import Wayland.Server.Interface (MessageDesc(..), InterfaceDesc, InterfaceCache, newInterfaceCache, destroyInterfaceCache, resolveInterface)
import qualified Data.Map as M
import Wayland.Server.Client (Client, unwrapClient)

-- | 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.
--
-- The request's type is erased and it is then processed by the display's event loop.
class (Typeable r, Show 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 (Typeable e, Show e) => EventClass e where

-- | A type-erased request.
data Request = forall r. RequestClass r => Request r

-- | 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     :: Int32                   -> 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 Int32

-- | 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. This is realized using an IORef which is set to nullPtr once the resource is destroyed. The IORef itself
-- is used for equality checks.
--
-- In addition, a custom type-erased cleanup action is stored which is invoked once the resource is destroyed.
data AnyObject
  = AnyObject { objName        :: StableName (IORef WLResource) -- ^ Name used for hashing and ordering.
              , objResource    :: IORef WLResource              -- ^ Reference to 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.
data NewObject o
  = NewObject { newClient  :: !Client            -- ^ The client this object is associated with.
              , newID      :: !Word32            -- ^ Object ID of the new object.
              , newVersion :: !CInt              -- ^ Protocol version to use for the newly created 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 fixed number of wl_arguments.
data ArgumentIterator
  = ArgumentIterator { _argNext          :: !(Ptr WLArgument) -- ^ Pointer to the next argument to be processed.
                     , _argRemaining     :: !Int              -- ^ Number of remaining arguments in the list.
                     }

-- | A state monad providing an instance of 'ArgumentIterator'.
--
-- The monad wraps a ContT monad to allow using alloca when wrapping arguments.
newtype ArgumentIteratorState r a
  = ArgumentIteratorState { runArgumentIteratorState :: StateT ArgumentIterator (ContT r IO) a }
  deriving (Functor, Applicative, Monad, MonadState ArgumentIterator, MonadIO)

makeLenses ''ObjectCache
makeLenses ''ArgumentIterator

instance Show Request where
  show (Request r) = show r

-- | Try to cast a type-erased request into a concrete 'RequestClass'.
fromRequest :: RequestClass a => Request -> Maybe a
fromRequest (Request x) = cast x

instance Eq AnyObject where
  a == b = objResource a == objResource b

instance Ord AnyObject where
  a <= b = hash (objName a) <= hash (objName b)

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

instance Show AnyObject where
  show o = "AnyObject " ++ show (hash (objName 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 -> CInt -> IO (NewObject o)
createNewObject client objID version = do
  obj <- newIORef Nothing
  pure $ NewObject { newClient  = client
                   , newID      = objID
                   , newVersion = version
                   , 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 = fetch (toAny o)
  where fetch obj = readIORef (objResource obj) >>= check
        check res@(WLResource ptr) | ptr == nullPtr = error "Invalid object handle."
                                   | otherwise      = pure res

-- | 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
  writeIORef (objResource obj) (WLResource nullPtr)
  freeStablePtr objPtr
  -- TODO: Prevent destroy hook from being invoked while in the request handler.
  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 destroyReq. Object o
  => RequestClass destroyReq
  => (Request -> IO ()) -- ^ Request handler used for dispatching object requests and the destroy hook.
  -> NewObject o        -- ^ Handle describing the object to be created
  -> (o -> destroyReq)  -- ^ Function creating a request dispatched after the object was destroyed.
  -> ObjectCacheState o -- ^ State monad returning the newly created object.
createResource handle newObj destroyHook = do
  existing <- liftIO $ readIORef (newObject newObj)
  case existing of
    Just obj -> pure obj
    Nothing  -> do typeInfo <- resolveTypeInfo (Proxy :: Proxy o)
                   destroyFunc <- use cacheDestroyFunc
                   liftIO $ do
                     client <- unwrapClient (newClient newObj)
                     res <- wlResourceCreate client (typeInterface typeInfo) (newVersion newObj) (newID newObj)
                     resRef <- newIORef res
                     name <- makeStableName resRef
                     let obj = fromAny $ AnyObject { objName        = name
                                                   , objResource    = resRef
                                                   , objDestroyHook = handle . Request . destroyHook . fromAny
                                                   }
                     objPtr <- newStablePtr obj
                     wlResourceSetDispatcher res (typeDispatchFunc typeInfo) nullPtr (castStablePtrToPtr objPtr) destroyFunc
                     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 ->
    let s = ArgumentIterator { _argNext      = buffer
                             , _argRemaining = paramCount
                             }
        handler (opcode, _) = wlResourcePostEventArray res (fromIntegral opcode) buffer
    in runContT (runStateT (runArgumentIteratorState (marshalEvent (Proxy :: Proxy o) e)) s) handler

-- | Request handler for a typed object.
--
-- This provides a WLDispatcherFunc for an object type.
dispatchResourceRequest :: forall o. Object o => Proxy o -> Ptr () -> WLResource -> Word32 -> Ptr WLMessage -> Ptr WLArgument -> IO CInt
dispatchResourceRequest _ _ res _ _ _ = do
  _ :: o <- wrapResource res
  -- TODO: Handle request
  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 -> ObjectCacheState TypeInfo
resolveTypeInfo proxy = do
  existing <- use cacheTypeInfo
  case M.lookup (typeRep proxy) existing of
    Just info -> pure info
    Nothing   -> do interface    <- zoom cacheInterfaces $ resolveInterface (describeInterface proxy)
                    dispatchFunc <- liftIO $ wrapWLDispatcherFunc (dispatchResourceRequest proxy)
                    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 -> ObjectCacheState (Ptr WLInterface)
resolveTypeInterface proxy = typeInterface <$> resolveTypeInfo proxy

-- | Get the next argument in an argument iterator.
--
-- This function may raise an error if there is no more argument remaining.
advanceIterator :: ArgumentIteratorState r (Ptr WLArgument)
advanceIterator = do
  next      <- argNext      <<%= \p -> plusPtr p (sizeOf (undefined :: WLArgument))
  remaining <- argRemaining <<%= (+(-1))
  when (remaining == 0) (error "Out of range for argument")
  pure next

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

instance ArgumentMarshaller (ArgumentIteratorState 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 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
    -- TODO
    undefined
  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 (ArgumentIteratorState r) 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
    -- TODO
    undefined
  demarshalArray = do
    ptr <- advanceIterator >>= liftIO . wlArgumentPeekArray
    arr <- liftIO $ peek ptr
    liftIO $ peekArray (fromIntegral (arrSize arr)) (castPtr (arrData arr))
  demarshalFD = advanceIterator >>= liftIO . wlArgumentPeekFD