{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE OverloadedStrings #-}

module Hsm.Core.Fsm
  ( FsmState(FsmState)
  , FsmOutput(FsmOutput)
  , FsmResult(FsmResult)
  , fsm
  ) where

import Data.Maybe (fromJust, isJust)
import Data.Text (Text)
import Effectful (Eff, (:>))
import Effectful.Log (Log, LogLevel, localDomain, logAttention_, logTrace_)
import Effectful.Reader.Static (Reader, ask)
import Effectful.State.Static.Local (State, get, put)
import Hsm.Core.Log (logTup)
import Streamly.Data.Stream qualified as S (Stream, mapM, takeWhile)

data FsmState i o env sta =
  FsmState Text (i -> env -> sta -> FsmOutput i o env sta)

data FsmOutput i o env sta =
  FsmOutput (Maybe (FsmResult i o env sta)) [(LogLevel, Text)]

data FsmResult i o env sta =
  FsmResult o sta (FsmState i o env sta)

-- Finite state machines allow processing of stream elements using pure
-- functions. One or more FSMs can be included within a `Streamly` pipeline.
fsm ::
     forall i o env sta es.
     ( Log :> es
     , Reader env :> es
     , State (FsmState i o env sta) :> es
     , State sta :> es
     )
  => S.Stream (Eff es) i
  -> S.Stream (Eff es) o
fsm = S.mapM (return . fromJust) . S.takeWhile isJust . S.mapM run
  where
    exit = do
      logAttention_ "No state returned, exiting FSM"
      return Nothing
    push (FsmResult out sta next) = do
      put sta
      put next
      return $ Just out
    run input =
      localDomain "fsm" $ do
        FsmState name action <- get
        sta <- get @sta
        env <- ask @env
        logTrace_ $ "Entering state " <> name
        FsmOutput res logs <- return $ action input env sta
        localDomain name $ mapM_ logTup logs
        logTrace_ $ "Exiting state " <> name
        maybe exit push res