{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE PatternSynonyms #-}
module DependencyRunner
  ( DepRunM
  , runDeps
  , runDepRunMIO
  ) where

import Types (evalFunction, evalFunctionIO)
import Types.Value
import Types.Token
import Types.Dependency

import Type.Reflection (Typeable, TypeRep, typeRep, pattern App)
import Data.Map (Map)
import qualified Data.Map as M
import Control.Monad (void, forM)
import Control.Monad.State (MonadState, MonadIO, StateT, evalStateT, get, modify, liftIO)

data ValueExistence = Evaluated Value
                    | NotEvaluated (DepRunM Value)

newtype DepRunM a = DepRunM { unDepRunM :: StateT (Map Int ValueExistence) IO a }
  deriving (Functor, Applicative, Monad, MonadIO, MonadState (Map Int ValueExistence))

runDeps :: [Dependency] -> DepRunM ()
runDeps = mapM_ runDep

runDepRunMIO :: DepRunM a -> IO a
runDepRunMIO m = evalStateT (unDepRunM m) M.empty

evaluate :: ValueExistence -> DepRunM Value
evaluate = \case
  Evaluated v -> pure v
  NotEvaluated m -> m

runDep :: Dependency -> DepRunM ()
runDep (Dependency _ a action _ b) =
  if actionTouchesFilesystem action
  then void m
  else putTokenValue b $ NotEvaluated m
  where m :: DepRunM Value
        m = do
          input <- getTokenValue a
          result <- runAction action input
          putTokenValue b $ Evaluated result
          liftIO $ do
            putStrLn ("input: " ++ show input)
            putStrLn ("action: " ++ show action)
            putStrLn ("output: " ++ show result)
            putStrLn "----------"
          pure result

getTokenValue :: (Typeable a, Show a) => Token a -> DepRunM Value
getTokenValue token = case token of
  Token i -> do
    m <- get
    evaluate (m M.! i)
  TupleToken a b -> do
    va <- getTokenValue a
    vb <- getTokenValue b
    pure $ toValueRep (tokenTypeRep token) (fromValue va, fromValue vb)
  ZipToken a b -> do
    va <- getTokenValue a
    vb <- getTokenValue b
    pure $ toValueRep (tokenTypeRep token) $ zip (fromValue va) (fromValue vb)
  ListToken ts -> do
    vs <- mapM getTokenValue ts
    pure $ toValueRep (tokenTypeRep token) (map fromValue vs)
  NoToken ->
    pure $ toValue ()

putTokenValue :: Token a -> ValueExistence -> DepRunM ()
putTokenValue t e = case t of
  Token i ->
    modify $ M.insert i e
  NoToken ->
    pure ()
  _ ->
    error "unexpected"

tupApp :: TypeRep a -> TypeRep b -> TypeRep (a, b)
tupApp ta tb = App (App (typeRep @(,)) ta) tb

listApp :: TypeRep a -> TypeRep [a]
listApp ta = App (typeRep @[]) ta

runAction :: Action -> Value -> DepRunM Value
runAction action input = case action of
  Function f ->
    pure $ toValue $ evalFunction f $ fromValue input
  InlineFunction ta tb (F f) ->
    pure $ toValueRep tb $ f $ fromValueRep ta input
  FunctionIO f ->
    liftIO (toValue <$> evalFunctionIO f (fromValue input))
  Inject v ->
    pure v
  FilterComp t ->
    let tl = listApp t
        (vs, mask) = fromValueRep (tupApp tl (typeRep @[Bool])) input
    in pure $ toValueRep tl $ map fst $ filter snd $ zip vs mask
  UntupleFst ta tb ->
    pure $ toValueRep ta $ fst $ fromValueRep (tupApp ta tb) input
  UntupleSnd ta tb ->
    pure $ toValueRep tb $ snd $ fromValueRep (tupApp ta tb) input
  UnzipFst ta tb ->
    pure $ toValueRep (listApp ta) $ map fst $ fromValueRep (listApp (tupApp ta tb)) input
  UnzipSnd ta tb ->
    pure $ toValueRep (listApp tb) $ map snd $ fromValueRep (listApp (tupApp ta tb)) input
  MapComp ta tb subDeps innerInput innerOutput ->
    (toValueRep (listApp tb) <$>) $ forM (fromValueRep (listApp ta) input) $ \x -> do
      putTokenValue innerInput $ Evaluated $ toValueRep ta x
      runDeps subDeps
      fromValueRep tb <$> getTokenValue innerOutput