module Quox.EffExtra

import public Control.Eff

import Data.IORef


export
localAt : (0 lbl : tag) -> Has (StateL lbl s) fs =>
          (s -> s) -> Eff fs a -> Eff fs a
localAt lbl f act = do
  old <- getAt lbl
  modifyAt lbl f *> act <* putAt lbl old

export %inline
localAt_ : (0 lbl : tag) -> Has (StateL lbl s) fs =>
           s -> Eff fs a -> Eff fs a
localAt_ lbl x = localAt lbl $ const x

export %inline
local : Has (State s) fs => (s -> s) -> Eff fs a -> Eff fs a
local = localAt ()

export %inline
local_ : Has (State s) fs => s -> Eff fs a -> Eff fs a
local_ = localAt_ ()


public export
record StateRes s a where
  constructor SR
  state  : s
  result : a

export
stateAt : (0 lbl : tag) -> Has (StateL lbl s) fs =>
          (s -> StateRes s a) -> Eff fs a
stateAt lbl f = do
  s <- getAt lbl
  let out = f s
  putAt lbl out.state
  pure out.result

export %inline
state : Has (State s) fs => (s -> StateRes s a) -> Eff fs a
state = stateAt ()


export
hasDrop : (0 neq : Not (a = b)) ->
          (ha : Has a fs) => (hb : Has b fs) =>
          Has a (drop fs hb)
hasDrop neq {ha = Z}    {hb = Z}    = void $ neq Refl
hasDrop neq {ha = S ha} {hb = Z}    = ha
hasDrop neq {ha = Z}    {hb = S hb} = Z
hasDrop neq {ha = S ha} {hb = S hb} = S $ hasDrop neq {ha, hb}

private
0 ioNotState : Not (IO = StateL _ _)
ioNotState Refl impossible

export
runStateIORefAt : (0 lbl : tag) -> (Has IO fs, Has (StateL lbl s) fs) =>
                  IORef s -> Eff fs a -> Eff (fs - StateL lbl s) a
runStateIORefAt lbl ref act = do
  let hh : Has IO (fs - StateL lbl s) := hasDrop ioNotState
  (val, st) <- runStateAt lbl !(readIORef ref) act
  writeIORef ref st $> val

export %inline
runStateIORef : (Has IO fs, Has (State s) fs) =>
                IORef s -> Eff fs a -> Eff (fs - State s) a
runStateIORef = runStateIORefAt ()


export %inline
evalStateAt : (0 lbl : tag) -> Has (StateL lbl s) fs =>
              s -> Eff fs a -> Eff (fs - StateL lbl s) a
evalStateAt lbl s act = map fst $ runStateAt lbl s act

export %inline
evalState : Has (State s) fs => s -> Eff fs a -> Eff (fs - State s) a
evalState = evalStateAt ()



public export
data Length : List a -> Type where
  Z : Length []
  S : Length xs -> Length (x :: xs)

export
subsetWith : Length xs => (forall z. Has z xs -> Has z ys) ->
             Subset xs ys
subsetWith @{Z}     f = []
subsetWith @{S len} f = f Z :: subsetWith (f . S)

export
subsetSelf : Length xs => Subset xs xs
subsetSelf = subsetWith id

export
subsetTail : Length xs => Subset xs (x :: xs)
subsetTail = subsetWith S


-- [fixme] allow the error to be anywhere in the effect list
export
wrapErrAt : Length fs => (0 lbl : tag) -> (e -> e) ->
            Eff (ExceptL lbl e :: fs) a -> Eff (ExceptL lbl e :: fs) a
wrapErrAt lbl f act =
  rethrowAt lbl . mapFst f =<< lift @{subsetTail} (runExceptAt lbl act)

export %inline
wrapErr : Length fs => (e -> e) ->
          Eff (Except e :: fs) a -> Eff (Except e :: fs) a
wrapErr = wrapErrAt ()


export %inline
runIO : (MonadRec io, HasIO io) => Eff [IO] a -> io a
runIO act = runEff act [liftIO]