{-# LANGUAGE OverloadedStrings #-}

-- Server.hs - UNIX socket server for RDE remote control
--
-- Note: This module requires the 'network' and 'stm' packages to be installed.
-- Add them to package.yaml:
--   dependencies:
--     - network
--     - stm

module Rahm.Desktop.Server where

import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception (finally, throwIO, try)
import Control.Monad (forever, void, when)
import Control.Monad.Fix
import Control.Monad.IO.Class (liftIO)
import Control.Monad.RWS (All (All))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC
import Data.Maybe (fromMaybe)
import Graphics.X11.Xlib
import Graphics.X11.Xlib.Extras (Event (..), currentTime, setClientMessageEvent, setEventType)
import Network.Socket
import Network.Socket.ByteString (recv, sendAll)
import System.Directory
import System.Environment (lookupEnv)
import System.FilePath ((</>))
import System.IO
import System.IO.Error (isEOFError)
import XMonad (X, XConfig (handleEventHook), io)

data CommandHandler c where
  CommandHandler ::
    { parser :: String -> Maybe c,
      handler :: c -> X ()
    } ->
    CommandHandler c

(<<>>) :: CommandHandler c1 -> CommandHandler c2 -> CommandHandler (Either c1 c2)
(<<>>) c1 c2 =
  CommandHandler
    ( \str -> case (parser c1 str, parser c2 str) of
        (Just l, _) -> Just (Left l)
        (_, Just r) -> Just (Right r)
        _ -> Nothing
    )
    ( \case
        Right r -> handler c2 r
        Left l -> handler c1 l
    )

data ServerContext c where
  ServerContext ::
    { display :: Display,
      rootWin :: Window,
      eventType :: Atom,
      commandHandler :: CommandHandler c,
      channel :: TChan c
    } ->
    ServerContext c

enableServerMode :: CommandHandler c -> XConfig l -> IO (XConfig l)
enableServerMode handler xconfig = do
  serverCtx <- forkServer handler
  return $
    xconfig
      { handleEventHook =
          eventHandler serverCtx <> handleEventHook xconfig
      }

socketPath :: IO FilePath
socketPath = do
  dir <- fromMaybe "/tmp/rde" <$> lookupEnv "XDG_RUNTIME_DIR"
  display <- fromMaybe ":0" <$> lookupEnv "DISPLAY"
  createDirectoryIfMissing True dir
  return (dir </> "rde-" <> dropWhile (== ':') display <> ".sock")

forkServer :: CommandHandler c -> IO (ServerContext c)
forkServer handler = do
  chan <- atomically newTChan
  path <- socketPath

  -- Remove stale socket if it exists
  exists <- doesPathExist path
  when exists $ removeFile path

  sock <- socket AF_UNIX Stream defaultProtocol
  bind sock (SockAddrUnix path)
  listen sock 5

  putStrLn $ "Listening on " ++ path

  dpy <- getDefaultDisplay
  let screen = defaultScreen dpy
  win <- rootWindow dpy screen
  eventType <- internAtom dpy "RDE_COMMAND_NOTIFY" True

  let ctx = ServerContext dpy win eventType handler chan

  forkIO $ do
    hFlush stdout
    forever $ do
      (conn, _) <- accept sock
      hFlush stdout
      void $ forkIO $ handleClient ctx conn

  return ctx

handleClient :: ServerContext c -> Socket -> IO ()
handleClient chan conn =
  finally (loop BS.empty) (close conn)
  where
    loop buf = do
      chunk <- recv conn 4096
      if BS.null chunk
        then return () -- connection closed
        else do
          let buf' = buf <> chunk
              (lines, rest) = splitLines buf'
          mapM_ (handleLine chan . BSC.unpack) lines
          loop rest

    splitLines :: BS.ByteString -> ([BS.ByteString], BS.ByteString)
    splitLines bs =
      case BSC.break (== '\n') bs of
        (line, rest)
          | BS.null rest -> ([], bs)
          | otherwise ->
              let remaining = BS.drop 1 rest
                  (ls, r) = splitLines remaining
               in (line : ls, r)

handleLine :: ServerContext c -> String -> IO ()
handleLine ctx line = case parser (commandHandler ctx) line of
  Just s -> pushChan s
  _ -> hPutStrLn stderr $ "Unable to parse command `" ++ line ++ "`"
  where
    pushChan c = do
      atomically $ writeTChan (channel ctx) c
      raiseNotifyEvent ctx

raiseNotifyEvent :: ServerContext c -> IO ()
raiseNotifyEvent ServerContext {display = dpy, rootWin = root, eventType = ty} =
  allocaXEvent $ \evPtr -> do
    m <- internAtom dpy "" False
    setEventType evPtr clientMessage
    setClientMessageEvent evPtr root ty 32 m currentTime
    sendEvent dpy root False structureNotifyMask evPtr
    flush dpy

eventHandler :: ServerContext c -> Event -> X All
eventHandler
  ServerContext
    { display = dpy,
      rootWin = root,
      eventType = ty,
      channel = chan,
      commandHandler = commandHandler
    }
  e@ClientMessageEvent {ev_message_type = atom} = do
    if atom == ty
      then do
        mapM_ (handler commandHandler) =<< io (atomically (tryReadTChan chan))
        return (All True)
      else return (All True)
eventHandler _ _ = return (All True)

getDefaultDisplay :: IO Display
getDefaultDisplay = openDisplay ""