-- | Minimal monadic push-based reactivity, modeled after JavaScript example in
--   https://dev.to/ryansolid/building-a-reactive-library-from-scratch-1i0p.
module Main where

import Control.Monad.State.Strict
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Control.Monad (forM_, unless)

-- | Reactive monad, wrapping StateT over IO for state and effects.
newtype Reactive a = Reactive (StateT Context IO a)
  deriving ( Functor
           , Applicative
           , Monad
           , MonadState Context
           )

instance MonadIO Reactive where
  liftIO io = Reactive (lift io)
  
-- | Reactive state context.
data Context = Context
  { nextId              :: Int
  , subscriptions       :: Map SignalId (Int, Set ComputationId)
  , dependencies        :: Map ComputationId ([SignalId], Reactive ())
  , activeSubscriptions :: Set ComputationId
  }

type SignalId = Int
type ComputationId = Int

-- | Runs a Reactive computation.
compute :: Reactive a -> IO a
compute (Reactive m) = evalStateT m (Context 0 Map.empty Map.empty Set.empty)

-- | Creates a signal with an initial value.
createSignal :: Int -> Reactive SignalId
createSignal initialValue = do
  s <- get
  let id = nextId s
  put s { nextId = id + 1, subscriptions = Map.insert id (initialValue, Set.empty) (subscriptions s) }
  return id

-- | Reads a signal's value.
readSignal :: SignalId -> Reactive Int
readSignal signalId = gets $ \s -> fst (subscriptions s Map.! signalId)

-- | Updates a signal and triggers effects.
writeSignal :: SignalId -> Int -> Reactive ()
writeSignal signalId newValue = do
  modify $ \s -> s { subscriptions = Map.adjust (\(_, subs) -> (newValue, subs)) signalId (subscriptions s) }
  modify $ \s -> s { activeSubscriptions = Set.empty }  -- Clear for new cycle
  dependencies <- gets dependencies
  mapM_ execute $ Map.toList dependencies
  where
    execute (computationId, (deps, computation))
      | signalId `elem` deps = do
          isRunning <- gets $ \s -> Set.member computationId (activeSubscriptions s)
          unless isRunning $ do
            modify $ \s -> s { activeSubscriptions = Set.insert computationId (activeSubscriptions s) }
            computation
      | otherwise = return ()

-- | Creates an effect that runs on signal changes.
createEffect :: [SignalId] -> Reactive () -> Reactive ()
createEffect signalIds computation = do
  s <- get
  let id = nextId s
  put s { nextId = id + 1 }
  modify $ \s -> s { dependencies = Map.insert id (signalIds, computation) (dependencies s) }
  modify $ \s -> s { subscriptions = foldr (\sid -> Map.adjust (\(val, subs) -> (val, Set.insert id subs)) sid) (subscriptions s) signalIds }
  computation

-- | Creates a memoized signal from a computation.
createMemo :: [SignalId] -> Reactive Int -> Reactive SignalId
createMemo signalIds computation = do
  signalId <- createSignal 0
  createEffect signalIds $ do
    value <- computation
    writeSignal signalId value
  return signalId

-- | A classic "counter" example.
main :: IO ()
main = compute $ do
  count <- createSignal 0
  doubleCount <- createMemo [count] $ do
    c <- readSignal count
    return (c * 2)

  -- Prints "Count: 0, Double: 0"
  createEffect [count, doubleCount] $ do
    c <- readSignal count
    dc <- readSignal doubleCount
    liftIO $ putStrLn $ "Count: " ++ show c ++ ", Double: " ++ show dc

  -- Prints "Count: 1, Double: 2"
  writeSignal count 1
  -- Prints "Count: 2, Double: 4"
  writeSignal count 2