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

module Hsm.Core.Fsm
  ( FsmState (FsmState)
  , FsmResult (FsmResult)
  , FsmOutput (FsmOutput)
  , fsm
  , fsmStream
  , pLogAttention
  , pLogInfo
  , pLogTrace
  )
where

import Control.Monad (forM_)
import Data.Aeson.Types (emptyObject)
import Data.Composition ((.:))
import Data.Function.Slip (slipl)
import Data.List (singleton)
import Data.Maybe (fromJust, isJust)
import Data.Text (Text)
import Effectful (Eff, (:>))
import Effectful.Log qualified as L
import Effectful.Reader.Static (Reader, ask)
import Effectful.State.Static.Local (State, get, put)
import Streamly.Data.Stream (Stream, mapM, takeWhile)
import Prelude hiding (mapM, takeWhile)

data FsmState i o e s
  = FsmState Text (i -> e -> s -> FsmOutput i o e s)

data FsmResult i o e s
  = FsmResult o s (FsmState i o e s)

data FsmOutput i o e s
  = FsmOutput (Maybe (FsmResult i o e s)) [(L.LogLevel, Text)]

type FsmConstraint i o e s es =
  ( L.Log :> es
  , Reader e :> es
  , State (FsmState i o e s) :> es
  , State s :> es
  )

fsm :: forall i o e s es. FsmConstraint i o e s es => i -> Eff es (Maybe o)
fsm input =
  L.localDomain "fsm" $ do
    FsmState name action <- get
    state <- get
    env <- ask
    L.logTrace_ $ "Entering state " <> name
    FsmOutput res logs <- return $ action input env state
    L.localDomain name $ forM_ logs $ uncurry $ slipl L.logMessage emptyObject
    L.logTrace_ $ "Exiting state " <> name
    maybe exit push res
 where
  push :: FsmResult i o e s -> Eff es (Maybe o)
  push (FsmResult output state next) = do
    put state
    put next
    return $ Just output

  exit :: Eff es (Maybe o)
  exit = do
    L.logAttention_ "No state returned, exiting FSM"
    return Nothing

fsmStream
  :: forall i o e s es
   . FsmConstraint i o e s es
  => Stream (Eff es) i
  -> Stream (Eff es) o
fsmStream = mapM (return . fromJust) . takeWhile isJust . mapM (fsm @_ @_ @e @s)

type LogList = [(L.LogLevel, Text)]

pLog :: L.LogLevel -> Text -> LogList
pLog = singleton .: (,)

pLogAttention :: Text -> LogList
pLogAttention = pLog L.LogAttention

pLogInfo :: Text -> LogList
pLogInfo = pLog L.LogInfo

pLogTrace :: Text -> LogList
pLogTrace = pLog L.LogTrace