{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Montis.Core.Monad where

import Control.Monad.Identity (Identity (Identity))
import Control.Monad.Reader
import Control.Monad.State (MonadState, StateT (runStateT), gets, modify')
import Data.Default.Class (Default (def))
import Data.Map qualified as Map
import Data.Maybe (fromMaybe)
import Data.Typeable
import Foreign (StablePtr)
import Montis.Core.Extensions (Extension (Extension), typeRepr)
import Montis.Core.State

-- | A Config type for the Montis monad.
type MontisConfig = Config Montis

-- | A Context type specific for the Montis monad.
type MontisContext = Context Montis

-- | A State type for the Montis monad.
type MontisState = State Montis

-- | The Opaque State Type passed between the world and the runtime. The
-- OpqStT *is* the opq_st_t from the runtime code.
type OpqStT = StablePtr (MontisContext, MontisState)

-- | The Montis monad is a Reader over an immutable context plus a State over
--   mutable runtime state, all ultimately running in IO for host effects.
newtype Montis a where
  Montis :: (ReaderT MontisContext (StateT MontisState IO) a) -> Montis a
  deriving (Functor, Applicative, Monad, MonadState MontisState, MonadIO)

-- | Reader access is scoped to the config portion of the full context; this
--   keeps world code from mutating the context while still allowing read-only
--   access to configuration.
instance MonadReader MontisConfig Montis where
  ask :: Montis MontisConfig
  ask = Montis $ asks ctxConfig

  reader :: (MontisConfig -> a) -> Montis a
  reader fn = Montis $ asks (fn . ctxConfig)

  local :: (MontisConfig -> MontisConfig) -> Montis a -> Montis a
  local cfn (Montis fn) =
    Montis $ local (\ctx -> ctx {ctxConfig = cfn (ctxConfig ctx)}) fn

-- | Access the world self pointer stored in the context.
getSelfPtr :: Montis SelfPtr
getSelfPtr = Montis $ asks ctxSelfPtr

-- | Run a Montis action with a fixed context and initial state, returning the
--   result value and the updated state.
runMontis :: MontisContext -> MontisState -> Montis a -> IO (a, MontisState)
runMontis ctx initState (Montis m) = runStateT (runReaderT m ctx) initState

-- | The standard default config.
defaultConfig :: MontisConfig
defaultConfig =
  Config
    { startingHooks =
        Hooks
          { keyHook = const (return ()),
            surfaceHook = const (return ()),
            buttonHook = const (return ()),
            motionHook = const (return ())
          },
      -- Default hooks are no-ops except for basic printing, which makes the
      -- system usable without extra wiring during development.
      startupHook = return (),
      resetHook = return (),
      -- Extensions start empty; callers can register config extensions as needed.
      configExtensions = mempty
    }

-- | Store a typed extension in the extensible state map under its TypeRep.
--   The value is wrapped so the map remains heterogeneously typed.
xStatePut :: forall a. (StateExtension a) => a -> Montis ()
xStatePut xst = do
  modify'
    ( \st ->
        st
          { extensibleState =
              Map.insert
                (typeRepr (Identity xst))
                (Right (Extension xst))
                (extensibleState st)
          }
    )

-- | Retrieve a typed extension, demarshalling it if needed and caching it back.
--   When the extension is stored in marshalled form, it is decoded and then
--   reinserted so future lookups are fast.
xStateGet :: forall a. (StateExtension a) => Montis a
xStateGet = do
  mp <- gets extensibleState
  case lookupByType (Proxy :: Proxy a) mp of
    Nothing -> return initialValue
    Just (Right (Extension v)) -> return $ fromMaybe initialValue (cast v)
    Just (Left s) -> do
      let x = (demarshalExtension s :: Maybe a)
       in forM_ x xStatePut >> return (fromMaybe initialValue x)

-- | Modifies the typed extension bi the given function.
xStateModify :: forall a. (StateExtension a) => (a -> a) -> Montis ()
xStateModify fn = do
  s <- xStateGet
  (xStatePut . fn) s

xStateModifyM :: forall a. (StateExtension a) => (a -> Montis a) -> Montis ()
xStateModifyM fn = do
  s <- xStateGet
  xStatePut =<< fn s

-- | Retrieve a typed configuration extension or return the default
--   instance if the extension had not been configured.
xConfigGet :: forall a. (Typeable a, Default a) => Montis a
xConfigGet = do
  exts <- asks configExtensions
  return $
    fromMaybe def $
      Map.lookup (typeRepr (Proxy :: Proxy a)) exts
        >>= (\(Extension a) -> cast a)


mio :: IO a -> Montis a
mio = liftIO