{-# LANGUAGE DuplicateRecordFields #-}

module Wetterhorn.Core.W where

import Control.Arrow (Arrow (first))
import Control.Monad ((<=<))
import Control.Monad.RWS (MonadIO (liftIO), MonadReader (..), MonadState, modify)
import Control.Monad.Reader (ReaderT (runReaderT))
import Control.Monad.State (StateT (runStateT), gets, modify')
import Control.Monad.Trans.Class
import Control.Monad.Trans.Maybe
import Data.Data (TypeRep, Typeable, cast, tyConModule, tyConName, tyConPackage)
import Data.Default.Class (Default, def)
import Data.Kind (Constraint, Type)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Proxy
import Data.Set (Set)
import qualified Data.Set as Set
import Foreign (Ptr, StablePtr, intPtrToPtr, ptrToIntPtr)
import Text.Printf (printf)
import Text.Read hiding (lift)
import Type.Reflection (someTypeRep, someTypeRepTyCon)
import Wetterhorn.Core.ButtonEvent (ButtonEvent)
import Wetterhorn.Core.KeyEvent
import Wetterhorn.Core.SurfaceEvent
import Wetterhorn.Foreign.ForeignInterface (ForeignInterface)
import qualified Wetterhorn.Foreign.ForeignInterface as ForeignInterface
import Wetterhorn.Foreign.WlRoots (Surface, WlrSeat)
import Wetterhorn.StackSet hiding (layout)
import qualified Wetterhorn.StackSet as StackSet

data RationalRect = RationalRect Rational Rational Rational Rational

-- | Wrapper for a message. Messages are sent to layout and layouts are supposed
-- to handle them. This hides a typeable parameter.
data Message where
  Message :: (Typeable a) => a -> Message

-- | casts a message to a type.
fromMessage :: (Typeable a) => Message -> Maybe a
fromMessage (Message t) = cast t

-- | Wraps a type in a message.
toMessage :: (Typeable a) => a -> Message
toMessage = Message

class (Typeable l) => HandleMessage l where
  handleMessage :: Message -> l -> MaybeT W l
  handleMessage _ = return

newtype Window = Window
  { surface :: Surface
  }
  deriving (Show, Ord, Eq, Read)

-- | Types of this class "lay out" windows by assigning rectangles and handle
-- messages.
class (Typeable l, HandleMessage l) => LayoutClass l where
  -- | Constraints on the type to lay out. Sometimes a layout requires the 'a'
  -- type to be "Ord", other times "Eq", this is the mechanism by which this
  -- constraint is expressed.
  type LayoutConstraint l :: Type -> Constraint

  -- | Runs the layout in an impure way returning a modified layout and the list
  -- of windows to their rectangles under a monad.
  runLayout :: (LayoutConstraint l a) => Stack a -> l -> W (l, [(a, RationalRect)])

  readLayout :: String -> Maybe l
  default readLayout :: (Read l) => String -> Maybe l
  readLayout = readMaybe

  serializeLayout :: l -> String
  default serializeLayout :: (Show l) => l -> String
  serializeLayout = show

  description :: l -> String
  default description :: (Show l) => l -> String
  description = show
  {-# MINIMAL runLayout #-}

-- | Lifts a pure-layout implementation to a signature that complies with
-- 'runLayout'
pureLayout ::
  (Stack a -> l -> [(a, RationalRect)]) ->
  Stack a ->
  l ->
  W (l, [(a, RationalRect)])
pureLayout fn as l = return (l, fn as l)

-- A Layout which hides the layout parameter under an existential type and
-- asserts the layout hidden can work with Window types.
data WindowLayout
  = forall l a.
    (LayoutClass l, LayoutConstraint l a, a ~ Window) =>
    WindowLayout l

runWindowLayout :: Stack Window -> WindowLayout -> W (WindowLayout, [(Window, RationalRect)])
runWindowLayout as (WindowLayout l) = first WindowLayout <$> runLayout as l

handleWindowMessage :: Message -> WindowLayout -> MaybeT W WindowLayout
handleWindowMessage m (WindowLayout l) = WindowLayout <$> handleMessage m l

-- | Using the 'Layout' as a witness, parse existentially wrapped windows
-- from a 'String'.
readWindowLayout :: WindowLayout -> String -> WindowLayout
readWindowLayout (WindowLayout l) s
  | (Just x) <- readLayout s =
    WindowLayout (asTypeOf x l)
readWindowLayout l _ = l

-- | Serializes a window layout to a string.
serializeWindowLayout :: WindowLayout -> String
serializeWindowLayout (WindowLayout l) = serializeLayout l

type ScreenId = ()

type ScreenDetail = ()

type Tag = String

newtype ReadPtr a = ReadPtr (Ptr ())

instance Read (ReadPtr a) where
  readPrec = fmap (ReadPtr . intPtrToPtr) readPrec

instance Show (ReadPtr a) where
  show (ReadPtr ptr) = show (ptrToIntPtr ptr)

type Wetterhorn = StablePtr (Context, State)

data Context = Context
  { ctxForeignInterface :: ForeignInterface,
    ctxConfig :: Config WindowLayout
  }

defaultHooks :: Hooks
defaultHooks =
  Hooks
    { keyHook = \_ -> return (),
      surfaceHook = handleSurface,
      buttonHook = \_ -> return ()
    }

defaultConfig :: Config ()
defaultConfig =
  Config
    { hooks = defaultHooks,
      layout = (),
      resetHook = return ()
    }

data Hooks = Hooks
  { keyHook :: KeyEvent -> W (),
    surfaceHook :: SurfaceEvent -> W (),
    buttonHook :: ButtonEvent -> W ()
  }

data Config l = Config
  { layout :: l,
    hooks :: Hooks,
    resetHook :: W ()
  }

-- | Typeclass defining the set of types which can be used as state extensions
-- to the W monad. These state extensions may be persistent or not.
--
-- There are default implementations for all methods if the type implements
-- Read, Show and Default,
class (Typeable a) => ExtensionClass a where
  -- | The initial value used for the first time an extension is 'gotten' or
  -- demarshalling fails.
  initialValue :: a

  -- | Transforms a type into a string. If the type cannot be marshalled, this
  -- function should return Nothing.
  --
  -- If a type cannot be marshalled, it cannot persist across hot reloads.
  marshalExtension :: a -> Maybe String

  -- | Reads an extension from a string. If this type is not marshallable or
  -- reading fails, this function should return Nothing.
  demarshalExtension :: String -> Maybe a

  -- | If the type implements Default, use the default implementation.
  default initialValue :: (Default a) => a
  initialValue = def

  -- | If the type implements Show, use show for the marshalling.
  default marshalExtension :: (Show a) => a -> Maybe String
  marshalExtension = Just . show

  -- | If the type implements Read, use read for the demarshalling.
  default demarshalExtension :: (Read a) => String -> Maybe a
  demarshalExtension = readMaybe

data StateExtension where
  StateExtension :: (ExtensionClass a) => a -> StateExtension

-- | Puts a state extension.
xput :: forall a m. (ExtensionClass a, Wlike m) => a -> m ()
xput val = liftW $ do
  modify' $ \state@State {extensibleState = extensibleState} ->
    state
      { extensibleState =
          M.insert
            ( xRepr (Proxy :: Proxy a)
            )
            (Right $ StateExtension val)
            extensibleState
      }

-- | Modifies a state extension.
xmodify :: forall a m. (ExtensionClass a, Wlike m) => (a -> a) -> m ()
xmodify fn = xput . fn =<< xget

-- | Modifies a state extension in the monadic context.
xmodifyM :: forall a m. (ExtensionClass a, Wlike m) => (a -> m a) -> m ()
xmodifyM fn = (xput <=< fn) =<< xget

-- | Produces a string representation of a type used to key into the extensible
-- state map.
xRepr :: forall proxy a. (ExtensionClass a) => proxy a -> String
xRepr _ = tyconToStr $ someTypeRepTyCon (someTypeRep (Proxy :: Proxy a))
  where
    tyconToStr tc =
      printf "%s.%s.%s" (tyConPackage tc) (tyConModule tc) (tyConName tc)

-- | Gets a state extension.
xget :: forall a m. (ExtensionClass a, Wlike m) => m a
xget = do
  xs <- liftW $ gets extensibleState
  case M.lookup (xRepr (Proxy :: Proxy a)) xs of
    Just (Right (StateExtension a)) -> return (fromMaybe initialValue (cast a))
    Just (Left str) ->
      let v = fromMaybe initialValue (demarshalExtension str)
       in xput v >> return v
    Nothing ->
      xput (initialValue :: a) >> return initialValue

xgets :: forall a b m. (ExtensionClass a, Wlike m) => (a -> b) -> m b
xgets fn = fn <$> xget

-- State as it is marshalled. Used for derived instances of Show and Read.
data MarshalledState
  = MarshalledState
      (StackSet ScreenId ScreenDetail Tag String Window)
      (Set Window)
      [(String, String)]
  deriving (Show, Read)

data State = State
  { -- The datastructure containing the state of the windows.
    mapped :: StackSet ScreenId ScreenDetail Tag WindowLayout Window,
    -- | All the windows wetterhorn knows about, even if they are not mapped.
    allWindows :: Set Window,
    -- | Current set of hooks. The initial hooks are provided by the
    -- configuration, but the hooks can change during operation. This is how key
    -- sequences can be mapped.
    currentHooks :: Hooks,
    -- | Map from the typerep string to the state extension.
    extensibleState :: Map String (Either String StateExtension)
  }

-- | Initializes a "cold" state from a configuration. A cold state is the
-- initial state on startup. It is constrasted with a "hot" state, which is a
-- persisted state after a hot-reload.
initColdState :: Config WindowLayout -> IO State
initColdState Config {layout = layout, hooks = hooks} =
  return $
    State
      ( StackSet (Screen () () (Workspace "0" layout (Stack [] []))) [] []
      )
      mempty
      hooks
      mempty

-- | Marshals the serializable parts of the state to a string. This happens
-- during a hot-reload where some state must be saved to persist across hot
-- reloads.
marshalState :: State -> String
marshalState
  ( State
      { mapped = mapped,
        allWindows = allWindows,
        extensibleState = xs
      }
    ) =
    show $
      MarshalledState
        (mapLayout serializeWindowLayout mapped)
        allWindows
        (mapMaybe (\(k, v) -> (k,) <$> doMarshalEx v) (M.toList xs))
    where
      doMarshalEx (Left s) = Just s
      doMarshalEx (Right (StateExtension a)) = marshalExtension a

-- | Demarshals the string from "marshalState" into a state. Uses the provided
-- config to fill out non-persistent parts of the state.
demarshalState :: Config WindowLayout -> String -> State
demarshalState Config {hooks = hooks, layout = layout} str =
  State mapped allWindows hooks xs
  where
    ( MarshalledState
        (mapLayout (readWindowLayout layout) -> mapped)
        allWindows
        (fmap Left . M.fromList -> xs)
      ) = read str

-- | This is _the_ main monad used for Wetterhorn operations. Contains
-- everything required to operate. Contains the state, configuration and
-- interface to foreign code.
newtype W a = W (ReaderT Context (StateT State IO) a)
  deriving (Functor, Applicative, Monad, MonadState State, MonadIO)

-- | Let Config be the thing W is a reader for. There is already a way to get
-- the foreign interface in the context.
instance MonadReader (Config WindowLayout) W where
  local fn (W r) = W $ local (\(Context fi conf) -> Context fi (fn conf)) r
  ask = W $ ctxConfig <$> ask

runW :: W a -> (Context, State) -> IO (a, State)
runW (W fn) (ctx, st) = runStateT (runReaderT fn ctx) st

foreignInterface :: W ForeignInterface
foreignInterface = W $ ctxForeignInterface <$> ask

getSeat :: W (Ptr WlrSeat)
getSeat = (wio . ForeignInterface.getSeat) =<< foreignInterface

requestHotReload :: W ()
requestHotReload = (wio . ForeignInterface.requestHotReload) =<< foreignInterface

requestExit :: Int -> W ()
requestExit ec = (wio . flip ForeignInterface.requestExit ec) =<< foreignInterface

shellExec :: String -> W ()
shellExec = wio . ForeignInterface.doShellExec

wio :: IO a -> W a
wio = liftIO

-- | Type class to lift an arbitrary 'W' computation into another monad.
class (Monad m) => Wlike m where
  liftW :: W a -> m a

-- | Trivial instance of W for Wlike.
instance Wlike W where
  liftW = id

-- Default implementations for common handlers.

-- | handles a new surface event. This updates the state to reflect how it
-- should look in the harness.
handleSurface :: SurfaceEvent -> W ()
handleSurface (SurfaceEvent state (Window -> win)) =
  case state of
    Destroy ->
      modify $
        \st@State
           { allWindows = allWindows,
             mapped = mapped
           } ->
            st
              { allWindows = Set.delete win allWindows,
                mapped = StackSet.delete win mapped
              }
    Unmap -> modify $
      \st@State {mapped = mapped} ->
        st
          { mapped = StackSet.delete win mapped
          }
    Map -> modify $
      \st@State {mapped = mapped, allWindows = allWindows} ->
        st
          { mapped = StackSet.insertTiled win mapped,
            allWindows = Set.insert win allWindows
          }