{-# LANGUAGE MagicHash #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecursiveDo #-}
import System.Mem
import System.Mem.Weak
import Data.IORef
import GHC.IORef
import GHC.STRef
import GHC.IO
import GHC.Weak
import GHC.Prim
import Control.Monad.Primitive
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.Set (Set)
import qualified Data.Set as Set
import System.IO.Unsafe
import Control.Monad
import Control.Concurrent
import Data.Foldable
import Data.These
import Unsafe.Coerce
import Control.Monad.Fix
import Data.List.NonEmpty (NonEmpty (..))
import Data.Semigroup

import UIO
import Test.UIO

main :: IO ()
main = do
  testDualWeak
  testWeakChain
  testGraphChain
  testGraphX
  testGraphO

type UIO = IO

data WeakBagInput a
data WeakBagOutput a
data WeakMutVarInput a = WeakMutVarInput (WeakWithLifespan (IORef a))
data WeakMutVarOutput a = WeakMutVarOutput (IORef a)

data WeakWithLifespan a = WeakWithLifespan (Weak a) Lifespan

newWeakMutVar :: a -> UIO (WeakMutVarInput a, WeakMutVarOutput a)
newWeakMutVar a = do
  r <- newIORef a
  l <- lifespanOfIORef r
  w <- newWeakWithLifespan l r
  pure (WeakMutVarInput w, WeakMutVarOutput r)

writeWeakMutVar :: WeakMutVarInput a -> UIO (Effect a)
writeWeakMutVar target = mapEffect const =<< modifyWeakMutVar target

modifyWeakMutVar :: WeakMutVarInput a -> UIO (Effect (a -> a))
modifyWeakMutVar (WeakMutVarInput w) = Effect <$> do
  mapWeakWithLifespan' (\r f -> atomicModifyIORef' r $ \v -> (f v, ())) w

readWeakMutVar :: WeakMutVarOutput a -> Coeffect a
readWeakMutVar (WeakMutVarOutput r) = Coeffect $ readIORef r

testEffect1 :: IO ()
testEffect1 = do
  (i, o) <- newWeakMutVar "A"
  e <- traceEffect id =<< writeWeakMutVar i
  performGCUntilFinalizersQuiesce
  putStrLn "X"
  runEffect e "B"
  touch o
  _ <- runCoeffect (readWeakMutVar o)
  performGCUntilFinalizersQuiesce
  putStrLn "Y"
  runEffect e "C"
  putStrLn "Z"

testEffect2 :: IO ()
testEffect2 = do
  (ia, oa) <- newWeakMutVar (1 :: Int)
  (ib, ob) <- newWeakMutVar 1
  modifyA <- traceEffect (const "modifyA") =<< modifyWeakMutVar ia
  modifyB <- traceEffect (const "modifyB") =<< modifyWeakMutVar ib
  let readA = readWeakMutVar oa
  incrementA <- mapEffect (\() -> (+1)) modifyA
  addToB <- mapEffect (\n -> (+n)) modifyB
  addAToB <- withCoeffect readA addToB
  e <- incrementA `andThen` addAToB
  performGCUntilFinalizersQuiesce
  runEffect e ((), ())
  touch oa

runEffect :: Effect a -> a -> IO ()
runEffect (Effect e) a = deRefWeakWithLifespan e >>= \case
  Nothing -> pure ()
  Just f -> f a

traceEffect :: (a -> String) -> Effect a -> UIO (Effect a)
traceEffect f = weakEffect $ putStrLn . f

-- Run the given IO action if the provided Effect is still runnable.  Note that the IO action may run even if the effect is *not* runnable, depending on garbage collection timing.
weakEffect :: (a -> IO ()) -> Effect a -> UIO (Effect a)
weakEffect f (Effect e) = fmap Effect $ forWeakWithLifespan' e $ \fe a -> do
  f a
  fe a

runCoeffect :: Coeffect a -> IO a
runCoeffect (Coeffect c) = c


emptyWeak :: Weak a
emptyWeak = unsafeCoerce emptyWeakUnit

{-# NOINLINE emptyWeakUnit #-}
emptyWeakUnit :: Weak ()
emptyWeakUnit = unsafePerformIO $ do
  w <- mkWeakPtr () Nothing
  finalize w
  pure w

newWeakWithLifespan :: Lifespan -> a -> UIO (WeakWithLifespan a)
newWeakWithLifespan (Lifespan l) a = do
  deRefWeak l >>= \case
    Nothing -> pure $ WeakWithLifespan emptyWeak (Lifespan l)
    Just r -> do
      w <- mkWeakWithIORefKey r a
      pure $ WeakWithLifespan w (Lifespan l)

underlyingLifespan :: WeakWithLifespan a -> Lifespan
underlyingLifespan (WeakWithLifespan _ l) = l

deRefWeakWithLifespan :: WeakWithLifespan a -> UIO (Maybe a)
deRefWeakWithLifespan (WeakWithLifespan w _) = deRefWeak w

-- Applies the function strictly; usually this is what you want, so that extra data is not retained
forWeakWithLifespan' :: WeakWithLifespan a -> (a -> b) -> UIO (WeakWithLifespan b)
forWeakWithLifespan' = flip mapWeakWithLifespan'

-- Applies the function strictly; usually this is what you want, so that extra data is not retained
mapWeakWithLifespan' :: (a -> b) -> WeakWithLifespan a -> UIO (WeakWithLifespan b)
mapWeakWithLifespan' f (WeakWithLifespan w l) = do
  deRefWeak w >>= \case
    Nothing -> pure $ WeakWithLifespan emptyWeak l
    Just v -> newWeakWithLifespan l $! f v

bothAlive :: WeakWithLifespan a -> WeakWithLifespan b -> UIO (WeakWithLifespan (a, b))
bothAlive = undefined

-- This is impossible because it assumes that weaks *do actual work* when you retrieve them.  This isn't what we want.
alignAlive :: WeakWithLifespan a -> WeakWithLifespan b -> UIO (WeakWithLifespan (These a b))
alignAlive = undefined

data Effect a = Effect (WeakWithLifespan (a -> UIO ())) -- Take an `a` and do a side effect with it

data Coeffect a = Coeffect (UIO a)  -- Read an `a` without doing any side effects

-- Like `also`, but also guarantees order.  But I'm not sure what the semantics should really be here, since we could want coeffects ordered separately from effects
andThen :: Effect a -> Effect b -> UIO (Effect (a, b))
andThen = also

also :: Effect a -> Effect b -> UIO (Effect (a, b))
also (Effect ea) (Effect eb) = Effect <$> do
  myLifespan <- unionLifespan (underlyingLifespan ea) (underlyingLifespan eb)
  newWeakWithLifespan myLifespan $ \(a, b) -> do
    deRefWeakWithLifespan ea >>= \case
      Nothing -> pure ()
      Just fa -> fa a
    deRefWeakWithLifespan eb >>= \case
      Nothing -> pure ()
      Just fb -> fb b

mapEffect :: (b -> a) -> Effect a -> UIO (Effect b)
mapEffect f (Effect e) = Effect <$> mapWeakWithLifespan' (\fe -> fe . f) e

withCoeffect :: Coeffect a -> Effect a -> UIO (Effect ())
withCoeffect (Coeffect c) (Effect e) = Effect <$> do
  let f fe () = do
        v <- c
        fe v
  mapWeakWithLifespan' f e

coAlso :: Coeffect a -> Coeffect b -> Coeffect (a, b)
coAlso = undefined

--TODO: This should use a WeakBag
newtype Lifespan = Lifespan (Weak (IORef [LifespanBacklink]))
newtype LifespanBacklink = LifespanBacklink (IORef [LifespanBacklink])

modifyWeakIORef :: Weak (IORef a) -> (a -> a) -> IO ()
modifyWeakIORef w f = deRefWeak w >>= \case
  Nothing -> pure ()
  Just r -> atomicModifyIORef' r $ \v -> (f v, ())

lifespanOfIORef :: IORef a -> UIO Lifespan
lifespanOfIORef basis = do
  mine <- newIORef []
  w <- mkWeakWithIORefKey basis mine -- This exploits the fact that System.Mem.Weak references keep the value alive even when the weak reference itself dies.
  pure $ Lifespan w

-- Return a lifespan
unionLifespan :: Lifespan -> Lifespan -> UIO Lifespan
unionLifespan (Lifespan a) (Lifespan b) = do
  r <- newIORef []
  w <- mkWeakWithIORefKey r r
  modifyWeakIORef a (LifespanBacklink r :)
  modifyWeakIORef b (LifespanBacklink r :)
  pure $ Lifespan w

intersectionLifespan :: Lifespan -> Lifespan -> UIO Lifespan
intersectionLifespan a b = undefined

-- If we do readWeakMutVar v `bind` writeWeakMutVar v', we should only keep v alive if `v'`'s output side is alive
bind :: IO a -> (a -> IO b) -> IO b
bind = undefined

data Event a = Event
  { _event_items :: WeakBagInput (a -> IO ()) -- Allows adding items to the weak bag.  Does not keep the weak bag alive; if the bag is gone, adding an item does nothing.
  , _event_currentValue :: WeakMutVarOutput (Maybe a) -- Allows seeing the current state of the event: if Nothing, it either isn't firing or hasn't fired yet this frame; if Just, it has fired this frame.
  --TODO: Position in topological ordering; only needed when we introduce Merge
  }

data Trigger a = Trigger
  { _trigger_items :: WeakBagOutput (a -> IO ()) -- Allows retrieving the contents of the weak bag
  , _trigger_currentValue :: WeakMutVarInput (Maybe a)
  }

newEvent :: IO (Trigger a, Event a)
newEvent = undefined

fireTrigger :: Trigger a -> a -> IO ()
fireTrigger = undefined

data Hold a = Hold
  {
  }

newHold :: a -> Event a -> IO (Hold a)
newHold = undefined

-- Demonstrate that a Weak's key keeps its value alive, even if the Weak is dead
testWeakChain :: IO ()
testWeakChain = do
  a <- newIORef ()
  b <- newIORef ()
  _ <- mkWeakWithIORefKey a b
  w <- mkWeakWithIORefKey b "b is still alive"
  performGCUntilFinalizersQuiesce
  Just "b is still alive" <- deRefWeak w
  touch a

testGraphChain :: IO ()
testGraphChain = do
  do
    (nai, neo) <- buildGraphChain
    performGCUntilFinalizersQuiesce
    ["a", "b", "c", "d", "e"] <- toList <$> findForwardVals nai
    ["a", "b", "c", "d", "e"] <- toList <$> findBackwardVals neo
    touch (nai, neo)
  do
    (nai, _) <- buildGraphChain
    performGCUntilFinalizersQuiesce
    [] <- toList <$> findForwardVals nai
    touch nai
  do
    (_, neo) <- buildGraphChain
    performGCUntilFinalizersQuiesce
    [] <- toList <$> findBackwardVals neo
    touch neo

buildGraphChain :: IO (NodeInput String, NodeOutput String)
buildGraphChain = do
  (nai, nao) <- newNode "a"
  (nbi, nbo) <- newNode "b"
  (nci, nco) <- newNode "c"
  (ndi, ndo) <- newNode "d"
  (nei, neo) <- newNode "e"
  link nao nbi
  link nbo nci
  link nco ndi
  link ndo nei
  pure (nai, neo)

buildGraphX :: IO (NodeInput String, NodeInput String, NodeOutput String, NodeOutput String)
buildGraphX = do
  (nai, nao) <- newNode "a"
  (nbi, nbo) <- newNode "b"
  (nci, nco) <- newNode "c"
  (ndi, ndo) <- newNode "d"
  (nei, neo) <- newNode "e"
  link nao nci
  link nbo nci
  link nco ndi
  link nco nei
  pure (nai, nbi, ndo, neo)

testGraphX :: IO ()
testGraphX = do
  do
    (nai, nbi, ndo, neo) <- buildGraphX
    performGCUntilFinalizersQuiesce
    ["a", "c", "d", "e"] <- toList <$> findForwardVals nai
    ["b", "c", "d", "e"] <- toList <$> findForwardVals nbi
    ["a", "b", "c", "d"] <- toList <$> findBackwardVals ndo
    ["a", "b", "c", "e"] <- toList <$> findBackwardVals neo
    touch (nai, nbi, ndo, neo)
  do
    (nai, nbi, ndo, neo) <- buildGraphX
    performGCUntilFinalizersQuiesce
    ["a", "c", "d"] <- toList <$> findForwardVals nai
    ["b", "c", "d"] <- toList <$> findForwardVals nbi
    touch (nai, nbi, ndo)
  do
    (nai, nbi, ndo, neo) <- buildGraphX
    performGCUntilFinalizersQuiesce
    [] <- toList <$> findForwardVals nai
    [] <- toList <$> findForwardVals nbi
    touch (nai, nbi)

testGraphO :: IO ()
testGraphO = do
  do
    (nai, nao) <- buildGraphO
    performGCUntilFinalizersQuiesce
    ["a", "b", "c"] <- toList <$> findBackward nao
    ["a", "b", "c"] <- toList <$> findForward nai
    touch (nai, nao)
  do
    (_, nao) <- buildGraphO
    performGCUntilFinalizersQuiesce
    [] <- toList <$> findBackward nao
    touch nao
  do
    (nai, _) <- buildGraphO
    performGCUntilFinalizersQuiesce
    [] <- toList <$> findForward nai
    touch nai

buildGraphO :: IO (NodeInput String, NodeOutput String)
buildGraphO = do
  (nai, nao) <- newNode "a"
  (nbi, nbo) <- newNode "b"
  (nci, nco) <- newNode "c"
  link nao nbi
  link nbo nci
  link nco nai
  pure (nai, nao)

data NodeInput a = NodeInput
  { _nodeInput_key :: !Int
  , _nodeInput_contents :: !(Weak (IORef (Maybe a)))
  , _nodeInput_forwardLinks :: !(IORef (IntMap (NodeInput a)))
  , _nodeInput_backLinks :: !(Weak (IORef (IntMap (NodeOutput a))))
  }

data NodeOutput a = NodeOutput
  { _nodeOutput_key :: !Int
  , _nodeOutput_contents :: !(Weak (IORef (Maybe a)))
  , _nodeOutput_forwardLinks :: !(Weak (IORef (IntMap (NodeInput a))))
  , _nodeOutput_backLinks :: !(IORef (IntMap (NodeOutput a)))
  }

{-# NOINLINE globalNodeIdRef #-}
globalNodeIdRef :: IORef Int
globalNodeIdRef = unsafePerformIO $ newIORef 1

newNodeId :: IO Int
newNodeId = atomicModifyIORef' globalNodeIdRef $ \n -> (succ n, n)

newNode :: a -> IO (NodeInput a, NodeOutput a)
newNode v = do
  nodeId <- newNodeId
  backLinks <- newIORef mempty
  forwardLinks <- newIORef mempty
  contentsRef <- newIORef $ Just v
  w <- mkWeakWithIORefKey backLinks contentsRef
  wBack <- mkWeakWithIORefKey backLinks backLinks
  wForward <- mkWeakWithIORefKeyWithFinalizer forwardLinks forwardLinks $ do
    writeIORef finalizerDidRunRef True
    deRefWeak w >>= \case
      Nothing -> pure ()
      Just contentsRef' -> do
        writeIORef contentsRef' Nothing
    deRefWeak wBack >>= \case
      Nothing -> pure ()
      Just backLinks' -> do
        writeIORef backLinks' mempty
        --TODO: Clear out all the nodes forward of us
  pure
    ( NodeInput
      { _nodeInput_key = nodeId
      , _nodeInput_contents = w
      , _nodeInput_backLinks = wBack
      , _nodeInput_forwardLinks = forwardLinks
      }
    , NodeOutput
      { _nodeOutput_key = nodeId
      , _nodeOutput_contents = w
      , _nodeOutput_backLinks = backLinks
      , _nodeOutput_forwardLinks = wForward
      }
    )

getNodeContents :: Weak (IORef (Maybe a)) -> IO (Maybe a)
getNodeContents w = do
  deRefWeak w >>= \case
    Nothing -> pure Nothing
    Just r -> readIORef r

link :: NodeOutput a -> NodeInput a -> IO ()
link aOut bIn = do
  deRefWeak (_nodeInput_backLinks bIn) >>= \case
    Nothing -> pure ()
    Just backLinks -> deRefWeak (_nodeOutput_forwardLinks aOut) >>= \case
      Nothing -> pure ()
      Just forwardLinks -> do
        atomicModifyIORef' forwardLinks $ \m -> (IntMap.insert (_nodeInput_key bIn) bIn m, ())
        atomicModifyIORef' backLinks $ \m -> (IntMap.insert (_nodeOutput_key aOut) aOut m, ())

findForwardVals :: Ord a => NodeInput a -> IO (Set a)
findForwardVals i = Set.fromList . IntMap.elems <$> findForward i

findForward :: forall a. NodeInput a -> IO (IntMap a)
findForward i = go mempty (IntMap.singleton (_nodeInput_key i) i)
  where
    go :: IntMap a -> IntMap (NodeInput a) -> IO (IntMap a)
    go found toSearch = case IntMap.minViewWithKey toSearch of
      Nothing -> pure found
      Just ((k, thisIn), restToSearch) -> do
        getNodeContents (_nodeInput_contents thisIn) >>= \case
          Nothing -> go found restToSearch
          Just this -> do
            newLinks <- readIORef (_nodeInput_forwardLinks thisIn)
            go (IntMap.insert k this found) (IntMap.union restToSearch (newLinks `IntMap.difference` found))

findBackwardVals :: Ord a => NodeOutput a -> IO (Set a)
findBackwardVals i = Set.fromList . IntMap.elems <$> findBackward i

findBackward :: forall a. NodeOutput a -> IO (IntMap a)
findBackward o = go mempty (IntMap.singleton (_nodeOutput_key o) o)
  where
    go :: IntMap a -> IntMap (NodeOutput a) -> IO (IntMap a)
    go found toSearch = case IntMap.minViewWithKey toSearch of
      Nothing -> pure found
      Just ((k, thisOut), restToSearch) -> do
        getNodeContents (_nodeOutput_contents thisOut) >>= \case
          Nothing -> go found restToSearch
          Just this -> do
            newLinks <- readIORef (_nodeOutput_backLinks thisOut)
            go (IntMap.insert k this found) (IntMap.union restToSearch (newLinks `IntMap.difference` found))

testDualWeak :: IO ()
testDualWeak = do
  do
    target <- newIORef ()
    r <- mkWeakWithIORefKey target ()
    (w, a, b) <- newDualWeak target
    performGC
    threadDelay 1000000
    Just () <- deRefWeak r
    Just _ <- getDualWeak w
    touch (a, b)
    performGCUntilFinalizersQuiesce
  do
    target <- newIORef ()
    r <- mkWeakWithIORefKey target ()
    (w, a, _) <- newDualWeak target
    performGC
    threadDelay 1000000
    Nothing <- deRefWeak r
    Nothing <- getDualWeak w
    touch a
    performGCUntilFinalizersQuiesce
  do
    target <- newIORef ()
    r <- mkWeakWithIORefKey target ()
    (w, _, b) <- newDualWeak target
    performGC
    threadDelay 1000000
    performGC
    Nothing <- deRefWeak r
    Nothing <- getDualWeak w
    touch b
    performGCUntilFinalizersQuiesce

performGCUntilFinalizersQuiesce :: IO ()
performGCUntilFinalizersQuiesce = do
  writeIORef finalizerDidRunRef False
  performGC
  fence <- createFinalizerFence -- Based on my shaky memory, I think finalizers are processed in a queue, and therefore this fence will run after all the other finalizers in the GC we just finished.  However, it may run at ANY point in the subsequent GC, so we can't rely on it to know that *that* GC had no finalizers run.
  performGC
  waitForFinalizerFence fence
  readIORef finalizerDidRunRef >>= \case
    True -> performGCUntilFinalizersQuiesce
    False -> pure ()

{-# NOINLINE finalizerDidRunRef #-}
finalizerDidRunRef :: IORef Bool
finalizerDidRunRef = unsafePerformIO $ newIORef False

newtype FinalizerFence = FinalizerFence (MVar ())

createFinalizerFence :: IO FinalizerFence
createFinalizerFence = do
  r <- newIORef ()
  v <- newEmptyMVar
  _ <- mkWeakWithIORefKeyWithFinalizer r () $ putMVar v ()
  pure $ FinalizerFence v

waitForFinalizerFence :: FinalizerFence -> IO ()
waitForFinalizerFence (FinalizerFence v) = takeMVar v

newtype DualWeak a = DualWeak (Weak (Weak (IORef (Maybe a))))
newtype Ticket = Ticket (IORef ())

-- | Creates a weak reference to `a` which only remains alive if *both* tickets are alive
newDualWeak :: a -> IO (DualWeak a, Ticket, Ticket)
newDualWeak v = do
  vRef <- newIORef $ Just v
  tInner <- newIORef ()
  wInner <- mkWeakWithIORefKey tInner vRef
  tOuter <- newIORef ()
  wOuter <- mkWeakWithIORefKeyWithFinalizer tOuter wInner $ do
    writeIORef finalizerDidRunRef True
    deRefWeak wInner >>= \case
      Nothing -> pure ()
      Just vRef' -> do
        writeIORef vRef' Nothing
  pure (DualWeak wOuter, Ticket tOuter, Ticket tInner)

getDualWeak :: DualWeak a -> IO (Maybe a)
getDualWeak (DualWeak wOuter) = do
  deRefWeak wOuter >>= \case
    Nothing -> pure Nothing
    Just wInner -> deRefWeak wInner >>= \case
      Nothing -> pure Nothing
      Just vRef -> readIORef vRef

mkWeakWithIORefKey :: IORef a -> b -> IO (Weak b)
mkWeakWithIORefKey (IORef (STRef r#)) v = IO $ \s ->
    case mkWeakNoFinalizer# r# v s of (# s1, w #) -> (# s1, Weak w #)

mkWeakWithIORefKeyWithFinalizer :: IORef a -> b -> IO () -> IO (Weak b)
mkWeakWithIORefKeyWithFinalizer (IORef (STRef r#)) v (IO f) = IO $ \s ->
    case mkWeak# r# v f s of (# s1, w #) -> (# s1, Weak w #)