{-# LANGUAGE GADTs #-}
{-# LANGUAGE FunctionalDependencies #-}
module DependencyGenerator
  ( DepGenM
  , DepGenM'
  , TokenableTo(..)
  , toTupleToken
  , evalDepGenM
  , inject
  , runFunction
  , runFunctionIO
  , runFunctionIO_
  , mapDepGenM
  , mapDepGenM_
  , forDepGenM
  , forDepGenM_
  , filterDepGenM
  , filterDepGenM'
  , zipDepGenM
  , untupleFstDepGenM
  , untupleSndDepGenM
  , untupleDepGenM
  , unzipFstDepGenM
  , unzipSndDepGenM
  , unzipDepGenM
  ) where

import Prelude hiding (String, FilePath)

import Types.Token (Token(..))
import Types.Value (Valuable(..))
import Types.FunctionIO (IsFunctionIO(..))
import Types.Function (IsFunction(..))
import Types.Dependency (Action(..), Dependency, makeDependency, makeUToken)

import Control.Monad.State (MonadState, State, runState, put, get)
import Control.Monad.Writer (MonadWriter, WriterT, runWriterT, tell)

newtype DepGenM a = DepGenM { unDepGenM :: WriterT [Dependency] (State Int) a }
  deriving (Functor, Applicative, Monad, MonadState Int, MonadWriter [Dependency])

type DepGenM' a = DepGenM (Token a)

runDepGenM :: Int -> DepGenM a -> ((a, [Dependency]), Int)
runDepGenM top m = runState (runWriterT (unDepGenM m)) top

evalDepGenM :: DepGenM () -> [Dependency]
evalDepGenM = snd . fst . runDepGenM 0

tellDep :: Dependency -> DepGenM ()
tellDep dep = tell [dep]

newToken :: DepGenM (Token a)
newToken = do
  top <- get
  let top' = top + 1
      target = Token top'
  put top'
  pure target

genDependencyM :: (Token a -> DepGenM Dependency) -> DepGenM (Token a)
genDependencyM f = do
  target <- newToken
  result <- f target
  tellDep result
  pure target

genDependency :: (Token a -> Dependency) -> DepGenM (Token a)
genDependency f = genDependencyM (pure . f)

inject :: Valuable a => a -> DepGenM (Token a)
inject x = genDependency (makeDependency NoToken (Inject (toValue x)))

runFunction :: IsFunction f a b => f -> Token a -> DepGenM (Token b)
runFunction f input = genDependency (makeDependency input (Function f))

runFunctionIO :: IsFunctionIO f a b => f -> Token a -> DepGenM (Token b)
runFunctionIO f input = genDependency (makeDependency input (FunctionIO f))

runFunctionIO_ :: IsFunctionIO f a () => f -> Token a -> DepGenM ()
runFunctionIO_ f input = tellDep $ makeDependency input (FunctionIO f) NoToken

class TokenableTo t s | s -> t where
  toToken :: s -> DepGenM (Token t)

instance TokenableTo a (Token a) where
  toToken = pure

instance TokenableTo [a] [Token a] where
  toToken = pure . ListToken

instance TokenableTo a (DepGenM (Token a)) where
  toToken = id

instance TokenableTo [a] [DepGenM (Token a)] where
  toToken = fmap ListToken . sequence

toTupleToken :: (TokenableTo ta a, TokenableTo tb b) => a -> b -> DepGenM (Token (ta, tb))
toTupleToken a b = TupleToken <$> toToken a <*> toToken b

mapDepGenM :: TokenableTo [a] v => (Token a -> DepGenM (Token b)) -> v -> DepGenM (Token [b])
mapDepGenM f input = do
  input' <- toToken input
  genDependencyM $ \target -> do
    top <- get
    let (((innerInp, innerOutp), subDeps), top') = runDepGenM top $ do
          inp <- newToken
          outp <- f inp
          pure (inp, outp)
    put top'
    pure (makeDependency input' (MapComp subDeps (makeUToken innerInp) (makeUToken innerOutp)) target)

mapDepGenM_ :: TokenableTo [a] v => (Token a -> DepGenM ()) -> v -> DepGenM ()
mapDepGenM_ f input = do
  _ <- mapDepGenM (\x -> f x >> pure NoToken) input
  pure ()

forDepGenM :: TokenableTo [a] v => v -> (Token a -> DepGenM (Token b)) -> DepGenM (Token [b])
forDepGenM = flip mapDepGenM

forDepGenM_ :: TokenableTo [a] v => v -> (Token a -> DepGenM ()) -> DepGenM ()
forDepGenM_ = flip mapDepGenM_

filterDepGenM' :: (TokenableTo [Bool] v, TokenableTo [a] u) => v -> u -> DepGenM (Token [a])
filterDepGenM' mask input = do
  tup <- toTupleToken input mask
  genDependency (makeDependency tup FilterComp)

filterDepGenM :: TokenableTo [a] v => (Token a -> DepGenM (Token Bool)) -> v -> DepGenM (Token [a])
filterDepGenM f input = do
  mask <- mapDepGenM f input
  filterDepGenM' mask input

zipDepGenM :: (TokenableTo [a] v, TokenableTo [b] u) => v -> u -> DepGenM (Token [(a, b)])
zipDepGenM a b = do
  a' <- toToken a
  b' <- toToken b
  pure $ ZipToken a' b'

untupleFstDepGenM :: TokenableTo (a, b) t => t -> DepGenM (Token a)
untupleFstDepGenM t = do
  t' <- toToken t
  case t' of
    TupleToken a _ -> pure a
    Token _ -> genDependency (makeDependency t' UntupleFst)

untupleSndDepGenM :: TokenableTo (a, b) t => t -> DepGenM (Token b)
untupleSndDepGenM t = do
  t' <- toToken t
  case t' of
    TupleToken _ b -> pure b
    Token _ -> genDependency (makeDependency t' UntupleSnd)

untupleDepGenM :: TokenableTo (a, b) t => t -> DepGenM (Token a, Token b)
untupleDepGenM t = do
  t' <- toToken t
  a <- untupleFstDepGenM t'
  b <- untupleSndDepGenM t'
  pure (a, b)

unzipFstDepGenM :: TokenableTo [(a, b)] t => t -> DepGenM (Token [a])
unzipFstDepGenM t = do
  t' <- toToken t
  case t' of
    ZipToken a _ -> pure a
    _ -> genDependency (makeDependency t' UnzipFst)

unzipSndDepGenM :: TokenableTo [(a, b)] t => t -> DepGenM (Token [b])
unzipSndDepGenM t = do
  t' <- toToken t
  case t' of
    ZipToken _ b -> pure b
    _ -> genDependency (makeDependency t' UnzipSnd)

unzipDepGenM :: TokenableTo [(a, b)] t => t -> DepGenM (Token [a], Token [b])
unzipDepGenM t = do
  t' <- toToken t
  a <- unzipFstDepGenM t'
  b <- unzipSndDepGenM t'
  pure (a, b)