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

module Wayland.Server
  ( Request(..)
  , ObjectID
  , Global
  , ArgumentMarshaller(..)
  , ArgumentDemarshaller(..)
  , DisplayM(..)
  , ResourceClass(..)
  , EventClass
  , RequestClass
  , InterfaceDesc(..)
  , MessageDesc(..)
  , ArgumentDesc(..)
  , Nullability(..)
  , fromRequest
  , withDisplay
  , bindSocket
  , processRequests
  , flushEvents
  , sendEvent
  , destroyResource
  , registerGlobal
  , unregisterGlobal
  , registerSignalEvent
  , requestTermination
  ) where

import Data.Typeable (Typeable, Proxy(..), cast, TypeRep, typeRep)
import Data.Kind (Type)
import Control.Monad.State.Strict (StateT, MonadState(..), MonadIO(..), runStateT)
import Lens.Micro.Mtl.Internal (Zoom(..), Zoomed)
import Lens.Micro.TH (makeLenses)
import Lens.Micro.Mtl ((%=), use, (<<%=))
import qualified Wayland.Server.System as Sys
import Control.Monad.Trans (MonadTrans(..))
import Foreign.C (peekCString, newCString, withCString)
import Foreign.C.Types (CInt)
import Foreign (Ptr, nullPtr, new, newStablePtr, castStablePtrToPtr, castPtrToStablePtr, deRefStablePtr, freeStablePtr, Storable(..), Int32, castPtr, withArray, with, allocaArray0)
import Data.IORef (IORef, newIORef, writeIORef, readIORef)
import Control.Monad.Reader (ReaderT, MonadReader(..), runReaderT)
import Control.Exception (bracket)
import Data.Functor (($>))
import Control.Monad (when)
import Data.Maybe (isJust)
import qualified Data.Map as M
import Foreign.Marshal (newArray)
import Data.Word (Word32)
import System.Mem.StableName (StableName, makeStableName)
import Data.Hashable (Hashable(..))
import Foreign.Ptr (plusPtr)
import Control.Monad.Cont (ContT(..))

data Global = Global

-- IORef set to nullPtr once destroyed.
data ObjectID = ObjectID (StableName (IORef Sys.WLResource)) (IORef Sys.WLResource)
  deriving (Eq)

instance Ord ObjectID where
  (ObjectID nameA _) <= (ObjectID nameB _) = hash nameA <= hash nameB

instance Hashable ObjectID where
  hashWithSalt s (ObjectID name _) = hashWithSalt s name

instance Show ObjectID where
  show (ObjectID name _) = "ObjectID " ++ show (hash name)

data Nullability
  = AllowNull
  | NotNull
  deriving (Show, Eq, Ord)

data ArgumentDesc
  = ArgumentInt
  | ArgumentUInt
  | ArgumentFixed
  | ArgumentString !Nullability
  | ArgumentObject !Nullability !InterfaceDesc
  | ArgumentNewID  !InterfaceDesc
  | ArgumentArray
  | ArgumentFD
  deriving (Show, Eq, Ord)

data MessageDesc
  = MessageDesc { msgName              :: !String
                , msgSignature         :: ![ArgumentDesc]
                , msgSince             :: !Int
                , msgIsDestructor      :: !Bool
                }
  deriving (Show, Eq, Ord)

data InterfaceDesc
  = InterfaceDesc { ifName     :: !String
                  , ifVersion  :: !Int
                  , ifRequests :: ![MessageDesc]
                  , ifEvents   :: ![MessageDesc]
                  }
  deriving (Show, Eq, Ord)

class (Typeable r, Show r) => RequestClass r where

class (Typeable e, Show e) => EventClass e where

data Request = forall a. RequestClass a => Request a

newtype DisplayM s a = DisplayM { runDisplayM :: StateT s (StateT DisplayState (ReaderT Display IO)) a }
  deriving (Functor, Applicative, Monad, MonadIO, MonadState s, MonadFail)

type RequestHandler s = Request -> DisplayM s ()

newtype EventLoop = EventLoop { handleEvent :: Request -> IO () }

data Display
  = Display { display         :: !Sys.WLDisplay
            , activeEventLoop :: !(IORef (Maybe EventLoop))
            }

data DisplayState
  = DisplayState { _typeInterfaces   :: !(M.Map TypeRep       (Ptr Sys.WLInterface))
                 , _nativeInterfaces :: !(M.Map InterfaceDesc (Ptr Sys.WLInterface))
                 }

makeLenses ''DisplayState

class (Typeable r, RequestClass (ResourceRequest r), EventClass (ResourceEvent r)) => ResourceClass r where
  type ResourceRequest r :: Type
  type ResourceEvent   r :: Type
  objectID  :: r        -> ObjectID
  fromID    :: ObjectID -> r
  interface :: Proxy r  -> InterfaceDesc
  marshalEvent     :: ArgumentMarshaller m   => Proxy r -> ResourceEvent r -> m Int
  demarshalRequest :: ArgumentDemarshaller m => r -> Int -> m (ResourceRequest r)
  eventSignature   :: Proxy r -> ResourceEvent r   -> MessageDesc
  requestSignature :: Proxy r -> ResourceRequest r -> MessageDesc

class (Monad m, MonadFail m) => ArgumentMarshaller m where
  marshalInt    :: Int                  -> m ()
  marshalUInt   :: Word                 -> m ()
  marshalFixed  :: Float                -> m ()
  marshalString :: String               -> m ()
  marshalObject :: ResourceClass r => r -> m ()
  marshalNewID  :: ResourceClass r => r -> m ()
  marshalArray  :: Storable a => [a]    -> m ()
  marshalFD     :: Int32                -> m ()

class (Monad m, MonadFail m) => ArgumentDemarshaller m where
  demarshalInt    :: m Int
  demarshalUInt   :: m Word
  demarshalFixed  :: m Float
  demarshalString :: m String
  demarshalObject :: ResourceClass r => m r
  demarshalNewID  :: ResourceClass r => m r
  demarshalArray  :: Storable a => m [a]
  demarshalFD     :: m Int32

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

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

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

getDisplay :: DisplayM s Display
getDisplay = DisplayM $ lift $ lift ask

liftDisplayState :: StateT DisplayState (ReaderT Display IO) a -> DisplayM s a
liftDisplayState m = DisplayM $ lift m

fromRequest :: Typeable a => Request -> Maybe a
fromRequest (Request x) = cast x

withDisplay :: s -> DisplayM s a -> IO (a, s)
withDisplay initialState f = bracket Sys.wlDisplayCreate Sys.wlDisplayDestroy $ \sysDisp -> do
  eventLoop <- newIORef Nothing
  let disp = Display { display         = sysDisp
                     , activeEventLoop = eventLoop
                     }
      dispState = DisplayState { _typeInterfaces   = M.empty
                               , _nativeInterfaces = M.empty
                               }
  (res, _) <- runReaderT (runStateT (runStateT (runDisplayM f) initialState) dispState) disp
  pure res

bindSocket :: DisplayM s String
bindSocket = do
  disp <- getDisplay
  name <- liftIO $ Sys.wlDisplayAddSocketAuto (display disp)
  liftIO $ peekCString name

processRequests :: RequestHandler s -> DisplayM s ()
processRequests h = do
  disp <- getDisplay
  stateBefore <- get
  displayStateBefore <- liftDisplayState get
  stateRef <- liftIO $ newIORef (stateBefore, displayStateBefore)
  let handler :: Request -> IO ()
      handler r = do
        (s, ds) <- readIORef stateRef
        ((_, s'), ds') <- runReaderT (runStateT (runStateT (runDisplayM (h r)) s) ds) disp
        writeIORef stateRef (s', ds')
  loop <- liftIO $ readIORef (activeEventLoop disp)
  when (isJust loop) (fail "Cannot start an event loop from within an event handler.")
  liftIO $ writeIORef (activeEventLoop disp) (Just (EventLoop handler))
  liftIO $ Sys.wlDisplayRun (display disp)
  (stateAfter, displayStateAfter) <- liftIO $ readIORef stateRef
  put stateAfter
  liftDisplayState $ put displayStateAfter

registerGlobal :: forall r bindReq unbindReq s. ResourceClass r => RequestClass bindReq => RequestClass unbindReq => (r -> bindReq) -> (r -> unbindReq) -> Int -> DisplayM s Global
registerGlobal mkBindReq mkUnbindReq version = do
  disp <- getDisplay
  iface <- getInterface (Proxy :: Proxy r)
  let dispatcher :: Ptr () -> Sys.WLResource -> Word32 -> Ptr Sys.WLMessage -> Ptr Sys.WLArgument -> IO CInt
      dispatcher _ res opcode _ _ = do
        handlePtr <- castPtrToStablePtr <$> Sys.wlResourceGetUserData res
        objID@(ObjectID _ _) <- deRefStablePtr handlePtr
        putStrLn $ "Handling request " ++ show opcode ++ " for " ++ show objID
        -- TODO: Handle request
        pure 0
  let destroy :: Sys.WLResource -> IO ()
      destroy res = do
        handlePtr <- castPtrToStablePtr <$> Sys.wlResourceGetUserData res
        objID@(ObjectID _ handle) <- deRefStablePtr handlePtr
        writeIORef handle (Sys.WLResource nullPtr)
        freeStablePtr handlePtr
        let req = mkUnbindReq (fromID objID)
        runRequest disp (Request req)
        pure ()
  let createObject :: Sys.WLClient -> Ptr () -> Word32 -> Word32 -> IO ()
      createObject client _ ver protID = do
        res <- Sys.wlResourceCreate client iface (fromIntegral ver) protID
        handle <- newIORef res
        name <- makeStableName handle
        let objID = ObjectID name handle
        handlePtr <- newStablePtr objID
        dis <- Sys.wrapWLDispatcherFunc dispatcher
        des <- Sys.wrapWLResourceDestroyFunc destroy
        -- TODO: Free functions
        Sys.wlResourceSetDispatcher res dis nullPtr (castStablePtrToPtr handlePtr) des
        let req = mkBindReq (fromID objID)
        runRequest disp (Request req)
        pure ()
  liftIO $ do
    func <- Sys.wrapWLGlobalBindFunc createObject
    _ <- Sys.wlGlobalCreate (display disp) iface (fromIntegral version) nullPtr func
    -- TODO: Free func
    -- TODO: Return proper Global
    pure ()
  pure Global

unregisterGlobal :: Global -> DisplayM s ()
unregisterGlobal _ = pure ()

runRequest :: Display -> Request -> IO ()
runRequest disp r = do
  loop <- readIORef (activeEventLoop disp)
  case loop of
    Just handler -> handleEvent handler r
    Nothing      -> putStrLn $ "Lost request because no event loop is active: " ++ show r

registerSignalEvent :: RequestClass req => req -> CInt -> DisplayM s ()
registerSignalEvent r sig = do
  disp <- getDisplay
  liftIO $ do
    loop <- Sys.wlDisplayGetEventLoop (display disp)
    func <- Sys.wrapWLEventLoopSignalFunc $ \_ _ -> runRequest disp (Request r) $> 0
    _ <- Sys.wlEventLoopAddSignal loop sig func nullPtr
    -- TODO: Free func
    -- TODO: Return event source handle
    pure ()

requestTermination :: DisplayM s ()
requestTermination = do
  disp <- getDisplay
  liftIO $ Sys.wlDisplayTerminate (display disp)

wrapInterface :: InterfaceDesc -> DisplayM s (Ptr Sys.WLInterface)
wrapInterface i = do
  let mkArgumentType :: ArgumentDesc -> DisplayM s (Ptr Sys.WLInterface)
      mkArgumentType (ArgumentObject _ a) = wrapInterface a
      mkArgumentType (ArgumentNewID a)    = wrapInterface a
      mkArgumentType _                    = pure nullPtr
  let mkArgument :: ArgumentDesc -> String
      mkArgument ArgumentInt                  = "i"
      mkArgument ArgumentUInt                 = "u"
      mkArgument ArgumentFixed                = "f"
      mkArgument (ArgumentString NotNull)     = "s"
      mkArgument (ArgumentString AllowNull)   = "?s"
      mkArgument (ArgumentObject NotNull _)   = "o"
      mkArgument (ArgumentObject AllowNull _) = "?o"
      mkArgument (ArgumentNewID _)            = "n"
      mkArgument ArgumentArray                = "a"
      mkArgument ArgumentFD                   = "h"
  let mkMessage :: MessageDesc -> DisplayM s Sys.WLMessage
      mkMessage m = do
        name      <- liftIO $ newCString (msgName m)
        signature <- liftIO $ newCString (show (msgSince m) ++ (msgSignature m >>= mkArgument))
        types     <- mapM mkArgumentType (msgSignature m) >>= liftIO . newArray
        pure $ Sys.WLMessage { Sys.msgName      = name
                             , Sys.msgSignature = signature
                             , Sys.msgTypes     = types
                             }
  let mkInterface :: DisplayM s (Ptr Sys.WLInterface)
      mkInterface = do
        name    <- liftIO $ newCString (ifName i)
        methods <- mapM mkMessage (ifRequests i) >>= liftIO . newArray
        events  <- mapM mkMessage (ifEvents i)   >>= liftIO . newArray
        liftIO $ new $ Sys.WLInterface { Sys.ifName        = name
                                       , Sys.ifVersion     = fromIntegral $ ifVersion i
                                       , Sys.ifMethodCount = fromIntegral $ length (ifRequests i)
                                       , Sys.ifMethods     = methods
                                       , Sys.ifEventCount  = fromIntegral $ length (ifEvents i)
                                       , Sys.ifEvents      = events
                                       }
  existing <- liftDisplayState $ use nativeInterfaces
  case M.lookup i existing of
    Just ptr -> pure ptr
    Nothing  -> do ptr <- mkInterface
                   -- TODO: Free ptr
                   liftDisplayState $ nativeInterfaces %= M.insert i ptr
                   pure ptr

getInterface :: ResourceClass r => Proxy r -> DisplayM s (Ptr Sys.WLInterface)
getInterface p = do
  existing <- liftDisplayState $ use typeInterfaces
  case M.lookup (typeRep p) existing of
    Just ptr -> pure ptr
    Nothing  -> do ptr <- wrapInterface (interface p)
                   liftDisplayState $ typeInterfaces %= M.insert (typeRep p) ptr
                   pure ptr

data ArgumentIteratorState
  = ArgumentIteratorState { _argNext          :: !(Ptr Sys.WLArgument)
                          , _argRemaining     :: !Int
                          }
  deriving (Show)

makeLenses ''ArgumentIteratorState

newtype ArgumentIterator r a = ArgumentIterator { runArgumentIterator :: StateT ArgumentIteratorState (ContT r IO) a }
  deriving (Functor, Applicative, Monad, MonadFail, MonadState ArgumentIteratorState, MonadIO)

advanceIterator :: ArgumentIterator r (Ptr Sys.WLArgument)
advanceIterator = do
  next      <- argNext      <<%= \p -> plusPtr p (sizeOf (undefined :: Sys.WLArgument))
  remaining <- argRemaining <<%= (+(-1))
  when (remaining == 0) (fail "Out of range for argument")
  pure next

getObjectForHandle :: ResourceClass r => r -> IO (Ptr a)
getObjectForHandle r = fetch (objectID r)
  where fetch (ObjectID _ handle) = readIORef handle >>= check
        check (Sys.WLResource ptr) | ptr == nullPtr = fail "Invalid object handle."
                                   | otherwise      = pure (castPtr ptr)

liftCont :: ((a -> IO r) -> IO r) -> ArgumentIterator r a
liftCont = ArgumentIterator . lift . ContT

instance ArgumentMarshaller (ArgumentIterator r) where
  marshalInt    v = advanceIterator >>= (liftIO . Sys.wlArgumentSetInt   (fromIntegral v))
  marshalUInt   v = advanceIterator >>= (liftIO . Sys.wlArgumentSetUInt  (fromIntegral v))
  marshalFixed  v = advanceIterator >>= (liftIO . Sys.wlArgumentSetFixed (Sys.WLFixed (round (v * 256))))
  marshalFD     v = advanceIterator >>= (liftIO . Sys.wlArgumentSetFD    (fromIntegral v))
  marshalString v = do
    arg <- advanceIterator
    str <- liftCont $ withCString v
    liftIO $ Sys.wlArgumentSetString str arg
  marshalObject v = do
    arg <- advanceIterator
    obj <- liftIO $ getObjectForHandle v
    liftIO $ Sys.wlArgumentSetObject obj arg
  marshalNewID v = do
    arg <- advanceIterator
    obj <- liftIO $ getObjectForHandle v
    liftIO $ Sys.wlArgumentSetObject obj arg
  marshalArray v = do
    arg <- advanceIterator
    let len = fromIntegral $ length v * sizeOf (head v)
    buf <- liftCont $ withArray v
    arr <- liftCont $ with $ Sys.WLArray { Sys.arrSize  = len
                                         , Sys.arrAlloc = len
                                         , Sys.arrData  = castPtr buf
                                         }
    liftIO $ Sys.wlArgumentSetArray arr arg

sendEvent :: forall r s. ResourceClass r => r -> ResourceEvent r -> DisplayM s ()
sendEvent res e = liftIO $ do
  let params = msgSignature (eventSignature (Proxy :: Proxy r) e)
  ptr <- getObjectForHandle res
  allocaArray0 (length params) $ \buffer ->
    let s = ArgumentIteratorState { _argNext          = buffer
                                  , _argRemaining     = length params
                                  }
        handler (opcode, _) = Sys.wlResourcePostEventArray (Sys.WLResource ptr) (fromIntegral opcode) buffer
    in runContT (runStateT (runArgumentIterator (marshalEvent (Proxy :: Proxy r) e)) s) handler

flushEvents :: DisplayM s ()
flushEvents = do
  disp <- getDisplay
  liftIO $ Sys.wlDisplayFlushClients (display disp)

destroyResource :: ResourceClass r => r -> DisplayM s ()
destroyResource r = liftIO $ do
  ptr <- getObjectForHandle r
  Sys.wlResourceDestroy (Sys.WLResource ptr)