{-# 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 Types.Token (Token(..))
import Types.Value
import Types.Functions (IsFunction(), IsFunctionIO(..))
import Types.Dependency (Action(..), Dependency, makeDependency, unListType, unTupleType)

import Type.Reflection (Typeable, TypeRep, typeRep)
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 :: (Typeable a, Show a) => DepGenM (Token a)
newToken = do
  top <- get
  let top' = top + 1
      target = Token top'
  put top'
  pure target

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

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

inject :: (Show a, Typeable 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 (Show t, Typeable t) => TokenableTo t s | s -> t where
  toToken :: s -> DepGenM (Token t)
  tokenTypeRep :: s -> TypeRep t

instance (Show a, Typeable a) => TokenableTo a (Token a) where
  toToken = pure
  tokenTypeRep _ = typeRep

instance (Show a, Typeable a) => TokenableTo [a] [Token a] where
  toToken = pure . ListToken
  tokenTypeRep _ = typeRep

instance (Show a, Typeable a) => TokenableTo a (DepGenM (Token a)) where
  toToken = id
  tokenTypeRep _ = typeRep

instance (Show a, Typeable a) => TokenableTo [a] [DepGenM (Token a)] where
  toToken = fmap ListToken . sequence
  tokenTypeRep _ = typeRep

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, Typeable a, Show a, Typeable b, Show b) => (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 typeRep typeRep subDeps innerInp innerOutp) target)

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

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

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

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

filterDepGenM :: (TokenableTo [a] v, Typeable a, Show a) => (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, Typeable a, Show a, Typeable b, Show b) => 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, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token a)
untupleFstDepGenM t = do
  t' <- toToken t
  case t' of
    TupleToken a _ ->
      pure a
    Token _ ->
      let tr = tokenTypeRep t
          (ta, tb) = unTupleType tr
      in genDependency (makeDependency t' (UntupleFst ta tb))

untupleSndDepGenM :: (TokenableTo (a, b) t, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token b)
untupleSndDepGenM t = do
  t' <- toToken t
  case t' of
    TupleToken _ b ->
      pure b
    Token _ ->
      let tr = tokenTypeRep t
          (ta, tb) = unTupleType tr
      in genDependency (makeDependency t' (UntupleSnd ta tb))

untupleDepGenM :: (TokenableTo (a, b) t, Typeable a, Show a, Typeable b, Show b) => 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, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token [a])
unzipFstDepGenM t = do
  t' <- toToken t
  case t' of
    ZipToken a _ ->
      pure a
    _ ->
      let tr = tokenTypeRep t
          (ta, tb) = unTupleType $ unListType tr
      in genDependency (makeDependency t' (UnzipFst ta tb))

unzipSndDepGenM :: (TokenableTo [(a, b)] t, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token [b])
unzipSndDepGenM t = do
  t' <- toToken t
  case t' of
    ZipToken _ b ->
      pure b
    _ ->
      let tr = tokenTypeRep t
          (ta, tb) = unTupleType $ unListType tr
      in genDependency (makeDependency t' (UnzipSnd ta tb))

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