{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

module Glean.Util.Trace
  ( Listener
  , listener
  , recorder
  , notify

  , Expect
  , expect
  , reject
  , accept
  , want
  , wants
  , (>|<)
  , opt
  , alt
  , (|||)
  , parallel
  )
where

import Control.Monad
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as HashMap
import Data.Hashable (Hashable)
import Data.Typeable

import Util.STM

-- | A 'Listener' can be notified of 'Typeable' events. Listeners can be
-- combined via the 'Semigroup'/'Monoid' interface.
newtype Listener = Listener (forall a. Typeable a => a -> STM ())

instance Semigroup Listener where
  Listener f <> Listener g = Listener $ \x -> f x >> g x

instance Monoid Listener where
  mempty = Listener $ const $ return ()

-- | Create a 'Listener' for a particular type.
listener :: Typeable a => (a -> STM ()) -> Listener
listener f = Listener $ mapM_ f . cast

-- | Create a 'Listener' which writes all events to a queue.
recorder :: Typeable a => IO (Listener, TQueue a)
recorder = do
  events <- newTQueueIO
  return (listener $ writeTQueue events, events)

-- | Notify a listener about an event, ignoring all exceptions.
notify :: Typeable a => Listener -> a -> STM ()
notify (Listener f) = f


-- | 'Expect' encapsulates expectations about which events should happen in
-- what order. We can expect a particular event via 'want', sequentially
-- compose expectations via 'Semigroup'/'Monoid', encode alternatives via
-- '>|<' and interleave expectations via '|||'. Once an 'Expect' is constructed,
-- it can be applied to the contents of a 'TQueue' via 'expect'.
data Expect a = Expect
  (HashMap a (Expect a))  -- what to do with each event
  Bool                    -- is this 'Expect' optional
  deriving(Show)

instance (Eq a, Hashable a) => Semigroup (Expect a) where
  Expect xs a <> q@(Expect _ b) =
    alt $ Expect (fmap (<> q) xs) (a && b) : [q | a]

instance (Eq a, Hashable a) => Monoid (Expect a) where
  mempty = Expect mempty True

-- | An 'Expect' that always fails.
reject :: (Eq a, Hashable a) => Expect a
reject = Expect mempty False

-- | An 'Expect' that always succeeds.
accept :: (Eq a, Hashable a) => Expect a
accept = mempty

-- | Expect a specific event.
want :: (Eq a, Hashable a) => a -> Expect a
want x = Expect (HashMap.singleton x accept) False

-- | Expect a sequence of events.
wants :: (Eq a, Hashable a) => [a] -> Expect a
wants = foldr (\x p -> Expect (HashMap.singleton x p) False) accept

infixl 3 >|<

-- | Alternative
(>|<) :: (Eq a, Hashable a) => Expect a -> Expect a -> Expect a
Expect xs a >|< Expect ys b = Expect (HashMap.unionWith (>|<) xs ys) (a || b)

-- | Optional
opt :: (Eq a, Hashable a) => Expect a -> Expect a
opt p = p >|< accept

-- | Multiple alternatives
alt :: (Eq a, Hashable a) => [Expect a] -> Expect a
alt = foldr (>|<) reject

infixl 2 |||

-- | Parallel composition - the two expectations can be arbitrarily interleaved
-- (both ordering within each 'Expect' is preserved).
(|||) :: (Eq a, Hashable a) => Expect a -> Expect a -> Expect a
p@(Expect xs a) ||| q@(Expect ys b) = alt $
  [ Expect
      (HashMap.unionWith
        (>|<)
        (fmap (||| Expect ys False) xs)
        (fmap (Expect xs False |||) ys))
      (a && b) ]
  ++ [p | b]
  ++ [q | a]

-- | Interleave multiple 'Expect's.
parallel :: (Eq a, Hashable a) => [Expect a] -> Expect a
parallel = foldr (|||) accept

-- | Consume elements of the 'TQueue' until we hit an 'accept', failing on
-- unexpected elements. This will always consume the smallest possible number
-- of events.
expect :: (Eq a, Hashable a, Show a) => TQueue a -> Expect a -> IO ()
expect queue (Expect xs b)
  | HashMap.null xs = if b then return () else fail "reject"
  | otherwise = join $ atomically $ do
      x <- readTQueue queue
      case HashMap.lookup x xs of
        Just p -> return $ expect queue p
        Nothing
          | b -> do
              unGetTQueue queue x
              return $ return ()
          | otherwise -> return $ fail $ concat
              [ "unexpected ", show x
              , "; expected one of ", show (HashMap.keys xs) ]