module ComputationM
  ( ComputationM
  , evalComputationM
  , inject
  -- , inject'
  , mapComputationM
  , mapComputationM_
  , filterComputationM
  , runFunction
  , runFunctionIO
  ) where

import Types
import Dependency

import Control.Monad.State
import Control.Monad.Writer
import Language.Haskell.TH.Syntax

newtype ComputationM' a = ComputationM { unComputationM :: WriterT [DependencyUntyped] (State Int) a }
  deriving (Functor, Applicative, Monad, MonadState Int, MonadWriter [DependencyUntyped])

type ComputationM a = ComputationM' (Token a)

evalComputationM' :: Int -> ComputationM () -> ([DependencyUntyped], Int)
evalComputationM' top m = runState (execWriterT (unComputationM m)) top

evalComputationM :: ComputationM () -> [DependencyUntyped]
evalComputationM m = fst (evalComputationM' 0 m)

tellDep :: (Show a, Show b, Lift a, Lift b) => Dependency a b -> ComputationM' ()
tellDep dep = tell [makeUntyped dep]

genDependency' :: (Show u, Show a, Lift u, Lift a) => (Token a -> ComputationM' (Dependency u a)) -> ComputationM a
genDependency' f = do
  top <- get
  let top' = top + 1
      target = Token top'
  put top'
  result <- f target
  tellDep result
  pure target

genDependency :: (Show u, Show a, Lift u, Lift a) => (Token a -> Dependency u a) -> ComputationM a
genDependency f = genDependency' (pure . f)

-- inject :: (Show a, Lift a) => a -> ComputationM a
-- inject x = genDependency (Dependency NoToken (Inject x))

inject :: (Show a, Lift a, Valuable a) => a -> ComputationM a
inject x = genDependency (Dependency NoToken (Inject (toValue x)))

-- inject' :: ImageConversionSettings -> ComputationM ImageConversionSettings
-- -- inject' x = genDependency (Dependency NoToken (InjectImageConversionSettings x))
-- inject' x = genDependency (Dependency NoToken (Inject WImageConversionSettings x))

getListElem :: (Show a, Lift a) => Token [a] -> ComputationM a
getListElem outer = genDependency (Dependency outer GetListElem)

setListElem :: (Show a, Lift a) => Token a -> Token [a] -> ComputationM ()
setListElem a outer = do
  tellDep (Dependency a SetListElem outer)
  pure NoToken

-- runFunction :: (Show a, Show b, Lift a, Lift b, IsFunction f a b) => f -> Token a -> ComputationM b
-- runFunction f input = genDependency (Dependency input (Function f))

-- runFunctionIO :: (Show a, Show b, Lift a, Lift b, IsFunctionIO f a b) => f -> Token a -> ComputationM b
-- runFunctionIO f input = genDependency (Dependency input (FunctionIO f))

runFunction :: (Show a, Show b, Lift a, Lift b) => Function -> Token a -> ComputationM b
runFunction f input = genDependency (Dependency input (Function f))

runFunctionIO :: (Show a, Show b, Lift a, Lift b) => FunctionIO -> Token a -> ComputationM b
runFunctionIO f input = genDependency (Dependency input (FunctionIO f))

mapComputationM :: (Show a, Show b, Lift a, Lift b) => (Token a -> ComputationM b) -> Token [a] -> ComputationM [b]
mapComputationM f input = genDependency' $ \target -> do
  top <- get
  let (res, top') = evalComputationM' top $ do
        inp <- getListElem input
        outp <- f inp
        setListElem outp target
  put top'
  pure (Dependency input (MapComp res) target)

mapComputationM_ :: (Show a, Lift a) => (Token a -> ComputationM ()) -> Token [a] -> ComputationM ()
mapComputationM_ f input = do
  _ <- mapComputationM f input
  pure NoToken

filterComputationM :: (Show a, Lift a) => (Token a -> ComputationM Bool) -> Token [a] -> ComputationM [a]
filterComputationM f input = do
  conds <- mapComputationM f input
  genDependency (Dependency (TupleToken input conds) FilterComp)