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
  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
  watchedPathsVar <- STM.newTVarIO Map.empty
  previousFilesRef <- newIORef Map.empty
  pure
    WatchState
      { watchManager = mgr,
        latestEventVar = latestEventVar,
        watchedPathsVar = watchedPathsVar,
        allowPredicate = allow,
        previousFilesRef = 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
  canonPath <- FilePath -> IO FilePath
canonicalizePath FilePath
path
  isDir <- doesDirectoryExist canonPath
  isFile <- doesFileExist canonPath
  if not (isDir || isFile)
    then pure Nothing
    else do
      alreadyWatched <- atomically $ Map.member canonPath <$> STM.readTVar ws.watchedPathsVar
      if alreadyWatched
        then pure (Just canonPath) -- Already watching, consider it a success
        else do
          -- Create the handler that writes to our shared TVar
          let 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 =
                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
          stopListening <- watchAction

          -- Record that we're watching this path, with the actual stop action
          atomically $ STM.modifyTVar ws.watchedPathsVar (Map.insert canonPath stopListening)
          pure (Just 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
              var <- Int -> IO (TVar Bool)
registerDelay Int
50_000
              (join . atomically . asum)
                [ do
                    event1 <- readLatestEvent
                    pure (go event1),
                  do
                    STM.readTVar var >>= STM.check
                    pure (pure event0)
                ]
        event <- STM (FilePath, UTCTime) -> IO (FilePath, UTCTime)
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically STM (FilePath, UTCTime)
readLatestEvent
        go 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
        (file, t) <- IO (FilePath, UTCTime)
awaitEvent0
        tryAny (readUtf8 file) >>= \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
            previousFiles <- IORef (Map FilePath (Text, UTCTime))
-> IO (Map FilePath (Text, UTCTime))
forall a. IORef a -> IO a
readIORef WatchState
ws.previousFilesRef
            case Map.lookup file 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
  canonPath <- FilePath -> IO FilePath
canonicalizePath FilePath
path
  maybeStop <- atomically $ do
    paths <- STM.readTVar ws.watchedPathsVar
    case Map.lookup canonPath 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 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