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

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

import Control.Category (Category(..), (>>>))
import Prelude hiding ((.), id)
import Control.Arrow (Arrow(..), ArrowChoice(..))
import Data.Time (UTCTime (UTCTime), NominalDiffTime, diffUTCTime, addUTCTime, TimeZone, LocalTime, utcToLocalTime, Day (..), diffTimeToPicoseconds, picosecondsToDiffTime)
import Data.Either (fromLeft)
import Data.Bool (bool)
import Data.UUID (UUID)
import qualified Data.Set as S
import qualified Data.Text as T
import Data.Serialize (Get, Putter, runPut, Serialize (put), runGet, get)
import qualified Data.ByteString as B
import Control.Exception (IOException, handle, throwIO)
import System.IO.Error (isDoesNotExistError)
import GHC.Generics (Generic)
import Data.Sequence (Seq, (|>))
import qualified Data.Foldable as F
import Control.Monad.IO.Class (MonadIO, liftIO)


data Codec s = Codec { getter :: !(Get s), putter :: !(Putter s) }

data State s = State {state :: !s, dirty :: !Bool}
             deriving Functor


instance Semigroup s => Semigroup (State s) where
  s1 <> s2 = State (state s1 <> state s2) (dirty s1 || dirty s2)

instance Monoid s => Monoid (State s) where
  mempty = State mempty False

instance Applicative State where
  pure a = State a False
  s1 <*> s2 = State
      { state =
          let a = state s2
              f = state s1
          in f a
      , dirty = dirty s1 || dirty s2
      }

mergeState :: State s1 -> State s2 -> State (s1, s2)
mergeState s1 s2 = (,) <$> s1 <*> s2


mergeCodec :: Codec s -> Codec s1 -> Codec (s, s1)
mergeCodec (Codec agetter aputter) (Codec bgetter bputter) = Codec (mergeGet agetter bgetter) (mergePut aputter bputter)
  where
    mergePut :: Putter s -> Putter s1 -> Putter (s, s1)
    mergePut p1 p2 (s, s1) = p1 s >> p2 s1
    mergeGet :: Get s -> Get s' -> Get (s, s')
    mergeGet g1 g2 = (,) <$> g1 <*> g2

data Pair a b = Pair !a !b

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


data Auto m a b
  = Fun (Request -> a -> b) -- Stateless variant, needed at least for 'id'
  | forall s. Stateful !(Codec s) !(State s) !(State s -> Request -> a -> m (b, State s))  -- State is explicitly part of it


instance Monad m => Functor (Auto m a) where
  fmap f = \case
    Fun x -> Fun $ \req -> f . x req
    Stateful codec s x -> Stateful codec s $ \s' req a -> do
      (a',s'') <- x s' req a
      pure (f a', s'')

instance Monad m => Applicative (Auto m a) where
  pure a = Fun (\_req -> const a)
  fa <*> fb =
    case (fa,fb) of
         (Fun af, Fun bf) -> Fun $ \req -> (af req <*> bf req)
         (Stateful codec s af, Fun bf) -> Stateful codec s
          (\s' req x -> do
           let a = bf req x
           (h, s'') <- af s' req x
           pure (h a, s'')
          )
         (Fun af, Stateful codec s bf) -> Stateful codec s
          (\s' req x -> do
           (a, s'') <- bf s' req x
           let h = af req x
           pure (h a, s'')
          )
         (Stateful acodec as af, Stateful bcodec bs bf) -> Stateful (mergeCodec acodec bcodec) (mergeState as bs)
          (\s' req x -> do
           (a, as') <- bf (snd <$> s') req x
           (h, bs') <- af (fst <$> s') req x
           pure (h a, mergeState bs' as')
          )


instance (Monad m, Semigroup b) => Semigroup (Auto m a b) where
  fa <> fb =
    case (fa,fb) of
         (Fun af, Fun bf) -> Fun (af <> bf)
         (Stateful codec s af, Fun bf) -> Stateful codec s
           (\s' req a -> do
            (ab, s'') <- af s' req a
            let bb = bf req a
            pure (ab <> bb, s'')
           )
         (Fun af, Stateful codec s bf) -> Stateful codec s
           (\s' req a -> do
            let ab = af req a
            (bb, s'') <- bf s' req a
            pure (ab <> bb, s'')
           )
         (Stateful acodec as af , Stateful bcodec bs bf) -> Stateful (mergeCodec acodec bcodec) (mergeState as bs)
           (\s req a -> do
            (ab, as'') <- af (fst <$> s) req a
            (bb, bs'') <- bf (snd <$> s) req a
            pure (ab <> bb, mergeState as'' bs'')
           )


instance (Monad m, Monoid b) => Monoid (Auto m a b) where
  mempty = Fun $ \_req _ -> mempty


instance Monad m => Category (Auto m) where
  id = Fun $ \_ -> id
  af . ag =
    case (af, ag) of
         (Fun f, Fun g) -> Fun (\req -> f req . g req)
         (Stateful codec s f, Fun g) -> Stateful codec s (\s' req -> f s' req . g req)
         (Fun f, Stateful codec s g) -> Stateful codec s (\s' req -> fmap (first (f req)) . g s' req)
         (Stateful fcodec fs f , Stateful gcodec gs g) ->
           Stateful (mergeCodec fcodec gcodec) (mergeState fs gs) (\s req a -> do
                       (b, s') <- g (snd <$> s) req a
                       (c, s'') <- f (fst <$> s) req b
                       pure (c, mergeState s'' s'))



instance Monad m => Arrow (Auto m) where
  arr f = Fun $ const f
  first = \case
    Fun f -> Fun $ \req -> first (f req)
    Stateful codec s f -> Stateful codec s $ \s' req (b,d) -> do
      (c, s'') <- f s' req b
      pure ((c,d), s'')

instance Monad m => ArrowChoice (Auto m) where
  left = \case
    Fun f -> Fun $ \req ->
      \case
        Left b -> Left $ f req b
        Right d -> Right d
    Stateful codec s f -> Stateful codec s $ \s' req -> \case
      Right d -> pure (Right d, s')
      Left b -> do
        (c, s'') <- f s' req b
        pure (Left c, s'')

serialize :: Auto m a b -> B.ByteString
serialize = \case
  Fun _ -> runPut $ put ()
  Stateful Codec{putter} s _ -> runPut $ putter (state s)

data DecodedAuto m a b
  = Decoded (Auto m a b) -- decoded from serialized state
  | FailDecode String (Auto m a b) -- gives back the original + errmsg

deserialize :: B.ByteString -> Auto m a b -> DecodedAuto m a b
deserialize bs = \case
  Fun f -> Decoded (Fun f) -- no state to decode, success by default
  Stateful codec s f ->
    either
      (\err -> FailDecode err (Stateful codec s f))
      (\s' -> Decoded $ Stateful codec (pure s') f)
      $ runGet (getter codec) bs

save :: FilePath -> Auto m a b -> IO (Auto m a b)
save path s
  | isDirty s = do
    _ <- B.writeFile path $ serialize s
    pure $ cleanDirty s
  | otherwise = pure s
  where
    cleanDirty :: Auto m a b -> Auto m a b
    cleanDirty (Stateful codec s' f) = Stateful codec s'{dirty=False} f
    cleanDirty a = a
    isDirty :: Auto m a b -> Bool
    isDirty (Stateful _ s' _) = dirty s'
    isDirty _ = False


load :: forall m a b. FilePath -> Auto m a b -> IO (DecodedAuto m a b)
load path a = handle defaultOnMissingFile (flip deserialize a <$> B.readFile path)
  where
    defaultOnMissingFile :: IOException -> IO (DecodedAuto m a b)
    defaultOnMissingFile e
      | isDoesNotExistError e = pure $ FailDecode "State doesn't exist eyt" a
      | otherwise = throwIO e

-- | 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. Monad m => (forall x. eff x -> m x) -> Auto m a b
  }


instance Semigroup b => Semigroup (Mealy eff a b) where
  Mealy ast af <> Mealy bst bf = Mealy (ast <> bst) $ \nt -> af nt <> bf nt


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

eff :: (Request -> a -> eff b) -> Mealy eff a b
eff f = Mealy mempty $ \nt -> do
  Stateful (Codec get put) (State () False) $ \s req a -> do
    b <- nt (f req a)
    pure (b, s)

-- | 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 (\_ -> id)
  (Mealy ast f) . (Mealy bst g) = Mealy (ast <> bst) $ \nt -> do
    f nt . g nt

instance Arrow (Mealy eff) where
  arr f = Mealy mempty $ \_nt -> arr f
  first (Mealy st f) = Mealy st $ \nt -> first (f nt)

instance ArrowChoice (Mealy eff) where
  left (Mealy st f) = Mealy st $ \nt -> left (f nt)


instance Functor (Mealy eff a) where
  fmap f (Mealy st g) = Mealy st $ \nt -> fmap f (g nt)

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

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

instance Serialize a => Serialize (Event a)

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

instance Monoid (Event a) where
  mempty = Tick

hold :: (Serialize a, Eq a) => a -> Mealy m (Event a) a
hold def = preMapAccum step def id
  where
    step prev = \case
      Tick -> prev
      Event new -> new


-- hold :: a -> Mealy eff (Event a) a
-- hold a = Mealy mempty $ \_ _ -> \case
--   Tick -> pure (Pair a (hold a))
--   Event a' -> pure (Pair 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
--   Pair (b, ev) f' <- f nt t a
--   case ev of
--           Tick -> pure (Pair 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' :: forall m x a b. (Monad m, Eq x, Serialize x) => (x -> a -> x) -> x -> (x -> b) -> Auto m a b
preMapAccum' f x extract = Stateful (Codec get put) (State x False) (\s _req a -> pure $ step s a)
  where
    step :: State x -> a -> (b, State x)
    step s a = let s' = f (state s) a in (extract (state s), State s' (dirty s || state s /= s'))


mapAccum :: (Eq x, Serialize x) => (x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
mapAccum step x extract = Mealy mempty $ \_nt -> mapAccum' step x extract

preMapAccum :: (Eq x, Serialize x) => (x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
preMapAccum step x extract = Mealy mempty $ \_nt -> preMapAccum' step x extract

mapAccum' :: forall m x a b. (Monad m, Eq x, Serialize x) => (x -> a -> x) -> x -> (x -> b) -> Auto m a b
mapAccum' f x extract = Stateful (Codec get put) (State x False) (\s _req a -> pure $ step s a)
  where
    step :: State x -> a -> (b, State x)
    step s a = let s' = f (state s) a in (extract s', State s' (dirty s || state s /= s'))

preMapAccumRequest :: (Serialize x, Eq x) => (Request -> x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
preMapAccumRequest step x extract = Mealy mempty $ \_ -> preMapAccumRequest' step x extract

preMapAccumRequest' :: forall m x a b. (Serialize x, Eq x, Monad m) => (Request -> x -> a -> x) -> x -> (x -> b) -> Auto m a b
preMapAccumRequest' f x extract = Stateful (Codec get put) (State x False) (\s req a -> pure $ step s req a)
  where
    step :: State x -> Request -> a -> (b, State x)
    step s req a = let s' = f req (state s) a in (extract (state s), State s' (dirty s || state s /= s'))

mapAccumRequest :: (Serialize x, Eq x) => (Request -> x -> a -> x) -> x -> (x -> b) -> Mealy eff a b
mapAccumRequest step x extract = Mealy mempty $ \_ -> mapAccumRequest' step x extract

mapAccumRequest' :: forall m x a b. (Monad m, Serialize x, Eq x) => (Request -> x -> a -> x) -> x -> (x -> b) -> Auto m a b
mapAccumRequest' f x extract = Stateful (Codec get put) (State x False) (\s req a -> pure $ step s req a)
  where
    step :: State x -> Request -> a -> (b, State x)
    step s req a = let s' = f req (state s) a in (extract s', State s' (dirty s || state s /= s'))

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

instance (Serialize x, Serialize a) => Serialize (DelayState x a)

newtype SerializeUTCTime = SerializeUTCTime UTCTime
  deriving Eq

-- TODO: Needs '\x -> pure x == get (put x)' test
instance Serialize SerializeUTCTime where
  put (SerializeUTCTime (UTCTime day time)) = do
    put (toModifiedJulianDay day)
    put (diffTimeToPicoseconds time)

  get = do
    day  <- ModifiedJulianDay <$> get
    time <- picosecondsToDiffTime <$> get
    pure $ SerializeUTCTime (UTCTime day time)

delayEvent :: (Eq a, Serialize a) => 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 ++ [(SerializeUTCTime $ delay `addUTCTime` now, x)]

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

           _ ->
                 DelayState queued Tick

debounce :: (Serialize a, Eq a) => 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 (SerializeUTCTime $ delay `addUTCTime` now, x)
      in case held of
              Just (SerializeUTCTime due, x)
                | due <= now -> DelayState Nothing (Event x)
              _ -> DelayState held Tick

changes :: (Serialize a, 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 =
  mapAccum
    (\(_, current) new -> (current, new))
    (False, False)
    (\(old, current) ->
      if not old && current
        then Event ()
        else Tick)





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


-- | 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
  :: (Serialize a, Eq a)
  => 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
    go
      :: Request
      -> Either (Event [a]) (SerializeUTCTime, Int, Seq a, Event [a])
      -> Event a
      -> Either (Event [a]) (SerializeUTCTime, Int, Seq a, Event [a])

    go _ (Left _) Tick =
      Left Tick

    go req (Left _) (Event a) =
      Right
        ( SerializeUTCTime $ addUTCTime (fromIntegral seconds) (requestTime req)
        , 1
        , mempty
        , Event [a]
        )

    go req (Right (SerializeUTCTime end, n, acc, _)) Tick
      | requestTime req >= end =
          Left (Event $ F.toList acc)
      | otherwise =
          Right (SerializeUTCTime end, n, acc, Tick)

    go req (Right (SerializeUTCTime end, n, acc, _)) (Event a)
      | requestTime req >= end =
          Left (Event $ F.toList (acc |> a))
      | n < limit =
          Right (SerializeUTCTime end, n + 1, acc, Event [a])
      | otherwise =
          Right (SerializeUTCTime end, n + 1, acc |> a, Tick)


-- Sliding window into the events
sliding :: (Serialize a, Eq a) => 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



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


stepAuto :: Monad m => Auto m a b -> Request -> a -> m (b, Auto m a b)
stepAuto (Fun f) req a = pure (f req a, Fun f)
stepAuto (Stateful codec s f) req a = do
  (b, s') <- f s req a
  pure (b, Stateful codec s' f)


stepAutoSerializing :: MonadIO m => FilePath -> Auto m a b -> Request -> a -> m (b, Auto m a b)
stepAutoSerializing path f req a = do
  (b,x) <- stepAuto f req a
  y <- liftIO $ save path x
  pure (b,y)

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