{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE Arrows #-}

module AFRP
  ( Mealy(..)
  , eff
  , withEntities
  , Event(..)
  , hold
  , events
  , switch
  , preMapAccum
  , preMapAccumRequest
  , mapAccum
  , mapAccumRequest
  , changes
  , whenA
  , filterA
  , thenA
  , (>>|)
  , toEvent
  , lMerge
  , Request(..)
  , edge
  , dropFirst
  , duration
  , tag
  , isEvent
  , delayEvent
  , sample
  , rollup
  , sliding
  , fixed
  , debounce
  , currentTime
  , onEvent
  ) where

import Control.Category (Category(..), (>>>))
import Prelude hiding ((.), id)
import Control.Arrow (Arrow(..), ArrowChoice(..), ArrowLoop(..))
import Data.Time (UTCTime, NominalDiffTime, diffUTCTime, addUTCTime, TimeZone, LocalTime, utcToLocalTime)
import Control.Monad.Fix (MonadFix (mfix))
import Data.Either (fromLeft)
import Data.Bool (bool)
import Data.Monoid (Endo(..))
import Data.UUID (UUID)
import qualified Data.Set as S
import qualified Data.Text as T

data Request = Request
  { requestTime :: !UTCTime
  , requestTimeZone :: !TimeZone
  , requestTraceId :: !UUID
  } deriving (Show, Eq)

-- | The set of entity ids an arrow subscribes to. Static: it does not
-- change as the machine steps, so the runtime can read it once to build
-- trigger subscriptions.
data Mealy eff a b = Mealy
  { entities :: S.Set T.Text
  , runMealy :: forall m. MonadFix m => (forall x. eff x -> m x) -> Request -> a -> m (b, Mealy eff a b)
  }


instance Semigroup b => Semigroup (Mealy eff a b) where
  Mealy ast af <> Mealy bst bf = Mealy (ast <> bst) $ \nt r a -> do
    (x, af') <- af nt r a
    (x', bf') <- bf nt r a
    pure (x <> x', af' <> bf')


instance Monoid b => Monoid (Mealy eff a b) where
  mempty = Mealy mempty $ \_ _ _ -> pure (mempty, mempty)

eff :: (Request -> a -> eff b) -> Mealy eff a b
eff f = Mealy mempty $ \nt req x ->
  nt (f req x) >>= \b -> pure (b, eff f)

-- | Override the static entity set of an arrow. Use when a combinator
-- (e.g. 'switch') hides continuation entities from the runtime's
-- startup subscription scan.
withEntities :: S.Set T.Text -> Mealy eff a b -> Mealy eff a b
withEntities es (Mealy _ f) = Mealy es f

instance Category (Mealy eff) where
  id = Mealy mempty (\_ _ x -> pure (x, id))
  (Mealy ast f) . (Mealy bst g) = Mealy (ast <> bst) $ \nt t a -> do
    (b, g') <- g nt t a
    (c, f') <- f nt t b
    pure (c, f' . g')

instance Arrow (Mealy eff) where
  arr f = Mealy mempty $ \_ _ b -> pure (f b, arr f)
  first (Mealy st f) = Mealy st $ \nt t (b,d) -> do
    (c, f') <- f nt t b
    pure ((c, d), first f')

instance ArrowChoice (Mealy eff) where
  left m@(Mealy st f) = Mealy st $ \nt t -> \case
    Left b -> do
      (c, f') <- f nt t b
      pure (Left c, left f')
    Right d -> pure (Right d, left m)

instance ArrowLoop (Mealy eff) where
  loop (Mealy st f) = Mealy st $ \nt t b -> do
    ((c,_), f') <- mfix $ \((_,d), _) -> f nt t (b,d)
    pure (c, loop f')

instance Functor (Mealy eff a) where
  fmap f (Mealy st g) = Mealy st $ \nt t a -> do
    (b, g') <- g nt t a
    pure (f b, fmap f g')

instance Applicative (Mealy eff a) where
  pure b = Mealy mempty $ \_ _ _ -> pure (b, pure b)
  Mealy ast f <*> Mealy bst x = Mealy (ast <> bst) $ \nt t a -> do
    (f', fNext) <- f nt t a
    (x', xNext) <- x nt t a
    pure (f' x', fNext <*> xNext)

data Event a
  = Tick
  | Event a
  deriving (Show, Eq, Functor, Foldable, Traversable)

instance Semigroup (Event a) where
  (<>) = lMerge

instance Monoid (Event a) where
  mempty = Tick

hold :: a -> Mealy eff (Event a) a
hold a = Mealy mempty $ \_ _ -> \case
  Tick -> pure (a, hold a)
  Event a' -> pure (a', hold a')

events :: Mealy eff (Event a) (Either () a)
events = arr $ \case
  Tick -> Left ()
  Event a -> Right a

isEvent :: Event a -> Bool
isEvent Tick = False
isEvent _ = True

tag :: b -> Event a -> Event b
tag b ev = b <$ ev

switch :: Mealy eff a (b, Event c) -> (c -> Mealy eff a b) -> Mealy eff a b
switch (Mealy st f) s = Mealy st $ \nt t a -> do
  ((b, ev), f') <- f nt t a
  case ev of
          Tick -> pure (b, switch f' s)
          Event x -> runMealy (s x) nt t a

sample :: Mealy eff (a, Event b) (Event a)
sample = arr (uncurry tag)

preMapAccum :: (x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
preMapAccum f x extract = go x
  where
    go b = Mealy mempty $ \_ _ a ->
      let next = f b a
      in pure (extract b, go next)

preMapAccumRequest :: (Request -> x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
preMapAccumRequest f x extract = go x
  where
    go b = Mealy mempty $ \_ t a ->
      let next = f t b a
      in pure (extract b, go next)

mapAccum :: (x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
mapAccum f x extract = go x
  where
    go b = Mealy mempty $ \_ _ a ->
      let next = f b a
      in pure (extract next, go next)

mapAccumRequest :: (Request -> x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
mapAccumRequest f x extract = go x
  where
    go b = Mealy mempty $ \_ t a ->
      let next = f t b a
      in pure (extract next, go next)

data DelayState x a = DelayState
  { pending :: x
  , output  :: !(Event a)
  }


delayEvent :: NominalDiffTime -> Mealy eff (Event a) (Event a)
delayEvent delay =
  mapAccumRequest step initial output
  where
    initial = DelayState [] Tick

    step req st input =
      let now = requestTime req

          queued =
            case input of
              Tick    -> pending st
              Event x -> pending st ++ [(delay `addUTCTime` now, x)]

      in case queued of
           (due, x) : rest
             | due <= now ->
                 DelayState rest (Event x)

           _ ->
                 DelayState queued Tick

debounce :: NominalDiffTime -> Mealy eff (Event a) (Event a)
debounce delay =
  mapAccumRequest step initial output
  where
    initial = DelayState Nothing Tick
    step req st input =
      let now = requestTime req
          held = case input of
                      Tick -> pending st
                      Event x -> Just (delay `addUTCTime` now, x)
      in case held of
              Just (due, x)
                | due <= now -> DelayState Nothing (Event x)
              _ -> DelayState held Tick

changes :: Eq a => Mealy eff a (Event a)
changes = mapAccum go Nothing (maybe Tick snd)
  where
    go :: Eq a => Maybe (a, Event a) -> a -> Maybe (a, Event a)
    -- The first observed value is not a change I think
    go Nothing x = Just (x, Tick)
    go (Just (y, _)) x | x == y = Just (x, Tick)
                       | otherwise = Just (x, Event x)

whenA :: (a -> Bool) -> Mealy eff a () -> Mealy eff a ()
whenA predicate auto = arr (\a -> if predicate a then Left a else Right ()) >>> left auto >>> arr (fromLeft ())

filterA :: (a -> Bool) -> Mealy eff a (Either () a)
filterA f = arr $ \a -> bool (Left ()) (Right a) (f a)

thenA :: (ArrowChoice cat, Arrow cat) => cat a (Either b1 c) -> cat c (Either b1 b2) -> cat a (Either b1 b2)
thenA f g = f >>> arr Left ||| g

(>>|) :: (ArrowChoice cat, Arrow cat) => cat a (Either b1 c) -> cat c (Either b1 b2) -> cat a (Either b1 b2)
(>>|) = thenA

infixl 2 >>|

toEvent :: Mealy eff (Either () a) (Event a)
toEvent = arr (either (const Tick) Event)


lMerge :: Event a -> Event a -> Event a
lMerge Tick Tick = Tick
lMerge (Event a) _ = Event a
lMerge Tick (Event a) = Event a

edge :: Mealy eff Bool (Event ())
edge = go False
  where
    go True = Mealy mempty $ \_ _ -> \case
      True -> pure (Tick, go True)
      False -> pure (Tick, go False)
    go False = Mealy mempty $ \_ _ -> \case
      True -> pure (Event (), go True)
      False -> pure (Tick, go False)


-- | Drop the first 'Event' and pass through everything after. Useful for
-- ignoring a self-triggered event (e.g. a service call that changes the
-- very entity the arrow listens to).
dropFirst :: Mealy eff (Event a) (Event a)
dropFirst = go False
  where
    go seen = Mealy mempty $ \_ _ input ->
      case input of
        Event _ | not seen -> pure (Tick, go True)
        _ -> pure (input, go seen)


duration :: forall eff a. Mealy eff a NominalDiffTime
duration = mapAccumRequest go (Nothing @(UTCTime, NominalDiffTime)) (maybe 0 snd)
  where
    go :: Request -> Maybe (UTCTime, NominalDiffTime) -> a -> Maybe (UTCTime, NominalDiffTime)
    go req Nothing _ = Just (requestTime req, requestTime req `diffUTCTime` requestTime req)
    go req (Just (startTime, _)) _ = Just (startTime, requestTime req `diffUTCTime` startTime)



-- | Rollup, hold back bursty messages
--
-- Consider a case where you have a  bursty set of data. You care to get an immediate response,
-- but don't want to spam the output
rollup
  :: Int -- ^ How many items to pass through before burst protection
  -> Int -- ` How many seconds to collect the bursty data
  -> Mealy eff (Event a) (Event [a])
rollup limit seconds = mapAccumRequest go (Left Tick) (either id (\(_, _, _, ev) -> ev))
  where
    e a = Endo ([a] ++)
    go :: Request -> Either (Event [a]) (UTCTime, Int, Endo [a], Event [a]) -> Event a -> Either (Event [a]) (UTCTime, Int, Endo [a], Event [a])
    go _ (Left _) Tick = Left Tick
    go req (Left _) (Event a) = Right (addUTCTime (fromIntegral seconds) (requestTime req), 1, mempty, Event [a])
    go req (Right (end, n, acc, _)) Tick
      | requestTime req >= end = Left (Event $ appEndo acc [])
      | otherwise = Right (end, n, acc, Tick)
    go req (Right (end, n, acc, _)) (Event a)
      | requestTime req >= end = Left (Event $ appEndo acc [a])
      | n < limit = Right (end, n+1, acc, Event [a])
      | otherwise = Right (end, n+1, acc <> e a, Tick)

-- Sliding window into the events
sliding :: Int -> Mealy eff (Event a) [a]
sliding size = mapAccum go [] id
  where
    go :: [a] -> Event a -> [a]
    go acc Tick = acc
    go acc (Event a) = let xs = acc ++ [a] in drop (max 0 (length xs - size)) xs

fixed :: Int -> Mealy eff (Event a) [a]
fixed seconds = mapAccumRequest go Nothing (maybe [] ((`appEndo` []) . snd))
  where
    e a = Endo ([a] ++)
    go :: Request -> Maybe (UTCTime, Endo [a]) -> Event a -> Maybe (UTCTime, Endo [a])
    go req Nothing Tick = Just (addUTCTime (fromIntegral seconds) (requestTime req), mempty)
    go req Nothing (Event a) = Just (addUTCTime (fromIntegral seconds) (requestTime req), e a)
    go req (Just (end, acc)) ev =
      case ev of
           Tick | requestTime req >= end -> Just (addUTCTime (fromIntegral seconds) end, mempty)
                | otherwise -> Just (end, acc)
           Event a | requestTime req >= end -> Just (addUTCTime (fromIntegral seconds) end, e a)
                   | otherwise -> Just (end, acc <> e a)


currentTime :: Mealy eff a LocalTime
currentTime = Mealy mempty $ \_ Request{requestTime, requestTimeZone} _ ->
  pure (utcToLocalTime requestTimeZone requestTime, currentTime)


onEvent :: Mealy eff a () -> Mealy eff (Event a) ()
onEvent f = events >>> (arr (const ()) ||| f)