{-# LANGUAGE DeriveFunctor, GADTs #-}
module GEODE.Metadata.SplitContext
  ( Current(..)
  , Field(..)
  , HeadWords
  , SplitContext(..)
  , SplitContextT(..)
  , evalSplit
  , evalSplitT
  , next
  , runSplit
  , runSplitT ) where

import Control.Monad.Identity (Identity(..))
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans (MonadTrans(..))
import Data.Map as Map (Map, (!?), delete, empty, insert)
import Data.Text (Text)

type HeadWords = Map Text Int
data Current = Current { page :: Int, rank :: Int, headWords :: HeadWords }
data Field a where
  Page :: Field Int
  Rank :: Field Int
  HeadWord :: Text -> Field (Maybe Int)

class Monad m => SplitContext m where
  context :: m Current
  update :: (Current -> Current) -> m ()

  get :: Field a -> m a
  get Page = page <$> context
  get Rank = rank <$> context
  get (HeadWord h) = ((!? h) . headWords) <$> context

  set :: Field a -> a -> m ()
  set Page p = update $ \c -> c { page = p }
  set Rank r = update $ \c -> c { rank = r }
  set (HeadWord h) Nothing = update $ \c -> c { headWords = delete h $ headWords c }
  set (HeadWord h) (Just n) = update $ \c -> c { headWords = insert h n $ headWords c }

  editLookup :: Field a -> (a -> a) -> m a
  editLookup field f = get field >>= (\a -> set field a *> pure a) . f

  lookupEdit :: Field a -> (a -> a) -> m a
  lookupEdit field f = get field >>= \a -> set field (f a) *> pure a

  edit :: Field a -> (a -> a) -> m ()
  edit field = (() <$) . editLookup field

  {-# MINIMAL context, update #-}

newtype SplitContextT m a =
  SplitContextT { runWithContext :: Current -> m (a, Current) }
  deriving (Functor)

instance Monad m => Applicative (SplitContextT m) where
  pure a = SplitContextT $ \current -> pure (a, current)
  sCF <*> sCA = SplitContextT $ \current0 -> do
    (f, current1) <- runWithContext sCF current0
    (a, current2) <- runWithContext sCA current1
    pure (f a, current2)

instance Monad m => Monad (SplitContextT m) where
  sCA >>= f = SplitContextT $ \current0 -> do
    (a, current1) <- runWithContext sCA current0
    runWithContext (f a) current1

instance Monad m => SplitContext (SplitContextT m) where
  context = SplitContextT $ \current -> pure (current, current)
  update f = SplitContextT $ \current -> pure ((), f current)

instance MonadTrans SplitContextT where
  lift m = SplitContextT $ \current -> m >>= \a -> pure (a, current)

instance MonadIO m => MonadIO (SplitContextT m) where
  liftIO = lift . liftIO

next :: SplitContext m => Field a -> m Int
next f@(HeadWord _) = lookupEdit f (Just . maybe 1 (+1)) >>= pure . maybe 0 id
next Page = lookupEdit Page (+1)
next Rank = lookupEdit Rank (+1)

runSplitT :: SplitContextT m a -> m (a, Current)
runSplitT = flip runWithContext $
  Current { page = 1, rank = 1, headWords = Map.empty }

evalSplitT :: Functor m => SplitContextT m a -> m a
evalSplitT = fmap fst . runSplitT

runSplit :: SplitContextT Identity a -> (a, Current)
runSplit = runIdentity . runSplitT

evalSplit :: SplitContextT Identity a -> a
evalSplit = fst . runSplit