module AFRPSpec (spec) where

import Control.Arrow (arr)
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.Time (NominalDiffTime, UTCTime (..), addUTCTime, diffUTCTime)
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) nil

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

runPure :: Mealy Identity a b -> [a] -> [b]
runPure _ [] = []
runPure m (a : as) = case runIdentity (AFRP.runMealy m id fakeRequest a) of
  (b, m') -> b : runPure m' 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 _ [] = []
runTimed m ((s, a) : as) =
  case runIdentity (AFRP.runMealy m id (Request (sec s) nil) a) of
    (b, m') -> b : runTimed m' 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 s0 as = go m s0 as
  where
    go _ s [] = ([], s)
    go m' s (a : rest) =
      case unSt (AFRP.runMealy m' id fakeRequest a) s of
        ((b, m''), s') -> let (bs, s'') = go m'' s' rest in (b : bs, s'')

spec :: Spec
spec = describe "AFRP" $ do
  holdSpec
  eventsSpec
  isEventSpec
  tagSpec
  toEventSpec
  lMergeSpec
  changesSpec
  edgeSpec
  filterASpec
  slidingSpec
  mapAccumSpec
  preMapAccumSpec
  durationSpec
  delayEventSpec
  debounceSpec
  rollupSpec
  fixedSpec
  effSpec
  switchSpec
  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)

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]

fixedSpec :: Spec
fixedSpec = describe "fixed" $ do
  it "accumulates events within a window and rolls over on expiry" $
    runTimed (fixed 10)
             [ (0, Event 'a'), (1, Event 'b'), (2, Tick)
             , (12, Event 'c'), (13, Tick)
             ]
      `shouldBe` [ ['a'], ['a', 'b'], ['a', 'b']
                 , ['c'], ['c']
                 ]

  it "starts a window even on a leading Tick" $
    runTimed (fixed 10)
             [ (0, Tick), (1, Event 'a')
             , (12, Tick), (13, Tick)
             , (25, Event 'b')
             ]
      `shouldBe` [ [], ['a'], [], [], ['b'] ]

  it "output is always the current window's accumulated list" $
    hedgehog $ do
      w <- forAll $ Gen.int (Range.constant 1 10)
      evs <- forAll $ Gen.list (Range.linear 0 30) eventGen
      let out = runTimed (fixed w) (zip [0 ..] evs)
          expected =
            [ [ x | Event x <- take (i - lo + 1) (drop lo evs) ]
            | i <- [0 .. length evs - 1]
            , let lo = (i `div` w) * w
            ]
      out === expected

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

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, 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

switchSpec :: Spec
switchSpec = describe "switch" $ do
  it "switches to the continuation at the first Event" $
    runPure (switch (arr (\x -> (x, if x >= 3 then Event () else Tick)))
                     (const (arr (const 99))))
             [1, 2, 3, 4, 5]
      `shouldBe` [1, 2, 99, 99, 99]

  it "never switches if no Event is emitted" $
    runPure (switch (arr (\x -> (x, Tick :: Event ())))
                     (const (arr (const 99))))
             [1, 2, 3]
      `shouldBe` [1, 2, 3]

  it "prefix outputs come from the first arrow, suffix from the continuation" $
    hedgehog $ do
      threshold <- forAll $ Gen.int (Range.linear (-20) 20)
      xs <- forAll $ Gen.list (Range.linear 0 30) (Gen.int (Range.linear (-20) 20))
      let first = arr (\x -> (x, if x >= threshold then Event () else Tick))
          out = runPure (switch first (const (arr (const 99)))) xs
          (pre, _post) = break (>= threshold) xs
      take (length pre) out === pre
      drop (length pre) out === replicate (length xs - length pre) 99

mapAccumRequestSpec :: Spec
mapAccumRequestSpec = describe "mapAccumRequest" $ do
  it "accumulates request times, post-state extraction" $
    runTimed (mapAccumRequest (\req s _ -> s ++ [requestTime req]) [] id)
             [(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 ++ [requestTime req]) [] id)
                         [(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 ++ [requestTime req]) [] id)
             [(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 ++ [requestTime req]) [] id)
                         [(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 not (even 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 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 = zip (take n vals) (take n evs)
          out = runPure sample ps
      for_ (zip ps out) $ \((v, ev), o) ->
        o === tag v ev