module Unison.Codebase.Watch
  ( watchPath,
    WatchState (..),
    newWatchState,
    awaitEvent,
    unwatchPath,
    getWatchedPaths,
  )
where

import Control.Concurrent.STM (STM, TVar)
import Control.Concurrent.STM qualified as STM
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Map qualified as Map
import Data.Time.Clock (UTCTime, diffUTCTime)
import GHC.Conc (registerDelay)
import System.Directory (canonicalizePath, doesDirectoryExist, doesFileExist)
import System.FSNotify (Event (Added, Modified))
import System.FSNotify qualified as FSNotify
import System.FilePath (splitFileName)
import Unison.Prelude
import UnliftIO.Exception (tryAny)
import UnliftIO.STM (atomically)

-- | State for managing multiple watched paths.
data WatchState = WatchState
  { -- | The FSNotify watch manager
    WatchState -> WatchManager
watchManager :: FSNotify.WatchManager,
    -- | TVar containing the latest event from any watcher
    WatchState -> TVar (Maybe (FilePath, UTCTime))
latestEventVar :: TVar (Maybe (FilePath, UTCTime)),
    -- | Map from watched paths to their stop-watching actions
    WatchState -> TVar (Map FilePath (IO ()))
watchedPathsVar :: TVar (Map FilePath (IO ())),
    -- | Predicate for filtering files (e.g., .u files only)
    WatchState -> FilePath -> Bool
allowPredicate :: FilePath -> Bool,
    -- | Cache for debouncing file contents
    WatchState -> IORef (Map FilePath (Text, UTCTime))
previousFilesRef :: IORef (Map FilePath (Text, UTCTime))
  }

-- | Create a new watch state. The Ki scope is used for structured concurrency -
-- when the scope exits, all watcher threads are automatically cleaned up.
newWatchState :: FSNotify.WatchManager -> (FilePath -> Bool) -> IO WatchState
newWatchState :: WatchManager -> (FilePath -> Bool) -> IO WatchState
newWatchState WatchManager
mgr FilePath -> Bool
allow = do
  TVar (Maybe (FilePath, UTCTime))
latestEventVar <- Maybe (FilePath, UTCTime) -> IO (TVar (Maybe (FilePath, UTCTime)))
forall a. a -> IO (TVar a)
STM.newTVarIO Maybe (FilePath, UTCTime)
forall a. Maybe a
Nothing
  TVar (Map FilePath (IO ()))
watchedPathsVar <- Map FilePath (IO ()) -> IO (TVar (Map FilePath (IO ())))
forall a. a -> IO (TVar a)
STM.newTVarIO Map FilePath (IO ())
forall k a. Map k a
Map.empty
  IORef (Map FilePath (Text, UTCTime))
previousFilesRef <- Map FilePath (Text, UTCTime)
-> IO (IORef (Map FilePath (Text, UTCTime)))
forall a. a -> IO (IORef a)
newIORef Map FilePath (Text, UTCTime)
forall k a. Map k a
Map.empty
  WatchState -> IO WatchState
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    WatchState
      { $sel:watchManager:WatchState :: WatchManager
watchManager = WatchManager
mgr,
        $sel:latestEventVar:WatchState :: TVar (Maybe (FilePath, UTCTime))
latestEventVar = TVar (Maybe (FilePath, UTCTime))
latestEventVar,
        $sel:watchedPathsVar:WatchState :: TVar (Map FilePath (IO ()))
watchedPathsVar = TVar (Map FilePath (IO ()))
watchedPathsVar,
        $sel:allowPredicate:WatchState :: FilePath -> Bool
allowPredicate = FilePath -> Bool
allow,
        $sel:previousFilesRef:WatchState :: IORef (Map FilePath (Text, UTCTime))
previousFilesRef = IORef (Map FilePath (Text, UTCTime))
previousFilesRef
      }

-- | Add a file or directory to be watched. Returns the canonical path if successful, Nothing otherwise.
--
-- Each watched path spawns a background thread via Ki that manages the FSNotify watcher.
-- When the Ki scope exits, all watcher threads are automatically cleaned up.
watchPath :: WatchState -> FilePath -> IO (Maybe FilePath)
watchPath :: WatchState -> FilePath -> IO (Maybe FilePath)
watchPath WatchState
ws FilePath
path = do
  FilePath
canonPath <- FilePath -> IO FilePath
canonicalizePath FilePath
path
  Bool
isDir <- FilePath -> IO Bool
doesDirectoryExist FilePath
canonPath
  Bool
isFile <- FilePath -> IO Bool
doesFileExist FilePath
canonPath
  if Bool -> Bool
not (Bool
isDir Bool -> Bool -> Bool
|| Bool
isFile)
    then Maybe FilePath -> IO (Maybe FilePath)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe FilePath
forall a. Maybe a
Nothing
    else do
      Bool
alreadyWatched <- STM Bool -> IO Bool
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ FilePath -> Map FilePath (IO ()) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member FilePath
canonPath (Map FilePath (IO ()) -> Bool)
-> STM (Map FilePath (IO ())) -> STM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map FilePath (IO ())) -> STM (Map FilePath (IO ()))
forall a. TVar a -> STM a
STM.readTVar WatchState
ws.watchedPathsVar
      if Bool
alreadyWatched
        then Maybe FilePath -> IO (Maybe FilePath)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FilePath -> Maybe FilePath
forall a. a -> Maybe a
Just FilePath
canonPath) -- Already watching, consider it a success
        else do
          -- Create the handler that writes to our shared TVar
          let handler :: Event -> IO ()
              handler :: Event -> IO ()
handler = \case
                Added FilePath
fp UTCTime
t EventIsDirectory
FSNotify.IsFile | WatchState
ws.allowPredicate FilePath
fp -> STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar (Maybe (FilePath, UTCTime))
-> Maybe (FilePath, UTCTime) -> STM ()
forall a. TVar a -> a -> STM ()
STM.writeTVar WatchState
ws.latestEventVar ((FilePath, UTCTime) -> Maybe (FilePath, UTCTime)
forall a. a -> Maybe a
Just (FilePath
fp, UTCTime
t)))
                Modified FilePath
fp UTCTime
t EventIsDirectory
FSNotify.IsFile | WatchState
ws.allowPredicate FilePath
fp -> STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar (Maybe (FilePath, UTCTime))
-> Maybe (FilePath, UTCTime) -> STM ()
forall a. TVar a -> a -> STM ()
STM.writeTVar WatchState
ws.latestEventVar ((FilePath, UTCTime) -> Maybe (FilePath, UTCTime)
forall a. a -> Maybe a
Just (FilePath
fp, UTCTime
t)))
                Event
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

          -- Determine what to watch
          let watchAction :: IO (IO ())
watchAction =
                if Bool
isDir
                  then WatchManager
-> FilePath -> ActionPredicate -> (Event -> IO ()) -> IO (IO ())
FSNotify.watchDir WatchState
ws.watchManager FilePath
canonPath (Bool -> ActionPredicate
forall a b. a -> b -> a
const Bool
True) Event -> IO ()
handler
                  else do
                    -- For a single file, we watch the parent directory and filter for our file
                    let (FilePath
parentDir, FilePath
_) = FilePath -> (FilePath, FilePath)
splitFileName FilePath
canonPath
                    WatchManager
-> FilePath -> ActionPredicate -> (Event -> IO ()) -> IO (IO ())
FSNotify.watchDir WatchState
ws.watchManager FilePath
parentDir (\Event
e -> Event -> FilePath
eventPath Event
e FilePath -> FilePath -> Bool
forall a. Eq a => a -> a -> Bool
== FilePath
canonPath) Event -> IO ()
handler

          -- Start watching with FSNotify
          IO ()
stopListening <- IO (IO ())
watchAction

          -- Record that we're watching this path, with the actual stop action
          STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Map FilePath (IO ()))
-> (Map FilePath (IO ()) -> Map FilePath (IO ())) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
STM.modifyTVar WatchState
ws.watchedPathsVar (FilePath -> IO () -> Map FilePath (IO ()) -> Map FilePath (IO ())
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert FilePath
canonPath IO ()
stopListening)
          Maybe FilePath -> IO (Maybe FilePath)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FilePath -> Maybe FilePath
forall a. a -> Maybe a
Just FilePath
canonPath)
  where
    eventPath :: Event -> FilePath
    eventPath :: Event -> FilePath
eventPath = \case
      Added FilePath
p UTCTime
_time EventIsDirectory
_isDir -> FilePath
p
      Modified FilePath
p UTCTime
_time EventIsDirectory
_isDir -> FilePath
p
      FSNotify.Removed FilePath
p UTCTime
_time EventIsDirectory
_isDir -> FilePath
p
      FSNotify.ModifiedAttributes FilePath
p UTCTime
_time EventIsDirectory
_isDir -> FilePath
p
      FSNotify.WatchedDirectoryRemoved FilePath
p UTCTime
_time EventIsDirectory
_isDir -> FilePath
p
      FSNotify.CloseWrite FilePath
p UTCTime
_time EventIsDirectory
_isDir -> FilePath
p
      FSNotify.Unknown FilePath
p UTCTime
_time EventIsDirectory
_isDir FilePath
_eventString -> FilePath
p

-- | Await an event from any watched source.
--
-- This function implements debouncing with the following logic, intended to work around the tendency
-- for modern editors to create a flurry of rapid filesystem events when a file is saved:
--
-- 1. Block until an event arrives.
-- 2. Keep consuming events until 50ms elapse without an event.
-- 3. Return only the last event.
--
-- Note we don't have any smarts here for a flurry of events that are related to more than one file;
-- we just throw everything away except the last event. In practice, this has seemed to work fine.
--
-- Additionally, we keep in memory the file contents of previously-saved files, so that we can avoid
-- emitting events for files that last changed less than 500ms ago, and whose contents haven't changed.
awaitEvent :: WatchState -> IO (FilePath, Text)
awaitEvent :: WatchState -> IO (FilePath, Text)
awaitEvent WatchState
ws = do
  let awaitEvent0 :: IO (FilePath, UTCTime)
      awaitEvent0 :: IO (FilePath, UTCTime)
awaitEvent0 = do
        let go :: (FilePath, UTCTime) -> IO (FilePath, UTCTime)
            go :: (FilePath, UTCTime) -> IO (FilePath, UTCTime)
go (FilePath, UTCTime)
event0 = do
              TVar Bool
var <- Int -> IO (TVar Bool)
registerDelay Int
50_000
              (IO (IO (FilePath, UTCTime)) -> IO (FilePath, UTCTime)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO (FilePath, UTCTime)) -> IO (FilePath, UTCTime))
-> ([STM (IO (FilePath, UTCTime))] -> IO (IO (FilePath, UTCTime)))
-> [STM (IO (FilePath, UTCTime))]
-> IO (FilePath, UTCTime)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM (IO (FilePath, UTCTime)) -> IO (IO (FilePath, UTCTime))
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM (IO (FilePath, UTCTime)) -> IO (IO (FilePath, UTCTime)))
-> ([STM (IO (FilePath, UTCTime))] -> STM (IO (FilePath, UTCTime)))
-> [STM (IO (FilePath, UTCTime))]
-> IO (IO (FilePath, UTCTime))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [STM (IO (FilePath, UTCTime))] -> STM (IO (FilePath, UTCTime))
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum)
                [ do
                    (FilePath, UTCTime)
event1 <- STM (FilePath, UTCTime)
readLatestEvent
                    IO (FilePath, UTCTime) -> STM (IO (FilePath, UTCTime))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((FilePath, UTCTime) -> IO (FilePath, UTCTime)
go (FilePath, UTCTime)
event1),
                  do
                    TVar Bool -> STM Bool
forall a. TVar a -> STM a
STM.readTVar TVar Bool
var STM Bool -> (Bool -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
STM.check
                    IO (FilePath, UTCTime) -> STM (IO (FilePath, UTCTime))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((FilePath, UTCTime) -> IO (FilePath, UTCTime)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FilePath, UTCTime)
event0)
                ]
        (FilePath, UTCTime)
event <- STM (FilePath, UTCTime) -> IO (FilePath, UTCTime)
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically STM (FilePath, UTCTime)
readLatestEvent
        (FilePath, UTCTime) -> IO (FilePath, UTCTime)
go (FilePath, UTCTime)
event

      readLatestEvent :: STM (FilePath, UTCTime)
      readLatestEvent :: STM (FilePath, UTCTime)
readLatestEvent =
        TVar (Maybe (FilePath, UTCTime)) -> STM (Maybe (FilePath, UTCTime))
forall a. TVar a -> STM a
STM.readTVar WatchState
ws.latestEventVar STM (Maybe (FilePath, UTCTime))
-> (Maybe (FilePath, UTCTime) -> STM (FilePath, UTCTime))
-> STM (FilePath, UTCTime)
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe (FilePath, UTCTime)
Nothing -> STM (FilePath, UTCTime)
forall a. STM a
STM.retry
          Just (FilePath, UTCTime)
event -> do
            TVar (Maybe (FilePath, UTCTime))
-> Maybe (FilePath, UTCTime) -> STM ()
forall a. TVar a -> a -> STM ()
STM.writeTVar WatchState
ws.latestEventVar Maybe (FilePath, UTCTime)
forall a. Maybe a
Nothing
            (FilePath, UTCTime) -> STM (FilePath, UTCTime)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FilePath, UTCTime)
event

  -- Apply debouncing based on file contents cache
  let awaitEvent1 :: IO (FilePath, Text)
      awaitEvent1 :: IO (FilePath, Text)
awaitEvent1 = do
        (FilePath
file, UTCTime
t) <- IO (FilePath, UTCTime)
awaitEvent0
        IO Text -> IO (Either SomeException Text)
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny (FilePath -> IO Text
readUtf8 FilePath
file) IO (Either SomeException Text)
-> (Either SomeException Text -> IO (FilePath, Text))
-> IO (FilePath, Text)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          -- Somewhat-expected read error from a file that was just written. Just ignore the event and try again.
          Left SomeException
_ -> IO (FilePath, Text)
awaitEvent1
          Right Text
contents -> do
            Map FilePath (Text, UTCTime)
previousFiles <- IORef (Map FilePath (Text, UTCTime))
-> IO (Map FilePath (Text, UTCTime))
forall a. IORef a -> IO a
readIORef WatchState
ws.previousFilesRef
            case FilePath -> Map FilePath (Text, UTCTime) -> Maybe (Text, UTCTime)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup FilePath
file Map FilePath (Text, UTCTime)
previousFiles of
              Just (Text
contents0, UTCTime
t0) | Text
contents Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
contents0 Bool -> Bool -> Bool
&& (UTCTime
t UTCTime -> UTCTime -> NominalDiffTime
`diffUTCTime` UTCTime
t0) NominalDiffTime -> NominalDiffTime -> Bool
forall a. Ord a => a -> a -> Bool
< NominalDiffTime
0.5 -> IO (FilePath, Text)
awaitEvent1
              Maybe (Text, UTCTime)
_ -> do
                IORef (Map FilePath (Text, UTCTime))
-> Map FilePath (Text, UTCTime) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef WatchState
ws.previousFilesRef (Map FilePath (Text, UTCTime) -> IO ())
-> Map FilePath (Text, UTCTime) -> IO ()
forall a b. (a -> b) -> a -> b
$! FilePath
-> (Text, UTCTime)
-> Map FilePath (Text, UTCTime)
-> Map FilePath (Text, UTCTime)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert FilePath
file (Text
contents, UTCTime
t) Map FilePath (Text, UTCTime)
previousFiles
                (FilePath, Text) -> IO (FilePath, Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FilePath
file, Text
contents)

  IO (FilePath, Text)
awaitEvent1

-- | Stop watching a path. Returns True if the path was being watched.
unwatchPath :: WatchState -> FilePath -> IO Bool
unwatchPath :: WatchState -> FilePath -> IO Bool
unwatchPath WatchState
ws FilePath
path = do
  FilePath
canonPath <- FilePath -> IO FilePath
canonicalizePath FilePath
path
  Maybe (IO ())
maybeStop <- STM (Maybe (IO ())) -> IO (Maybe (IO ()))
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM (Maybe (IO ())) -> IO (Maybe (IO ())))
-> STM (Maybe (IO ())) -> IO (Maybe (IO ()))
forall a b. (a -> b) -> a -> b
$ do
    Map FilePath (IO ())
paths <- TVar (Map FilePath (IO ())) -> STM (Map FilePath (IO ()))
forall a. TVar a -> STM a
STM.readTVar WatchState
ws.watchedPathsVar
    case FilePath -> Map FilePath (IO ()) -> Maybe (IO ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup FilePath
canonPath Map FilePath (IO ())
paths of
      Maybe (IO ())
Nothing -> Maybe (IO ()) -> STM (Maybe (IO ()))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (IO ())
forall a. Maybe a
Nothing
      Just IO ()
stopAction -> do
        TVar (Map FilePath (IO ())) -> Map FilePath (IO ()) -> STM ()
forall a. TVar a -> a -> STM ()
STM.writeTVar WatchState
ws.watchedPathsVar (FilePath -> Map FilePath (IO ()) -> Map FilePath (IO ())
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete FilePath
canonPath Map FilePath (IO ())
paths)
        Maybe (IO ()) -> STM (Maybe (IO ()))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just IO ()
stopAction)
  case Maybe (IO ())
maybeStop of
    Maybe (IO ())
Nothing -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
    Just IO ()
stopAction -> do
      IO ()
stopAction
      Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

-- | Get the list of currently watched paths.
getWatchedPaths :: WatchState -> IO (Set FilePath)
getWatchedPaths :: WatchState -> IO (Set FilePath)
getWatchedPaths WatchState
ws = STM (Set FilePath) -> IO (Set FilePath)
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM (Set FilePath) -> IO (Set FilePath))
-> STM (Set FilePath) -> IO (Set FilePath)
forall a b. (a -> b) -> a -> b
$ Map FilePath (IO ()) -> Set FilePath
forall k a. Map k a -> Set k
Map.keysSet (Map FilePath (IO ()) -> Set FilePath)
-> STM (Map FilePath (IO ())) -> STM (Set FilePath)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map FilePath (IO ())) -> STM (Map FilePath (IO ()))
forall a. TVar a -> STM a
STM.readTVar WatchState
ws.watchedPathsVar