{-# LANGUAGE OverloadedStrings #-}

module AFRPSpec (spec) where

import Control.Arrow (arr, (&&&), first, left)
import Control.Category ((>>>))
import Control.Monad.Fix (MonadFix (..))
import Data.Foldable (for_)
import AFRP
import Data.Functor.Identity (Identity (..))
import Data.List (sort)
import Data.Serialize (get, put, runGet, runPut)
import qualified Data.Set as S
import qualified Data.Text as T
import Data.Time (Day (..), NominalDiffTime, LocalTime(..), TimeOfDay(..), UTCTime (..), picosecondsToDiffTime, utc)
import Data.UUID (nil)
import Hedgehog
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Test.Hspec
import Test.Hspec.Hedgehog

fakeRequest :: Request
fakeRequest = Request (sec 0) utc nil

sec :: Integer -> UTCTime
sec n = UTCTime (toEnum 0) (fromIntegral n)

untime :: SerializeUTCTime -> UTCTime
untime (SerializeUTCTime t) = t

runPure :: Mealy Identity a b -> [a] -> [b]
runPure m = go (runMealy m id)
  where
    go _ [] = []
    go w (a : as) = case runIdentity (stepAuto w fakeRequest a) of
      (b, w') -> b : go w' as

-- | Run a Mealy with a per-step wall clock (seconds since the day-0 epoch).
runTimed :: Mealy Identity a b -> [(Integer, a)] -> [b]
runTimed m = go (runMealy m id)
  where
    go _ [] = []
    go w ((s, a) : as) =
      case runIdentity (stepAuto w (Request (sec s) utc nil) a) of
        (b, w') -> b : go w' as

-- | A minimal State monad for observing effectful arrows (e.g. whenA gating).
newtype St a = St { unSt :: Int -> (a, Int) }

instance Functor St where
  fmap f (St g) = St $ \s -> let (a, s') = g s in (f a, s')

instance Applicative St where
  pure a = St (\s -> (a, s))
  St f <*> St x = St $ \s -> let (f', s') = f s; (a, s'') = x s' in (f' a, s'')

instance Monad St where
  St m >>= k = St $ \s -> let (a, s') = m s; (b, s'') = unSt (k a) s' in (b, s'')

instance MonadFix St where
  mfix f = St $ \s -> let (a, s') = unSt (f a) s in (a, s')

runStEff :: Mealy St a b -> Int -> [a] -> ([b], Int)
runStEff m = go (runMealy m id)
  where
    go _ s [] = ([], s)
    go w s (a : rest) = case unSt (stepAuto w fakeRequest a) s of
      ((b, w'), s') -> let (bs, s'') = go w' s' rest in (b : bs, s'')

spec :: Spec
spec = describe "AFRP" $ do
  entitiesSpec
  holdSpec
  eventsSpec
  isEventSpec
  tagSpec
  toEventSpec
  lMergeSpec
  changesSpec
  edgeSpec
  filterASpec
  slidingSpec
  mapAccumSpec
  preMapAccumSpec
  durationSpec
  delayEventSpec
  debounceSpec
  rollupSpec
  serializeSpec
  effSpec
  mapAccumRequestSpec
  preMapAccumRequestSpec
  whenASpec
  thenASpec
  sampleSpec

holdSpec :: Spec
holdSpec = describe "hold" $ do
  it "holds initial value until an Event arrives" $
    runPure (hold 'a') [Tick, Event 'b', Tick, Event 'c']
      `shouldBe` ['a', 'b', 'b', 'c']

  it "never changes on Tick" $
    runPure (hold (0 :: Int)) (replicate 5 Tick) `shouldBe` replicate 5 (0 :: Int)

eventsSpec :: Spec
eventsSpec = describe "events" $ do
  it "converts Tick to Left () and Event a to Right a" $
    runPure events [Tick, Event 'a', Tick, Event 'b']
      `shouldBe` [Left (), Right 'a', Left (), Right 'b']

isEventSpec :: Spec
isEventSpec = describe "isEvent" $ do
  it "returns False for Tick" $
    isEvent Tick `shouldBe` False

  it "returns True for Event x" $
    isEvent (Event ()) `shouldBe` True

tagSpec :: Spec
tagSpec = describe "tag" $ do
  it "replaces value preserving structure" $ do
    tag 'b' Tick `shouldBe` Tick
    tag 'b' (Event 'a') `shouldBe` Event 'b'

toEventSpec :: Spec
toEventSpec = describe "toEvent" $ do
  it "round-trips through events" $
    runPure toEvent [Left (), Right 'a', Left ()]
      `shouldBe` [Tick, Event 'a', Tick]

  it "is inverse of events modulo Event/Either" $
    runPure (events >>> toEvent) [Tick, Event 'a', Event 'b']
      `shouldBe` [Tick, Event 'a', Event 'b']

lMergeSpec :: Spec
lMergeSpec = describe "lMerge" $ do
  it "both Tick gives Tick" $
    lMerge (Tick :: Event Int) (Tick :: Event Int) `shouldBe` Tick

  it "prefers left Event" $
    lMerge (Event (1 :: Int)) (Event (2 :: Int)) `shouldBe` Event (1 :: Int)

  it "prefers right Event if left is Tick" $
    lMerge Tick (Event (2 :: Int)) `shouldBe` Event (2 :: Int)

  it "Tick is a left identity" $
    hedgehog $ do
      e <- forAll eventGen
      lMerge Tick e === e

  it "Tick is a right identity" $
    hedgehog $ do
      e <- forAll eventGen
      lMerge e Tick === e

  it "is associative" $
    hedgehog $ do
      a <- forAll eventGen
      b <- forAll eventGen
      c <- forAll eventGen
      lMerge a (lMerge b c) === lMerge (lMerge a b) c

changesSpec :: Spec
changesSpec = describe "changes" $ do
  it "first output is always Tick" $
    runPure changes "hello" !! 0 `shouldBe` Tick

  it "outputs Event only on value change" $
    runPure changes "aaaabbbcca"
      `shouldBe` [Tick, Tick, Tick, Tick
                 , Event 'b', Tick, Tick
                 , Event 'c', Tick
                 , Event 'a'
                 ]

  it "first output is Tick, subsequent outputs are Event iff value changed" $
    hedgehog $ do
      xs <- forAll $ Gen.list (Range.linear 0 100) Gen.alpha
      let out = runPure changes xs
      length out === length xs
      case out of
        [] -> pure ()
        (Tick : rest) -> do
          let triples = zip3 xs (drop 1 xs) rest
          for_ triples $ \(prev, curr, o) ->
            if prev /= curr
              then o === Event curr
              else o === Tick
        _ -> failure

edgeSpec :: Spec
edgeSpec = describe "edge" $ do
  it "emits Event () only on rising edge" $
    runPure edge [False, True, True, False, True]
      `shouldBe` [Tick, Event (), Tick, Tick, Event ()]

  it "starts from False, so first True is a rising edge" $
    runPure edge [True, False, True]
      `shouldBe` [Event (), Tick, Event ()]

  it "Event () only on False -> True transition" $
    hedgehog $ do
      bs <- forAll $ Gen.list (Range.linear 0 50) Gen.bool
      let out = runPure edge bs
      length out === length bs
      case (bs, out) of
        ([], []) -> pure ()
        (b : _, o : _) -> do
          if b then o === Event () else o === Tick
          let triples = zip3 bs (drop 1 bs) (drop 1 out)
          for_ triples $ \(prev, curr, o') ->
            if not prev && curr
              then o' === Event ()
              else o' === Tick
        _ -> failure


filterASpec :: Spec
filterASpec = describe "filterA" $ do
  it "lets through values matching predicate" $
    runPure (filterA (even @Int)) [1, 2, 3, 4]
      `shouldBe` [Left (), Right 2, Left (), Right 4]

  it "output is Right a iff predicate holds" $
    hedgehog $ do
      threshold <- forAll $ Gen.int (Range.linear (-10) 10)
      let p = (> threshold)
      xs <- forAll $ Gen.list (Range.linear 0 50) (Gen.int (Range.linear (-20) 20))
      let out = runPure (filterA p) xs
      length out === length xs
      for_ (zip xs out) $ \(x, o) ->
        if p x
          then o === Right x
          else o === Left ()

slidingSpec :: Spec
slidingSpec = describe "sliding" $ do
  it "accumulates events up to the window size" $
    runPure (sliding (3 :: Int)) [Event (1 :: Int), Event 2, Event 3, Event 4]
      `shouldBe` [[1], [1, 2], [1, 2, 3], [2, 3, 4]]

  it "Ticks don't change the accumulator" $
    runPure (sliding 2) [Event (1 :: Int), Tick, Event 2]
      `shouldBe` [[1], [1], [1, 2]]

  it "empty list stays empty" $
    runPure (sliding (5 :: Int)) ([] :: [Event Int]) `shouldBe` []

  it "output length never exceeds window size" $
    hedgehog $ do
      n <- forAll $ Gen.int (Range.constant 1 10)
      evs <- forAll $ Gen.list (Range.linear 0 20) (Gen.frequency
        [(3, Event <$> Gen.alpha), (1, pure Tick)])
      let out = runPure (sliding n) evs
      for_ out $ \xs -> assert (length xs <= n)

mapAccumSpec :: Spec
mapAccumSpec = describe "mapAccum" $ do
  it "running sum" $
    runPure (mapAccum (+) (0 :: Int) id) [1, 2, 3]
      `shouldBe` [1, 3, 6]

  it "post-state extraction: output uses state after applying f" $
    runPure (mapAccum (\s x -> s ++ [x]) ([] :: [Int]) id) [1, 2, 3]
      `shouldBe` [[1], [1, 2], [1, 2, 3]]

  it "output equals running sum of all inputs so far" $
    hedgehog $ do
      xs <- forAll $ Gen.list (Range.linear 0 50) (Gen.int (Range.linear (-100) 100))
      let out = runPure (mapAccum (+) (0 :: Int) id) xs
      length out === length xs
      for_ (zip3 [0 ..] xs out) $ \(i, _x, cur) ->
        cur === sum (take (i + 1) xs)

preMapAccumSpec :: Spec
preMapAccumSpec = describe "preMapAccum" $ do
  it "running sum with pre-state extraction" $
    runPure (preMapAccum (+) (0 :: Int) id) [1, 2, 3]
      `shouldBe` [0, 1, 3]

  it "pre-state extraction: output uses state before applying f" $
    runPure (preMapAccum (\s x -> s ++ [x]) ([] :: [Int]) id) [1, 2, 3]
      `shouldBe` [[], [1], [1, 2]]

  it "output equals running sum before current input" $
    hedgehog $ do
      xs <- forAll $ Gen.list (Range.linear 0 50) (Gen.int (Range.linear (-100) 100))
      let out = runPure (preMapAccum (+) (0 :: Int) id) xs
      length out === length xs
      for_ (zip3 [0 ..] xs out) $ \(i, _x, cur) ->
        cur === sum (take i xs)

durationSpec :: Spec
durationSpec = describe "duration" $ do
  it "first sample is 0, then elapsed time since first sample" $
    runTimed duration [(0, 'a'), (5, 'b'), (10, 'c')]
      `shouldBe` [0, 5, 10]

  it "measures from the first observation, not the most recent" $
    runTimed duration [(2, 'a'), (3, 'b'), (7, 'c')]
      `shouldBe` [0, 1, 5]

  it "output i equals times[i] - times[0]" $
    hedgehog $ do
      secs <- forAll $ Gen.list (Range.linear 0 50) (Gen.int (Range.linear 0 1000))
      let out = runTimed duration [(fromIntegral s, ()) | s <- secs]
      length out === length secs
      case secs of
        [] -> pure ()
        (t0 : _) -> for_ (zip secs out) $ \(s, d) ->
          d === fromIntegral (s - t0)

delayEventSpec :: Spec
delayEventSpec = describe "delayEvent" $ do
  let delay = 5 :: NominalDiffTime

  it "emits a queued event once the delay has elapsed" $
    runTimed (delayEvent delay)
             [(0, Event 'a'), (3, Tick), (6, Tick)]
      `shouldBe` [Tick, Tick, Event 'a']

  it "preserves order when multiple events are queued" $
    runTimed (delayEvent delay)
             [(0, Event 'a'), (1, Event 'b'), (10, Tick), (12, Tick)]
      `shouldBe` [Tick, Tick, Event 'a', Event 'b']

  it "emits nothing on pure Tick input" $
    runTimed (delayEvent delay) [(0, Tick :: Event Char), (10, Tick)]
      `shouldBe` [Tick, Tick]

  it "emits exactly one output Event per input Event" $
    hedgehog $ do
      evs <- forAll $ Gen.list (Range.linear 0 30) eventGen
      let evs' = evs ++ replicate 3 Tick
          n = length [() | Event _ <- evs]
          out = runTimed (delayEvent (1 :: NominalDiffTime))
                         (zip [0, 2 ..] evs')
      length [() | Event _ <- out] === n

debounceSpec :: Spec
debounceSpec = describe "debounce" $ do
  let delay = 5 :: NominalDiffTime

  it "fires the last event after the quiet period" $
    runTimed (debounce delay)
             [(0, Event 'a'), (3, Tick), (6, Tick)]
      `shouldBe` [Tick, Tick, Event 'a']

  it "a newer event before firing resets the timer" $
    runTimed (debounce delay)
             [(0, Event 'a'), (3, Event 'b'), (6, Tick), (8, Tick)]
      `shouldBe` [Tick, Tick, Tick, Event 'b']

  it "collapses a burst into a single emission" $
    runTimed (debounce delay)
             [(0, Event 'a'), (1, Event 'b'), (2, Event 'c'), (10, Tick)]
      `shouldBe` [Tick, Tick, Tick, Event 'c']

  it "emits no more output Events than input Events" $
    hedgehog $ do
      evs <- forAll $ Gen.list (Range.linear 0 30) eventGen
      let nIn = length [() | Event _ <- evs]
          out = runTimed (debounce (1 :: NominalDiffTime))
                         (zip [0, 2 ..] evs)
          nOut = length [() | Event _ <- out]
      assert (nOut <= nIn)

rollupSpec :: Spec
rollupSpec = describe "rollup" $ do
  it "passes the first `limit` events through immediately, then bursts" $
    runTimed (rollup 2 10)
             [ (0, Event 'a'), (1, Event 'b')
             , (2, Event 'c'), (3, Event 'd')
             , (12, Tick)
             ]
      `shouldBe` [ Event ['a'], Event ['b']
                 , Tick, Tick
                 , Event ['c', 'd']
                 ]

  it "flushes the accumulator when the window ends on a Tick" $
    runTimed (rollup 1 10)
             [(0, Event 'a'), (1, Event 'b'), (2, Tick), (12, Tick)]
      `shouldBe` [Event ['a'], Tick, Tick, Event ['b']]

  it "is idle (Tick) until the first Event" $
    runTimed (rollup 2 10) [(0, Tick :: Event Char), (1, Tick)]
      `shouldBe` [Tick, Tick]

  it "every input Event appears exactly once across the outputs" $
    hedgehog $ do
      limit <- forAll $ Gen.int (Range.constant 1 5)
      evs <- forAll $ Gen.list (Range.linear 0 30) eventGen
      let evs' = evs ++ replicate 10 Tick
          times = [0 ..]
          out = runTimed (rollup limit 5) (zip times evs')
          emitted = concat [xs | Event xs <- out]
      sort emitted === sort [x | Event x <- evs]


eventGen :: Gen (Event Char)
eventGen = Gen.frequency
  [ (3, Event <$> Gen.alpha)
  , (1, pure Tick)
  ]

timeGen :: Gen SerializeUTCTime
timeGen = do
  day <- ModifiedJulianDay . fromIntegral <$> Gen.int (Range.linear 0 100000)
  pico <- picosecondsToDiffTime . fromIntegral
    <$> Gen.int (Range.linear 0 (86400 * 10 ^ (12 :: Int) - 1))
  pure $ SerializeUTCTime (UTCTime day pico)

localTimeGen :: Gen SerializeLocalTime
localTimeGen = do
  day <- ModifiedJulianDay . fromIntegral <$> genInt 0 100000
  tod <- TimeOfDay <$> genInt 0 23 <*> genInt 0 59 <*> (fromIntegral <$> genInt 0 60)
  pure $ SerializeLocalTime (LocalTime day tod)
  where
    genInt a b = Gen.int (Range.linear a b)

serializeSpec :: Spec
serializeSpec = do
  describe "SerializeUTCTime" $ do
    it "get (put x) == pure x" $
      hedgehog $ do
        x <- forAll timeGen
        tripping x (runPut . put) (runGet get)
  describe "SerializeLocalTime" $ do
    it "get (put x) == pure x" $
      hedgehog $ do
        x <- forAll localTimeGen
        tripping x (runPut . put) (runGet get)

effSpec :: Spec
effSpec = describe "eff" $ do
  it "lifts a pure effect function into a stateless Mealy" $
    runPure (eff (\_ x -> Identity (x + 1))) [1, 2, 3]
      `shouldBe` [2 :: Int, 3, 4]

  it "output equals f(input) for every step" $
    hedgehog $ do
      xs <- forAll $ Gen.list (Range.linear 0 50) (Gen.int (Range.linear (-100) 100))
      let out = runPure (eff (\_ x -> Identity (x * 2))) xs
      out === map (* 2) xs

mapAccumRequestSpec :: Spec
mapAccumRequestSpec = describe "mapAccumRequest" $ do
  it "accumulates request times, post-state extraction" $
    runTimed (mapAccumRequest (\req s _ -> s ++ [SerializeUTCTime (requestTime req)]) [] (map untime))
             [(0, 'a'), (5, 'b'), (10, 'c')]
      `shouldBe` [ [sec 0], [sec 0, sec 5], [sec 0, sec 5, sec 10] ]

  it "output i is every request time seen so far" $
    hedgehog $ do
      secs' <- forAll $ Gen.list (Range.linear 0 30) (Gen.int (Range.linear 0 100))
      let out = runTimed (mapAccumRequest (\req s _ -> s ++ [SerializeUTCTime (requestTime req)]) [] (map untime))
                         [(fromIntegral s, ()) | s <- secs']
          expected = [ map (sec . fromIntegral) (take (i + 1) secs') | i <- [0 .. length secs' - 1] ]
      out === expected

preMapAccumRequestSpec :: Spec
preMapAccumRequestSpec = describe "preMapAccumRequest" $ do
  it "accumulates request times, pre-state extraction" $
    runTimed (preMapAccumRequest (\req s _ -> s ++ [SerializeUTCTime (requestTime req)]) [] (map untime))
             [(0, 'a'), (5, 'b'), (10, 'c')]
      `shouldBe` [ [], [sec 0], [sec 0, sec 5] ]

  it "output i is every request time before the current step" $
    hedgehog $ do
      secs' <- forAll $ Gen.list (Range.linear 0 30) (Gen.int (Range.linear 0 100))
      let out = runTimed (preMapAccumRequest (\req s _ -> s ++ [SerializeUTCTime (requestTime req)]) [] (map untime))
                         [(fromIntegral s, ()) | s <- secs']
          expected = [ map (sec . fromIntegral) (take i secs') | i <- [0 .. length secs' - 1] ]
      out === expected

whenASpec :: Spec
whenASpec = describe "whenA" $ do
  let counter = eff (\_ (_ :: Int) -> St (\s -> ((), s + 1)))

  it "output is always () regardless of the predicate" $
    fst (runStEff (whenA (> 5) counter) 0 [1, 6, 2, 7])
      `shouldBe` [(), (), (), ()]

  it "runs the inner arrow only when the predicate holds" $
    snd (runStEff (whenA (> 5) counter) 0 [1, 6, 2, 7])
      `shouldBe` 2

  it "never runs the inner arrow when the predicate is always false" $
    snd (runStEff (whenA (const False) counter) 0 [1, 6, 2, 7])
      `shouldBe` 0

  it "runs the inner arrow on every input when the predicate is always true" $
    snd (runStEff (whenA (const True) counter) 0 [1, 6, 2, 7])
      `shouldBe` 4

thenASpec :: Spec
thenASpec = describe "thenA" $ do
  it "short-circuits on Left and continues on Right" $
    runPure (filterA (even @Int) `thenA` filterA (> 3)) [1..6]
      `shouldBe` [Left (), Left (), Left (), Right 4, Left (), Right 6]

  it "(>>|) is an infix alias for thenA" $
    runPure (filterA (even @Int) >>| filterA (> 3)) [1..6]
      `shouldBe` [Left (), Left (), Left (), Right 4, Left (), Right 6]

  it "second stage runs only when the first produces Right" $
    hedgehog $ do
      xs <- forAll $ Gen.list (Range.linear 0 30) (Gen.int (Range.linear (-20) 20))
      let out = runPure (filterA (even @Int) >>| filterA (> 0)) xs
          expected =
            [ if odd x then Left ()
              else if x > 0 then Right x
              else Left ()
            | x <- xs ]
      out === expected

sampleSpec :: Spec
sampleSpec = describe "sample" $ do
  it "tags the current value onto the Event structure" $
    runPure sample [(1, Tick), (2, Event 'a'), (3, Tick)]
      `shouldBe` [Tick, Event @Int 2, Tick]

  it "output is Event a iff the input event is present" $
    hedgehog $ do
      vals <- forAll $ Gen.list (Range.linear 0 30) (Gen.int (Range.linear 0 100))
      evs <- forAll $ Gen.list (Range.linear 0 30) eventGen
      let n = min (length vals) (length evs)
          ps = take n $ zip vals evs
          out = runPure sample ps
      for_ (zip ps out) $ \((v, ev), o) ->
        o === tag v ev

-- | A stateless arrow carrying a fixed entity set, for testing propagation.
subscribed :: S.Set T.Text -> Mealy Identity Int Int
subscribed ents = Mealy ents $ \_ -> Fun $ \_ a -> a

-- | Same as 'subscribed' but yields a function, for testing '<*>'.
subscribedF :: S.Set T.Text -> Mealy Identity Int (Int -> Int)
subscribedF ents = Mealy ents $ \_ -> Fun $ \_ a -> (a +)

entitiesSpec :: Spec
entitiesSpec = describe "entities" $ do
  it "id carries no entities" $
    entities (arr id :: Mealy Identity Int Int) `shouldBe` S.empty

  it "arr carries no entities" $
    entities (arr (+ 1) :: Mealy Identity Int Int) `shouldBe` S.empty

  it "eff carries no entities" $
    entities (eff (\_ x -> Identity (x + 1 :: Int))) `shouldBe` S.empty

  it "primitive combinators carry no entities" $ do
    entities (hold 'a') `shouldBe` S.empty
    entities (changes @Int) `shouldBe` S.empty
    entities edge `shouldBe` S.empty
    entities (sliding (3 :: Int) :: Mealy Identity (Event Int) [Int]) `shouldBe` S.empty

  it "Category (.) unions entity sets" $
    entities (subscribed (S.singleton "a") >>> subscribed (S.singleton "b"))
      `shouldBe` S.fromList ["a", "b"]

  it "Applicative (<*>) unions entity sets" $
    entities (subscribedF (S.singleton "a") <*> subscribed (S.singleton "b"))
      `shouldBe` S.fromList ["a", "b"]

  it "Arrow (&&&) unions entity sets" $
    entities (subscribed (S.singleton "a") &&& subscribed (S.singleton "b"))
      `shouldBe` S.fromList ["a", "b"]

  it "(>>>) unions entity sets" $
    entities (subscribed (S.singleton "a") >>> arr id >>> subscribed (S.singleton "b"))
      `shouldBe` S.fromList ["a", "b"]

  it "left preserves the entity set" $
    entities (left (subscribed (S.singleton "a")) :: Mealy Identity (Either Int Int) (Either Int Int))
      `shouldBe` S.singleton "a"

  it "first preserves the entity set" $
    entities (first (subscribed (S.singleton "a")) :: Mealy Identity (Int, Int) (Int, Int))
      `shouldBe` S.singleton "a"

  it "fmap preserves the entity set" $
    entities (fmap (+ 1) (subscribed (S.singleton "a")))
      `shouldBe` S.singleton "a"