content: module Cachix.Daemon.Progress
    ( UploadProgress,
      ProgressState,
      new,
      complete,
      fail,
      tick,
      update,
      newProgressState,
      displayPushEvent,
    )
  where
  
  import Cachix.Client.HumanSize (humanSize)
  import Cachix.Daemon.Types (PushRetryStatus (..))
  import Cachix.Daemon.Types.PushEvent (PushEvent (..), PushEventMessage (..))
  import Control.Concurrent.Async qualified as Async
  import Control.Concurrent.MVar
  import Data.HashMap.Strict qualified as HashMap
  import Data.IORef
  import Data.String (String)
  import Data.Text qualified as T
  import Protolude
  import System.Console.AsciiProgress qualified as Ascii
  import System.Console.AsciiProgress.Internal qualified as Ascii.Internal
  import System.Console.Pretty
  import System.Environment (lookupEnv)
  import System.IO (hIsTerminalDevice)
  
  data UploadProgress
    = ProgressBar
        { path :: String,
          size :: Int64,
          progressBar :: Ascii.ProgressBar
        }
    | FallbackText
        { path :: String,
          size :: Int64
        }
  
  -- | State for tracking multiple progress bars
  type ProgressState = IORef (HashMap.HashMap FilePath (UploadProgress, PushRetryStatus))
  
  -- | Create a new progress state for tracking multiple progress bars
  newProgressState :: IO ProgressState
  newProgressState = newIORef HashMap.empty
  
  new :: Handle -> String -> Int64 -> PushRetryStatus -> IO UploadProgress
  new hdl path size retryStatus = do
    isCI <- liftIO $ (== Just "true") <$> lookupEnv "CI"
    isTerminal <- liftIO $ hIsTerminalDevice hdl
    if isTerminal && not isCI
      then do
        progressBar <- newProgressBar path size retryStatus
        return $ ProgressBar {..}
      else do
        hPutStr stderr $ uploadStartFallback path size retryStatus
        return $ FallbackText {..}
  
  complete :: UploadProgress -> IO ()
  complete ProgressBar {..} = Ascii.complete progressBar
  complete _ = pure ()
  
  tick :: UploadProgress -> Int64 -> IO ()
  tick ProgressBar {..} deltaBytes = Ascii.tickN progressBar (fromIntegral deltaBytes)
  tick _ _ = pure ()
  
  fail :: UploadProgress -> IO ()
  fail ProgressBar {..} =
    clearAsciiWith progressBar (uploadFailed path)
  fail FallbackText {path} =
    hPutStr stderr $ uploadFailed path
  
  update :: UploadProgress -> PushRetryStatus -> IO UploadProgress
  update pg@ProgressBar {..} retryStatus = do
    let opts = newProgressBarOptions path size retryStatus
    pg' <- newAsciiProgressBarInPlace progressBar opts
  
    return $ pg {progressBar = pg'}
  update fbt _ = return fbt
  
  newProgressBar :: String -> Int64 -> PushRetryStatus -> IO Ascii.ProgressBar
  newProgressBar path size retryStatus =
    newAsciiProgressBar $ newProgressBarOptions path size retryStatus
  
  newProgressBarOptions :: String -> Int64 -> PushRetryStatus -> Ascii.Options
  newProgressBarOptions path size retryStatus = do
    let hSize = toS $ humanSize $ fromIntegral size
    let barLength =
          uploadTickBar path hSize retryStatus
            & toS
            & T.replace "[:bar]" ""
            & T.replace ":percent" "  0%"
            & T.length
    Ascii.def
      { Ascii.pgTotal = fromIntegral size,
        -- https://github.com/yamadapc/haskell-ascii-progress/issues/24
        Ascii.pgWidth = 20 + barLength,
        Ascii.pgOnCompletion = Just $ uploadComplete path size,
        Ascii.pgFormat = uploadTickBar path hSize retryStatus
      }
  
  retryText :: PushRetryStatus -> String
  retryText PushRetryStatus {retryCount} =
    if retryCount == 0
      then ""
      else color Yellow $ "retry #" <> show retryCount <> " "
  
  uploadComplete :: String -> Int64 -> String
  uploadComplete path size =
    let hSize = toS $ humanSize $ fromIntegral size
     in color Green "✓ " <> path <> " (" <> hSize <> ")"
  
  uploadFailed :: String -> String
  uploadFailed path =
    color Red "✗ " <> path
  
  uploadTickBar :: String -> String -> PushRetryStatus -> String
  uploadTickBar path hSize retryStatus =
    color Blue "[:bar] " <> retryText retryStatus <> toS path <> " (:percent of " <> hSize <> ")"
  
  uploadStartFallback :: String -> Int64 -> PushRetryStatus -> String
  uploadStartFallback path size retryStatus =
    let hSize = toS $ humanSize $ fromIntegral size
     in retryText retryStatus <> "Pushing " <> path <> " (" <> hSize <> ")\n"
  
  -- Internal
  
  newAsciiProgressBar :: Ascii.Options -> IO Ascii.ProgressBar
  newAsciiProgressBar opts = do
    region <- Ascii.openConsoleRegion Ascii.Linear
    newAsciiProgressBarWithRegion opts region
  
  -- | Create a new progress bar in place of an existing one, reusing the console region.
  newAsciiProgressBarInPlace :: Ascii.ProgressBar -> Ascii.Options -> IO Ascii.ProgressBar
  newAsciiProgressBarInPlace pg opts = do
    cancelUpdates pg
    newAsciiProgressBarWithRegion opts (Ascii.pgRegion pg)
  
  -- | Create a new progress bar using the provided console region.
  newAsciiProgressBarWithRegion :: Ascii.Options -> Ascii.ConsoleRegion -> IO Ascii.ProgressBar
  newAsciiProgressBarWithRegion opts region = do
    info <- Ascii.Internal.newProgressBarInfo opts
  
    -- Display initial progress-bar
    pgStr <- Ascii.pgGetProgressStr opts opts <$> Ascii.Internal.getInfoStats info
    Ascii.setConsoleRegion region pgStr
  
    future <- Async.async $ start info
    return $ Ascii.ProgressBar info future region
    where
      start info@Ascii.Internal.ProgressBarInfo {..} = do
        c <- readMVar pgCompleted
        unlessDone c $ do
          n <- readChan pgChannel
          _ <- handleMessage info n
          unlessDone (c + n) $ start info
        where
          unlessDone c action | c < Ascii.pgTotal opts = action
          unlessDone _ _ = do
            let fmt = fromMaybe (Ascii.pgFormat opts) (Ascii.pgOnCompletion opts)
            onCompletion <- Ascii.pgGetProgressStr opts opts {Ascii.pgFormat = fmt} <$> Ascii.Internal.getInfoStats info
            Ascii.setConsoleRegion region onCompletion
  
      handleMessage info n = do
        -- Update the completed tick count
        modifyMVar_ (Ascii.Internal.pgCompleted info) (\c -> return (c + n))
        -- Find and update the current and first tick times:
        stats <- Ascii.Internal.getInfoStats info
        let progressStr = Ascii.Internal.pgGetProgressStr opts opts stats
        Ascii.setConsoleRegion region progressStr
  
  -- | Cancel async updates.
  cancelUpdates :: Ascii.ProgressBar -> IO ()
  cancelUpdates (Ascii.ProgressBar _ future _) = Async.cancel future
  
  -- | Cancel async updates and clear the console region with the given string.
  clearAsciiWith :: Ascii.ProgressBar -> String -> IO ()
  clearAsciiWith pg@(Ascii.ProgressBar _ _ region) str = do
    cancelUpdates pg
    Ascii.setConsoleRegion region str
  
  -- | Display push events as progress bars, deduplicating by path and retry status
  displayPushEvent :: ProgressState -> PushEvent -> IO ()
  displayPushEvent statsRef PushEvent {eventMessage} = do
    stats <- readIORef statsRef
    case eventMessage of
      PushStorePathAttempt path pathSize retryStatus -> do
        case HashMap.lookup path stats of
          Just (progress, prevRetryStatus)
            | prevRetryStatus == retryStatus -> return ()
            | otherwise -> do
                newProgress <- update progress retryStatus
                writeIORef statsRef $ HashMap.insert path (newProgress, retryStatus) stats
          Nothing -> do
            progress <- new stderr (toS path) pathSize retryStatus
            writeIORef statsRef $ HashMap.insert path (progress, retryStatus) stats
      PushStorePathProgress path _ newBytes -> do
        case HashMap.lookup path stats of
          Nothing -> return ()
          Just (progress, _) -> tick progress newBytes
      PushStorePathDone path -> do
        mapM_ (complete . fst) (HashMap.lookup path stats)
      PushStorePathFailed path _ -> do
        mapM_ (fail . fst) (HashMap.lookup path stats)
      _ -> return ()
  