content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE QuasiQuotes #-}
  {-# LANGUAGE Rank2Types #-}
  
  {- This is a standalone module so it shouldn't depend on any CLI state like Env -}
  module Cachix.Client.Push
    ( -- * Pushing a single path
      pushSingleStorePath,
      uploadStorePath,
  
      -- * Push strategy and parameters
      PushParams (..),
      PushSecret (..),
      getAuthTokenFromPushSecret,
      PushStrategy (..),
      defaultWithXzipCompressor,
      defaultWithXzipCompressorWithLevel,
      defaultWithZstdCompressor,
      defaultWithZstdCompressorWithLevel,
  
      -- * Path info
      PathInfo (..),
      newPathInfoFromStorePath,
      newPathInfoFromNarInfo,
  
      -- * Streaming upload
      Push.S3.UploadMultipartResult (..),
      UploadNarDetails (..),
      streamUploadNar,
      streamCopy,
      completeNarUpload,
  
      -- * Narinfo
      narinfoExists,
      newNarInfoCreate,
      queryNarInfoBulk,
  
      -- * Pushing a closure of store paths
      pushClosure,
      computeClosure,
      getMissingPathsForClosure,
      mapConcurrentlyBounded,
    )
  where
  
  import Cachix.API qualified as API
  import Cachix.API.Error
  import Cachix.API.Signing (fingerprint, passthroughHashSink, passthroughHashSinkB16, passthroughSizeSink)
  import Cachix.Client.Exception (CachixException (..))
  import Cachix.Client.Push.S3 qualified as Push.S3
  import Cachix.Client.Retry (retryAll, retryHttp)
  import Cachix.Client.Secrets
  import Cachix.Client.Servant
  import Cachix.Types.BinaryCache qualified as BinaryCache
  import Cachix.Types.MultipartUpload qualified as Multipart
  import Cachix.Types.NarInfoCreate qualified as Api
  import Cachix.Types.NarInfoHash qualified as NarInfoHash
  import Conduit (MonadUnliftIO)
  import Control.Concurrent.Async (mapConcurrently)
  import Control.Concurrent.QSem qualified as QSem
  import Control.Exception.Safe (MonadCatch, MonadMask, throwM)
  import Control.Monad.Trans.Resource (ResourceT)
  import Control.Retry (RetryStatus)
  import Crypto.Sign.Ed25519
  import Data.ByteString.Base64 qualified as B64
  import Data.Conduit
  import Data.Conduit.ByteString (ChunkSize)
  import Data.Conduit.Lzma qualified as Lzma (compress)
  import Data.Conduit.Zstd qualified as Zstd (compress)
  import Data.IORef
  import Data.List.Extra (chunksOf)
  import Data.Set qualified as Set
  import Data.String.Here (iTrim)
  import Data.Text qualified as T
  import Hercules.CNix (StorePath)
  import Hercules.CNix.Std.Set qualified as Std.Set
  import Hercules.CNix.Store (Store)
  import Hercules.CNix.Store qualified as Store
  import Network.HTTP.Types (status401, status404)
  import Nix.NarInfo qualified as NarInfo
  import Protolude hiding (toS)
  import Protolude.Conv
  import Servant.API
  import Servant.Auth ()
  import Servant.Auth.Client
  import Servant.Client.Streaming
  import Servant.Conduit ()
  import System.Nix.Base32 qualified
  import System.Nix.Nar
  
  -- | A secret for authenticating with a cache.
  data PushSecret
    = -- | An auth token. Could be a personal, cache, or agent token.
      PushToken Token
    | -- | An auth token with a signing key for pushing to self-signed caches.
      PushSigningKey Token SigningKey
  
  getAuthTokenFromPushSecret :: PushSecret -> Maybe Token
  getAuthTokenFromPushSecret = \case
    PushToken token -> nullTokenToMaybe token
    PushSigningKey token _ -> nullTokenToMaybe token
    where
      nullTokenToMaybe (Token "") = Nothing
      nullTokenToMaybe token = Just token
  
  -- | Parameters for pushing a closure of store paths, to be passed to 'pushClosure'.
  -- This also contains the parameters for pushing a single path, in 'pushParamStrategy'.
  data PushParams m r = PushParams
    { pushParamsName :: Text,
      pushParamsSecret :: PushSecret,
      -- | An action to run on each closure push attempt.
      -- It receives a list of all paths in the closure and a subset of those paths missing from the cache.
      -- It should return a list of paths to push to the cache.
      pushOnClosureAttempt ::
        -- All paths in the closure
        [StorePath] ->
        -- Paths that are missing in the cache
        [StorePath] ->
        -- Paths to push to the cache
        m [StorePath],
      -- | Hooks into different stages of the single path push process.
      pushParamsStrategy :: StorePath -> PushStrategy m r,
      -- | cachix base url, connection manager, see 'Cachix.Client.URI.defaultCachixBaseUrl', 'Servant.Client.mkClientEnv'
      pushParamsClientEnv :: ClientEnv,
      pushParamsStore :: Store
    }
  
  -- | Parameters for pushing a single store path. See 'pushSingleStorePath'
  data PushStrategy m r = PushStrategy
    { -- | An action to call when a path is already in the cache.
      onAlreadyPresent :: m r,
      -- | An action to run on each push attempt.
      onAttempt :: RetryStatus -> Int64 -> m (),
      -- | A conduit to inspect the NAR stream before compression.
      -- This is useful for tracking progress and computing hashes. The conduit must not modify the stream.
      onUncompressedNARStream :: RetryStatus -> Int64 -> ConduitT ByteString ByteString (ResourceT m) (),
      -- | An action to run when authentication fails.
      on401 :: ClientError -> m r,
      -- | An action to run when an error occurs.
      onError :: ClientError -> m r,
      -- | An action to run after the path is pushed successfully.
      onDone :: m r,
      -- | The compression method to use.
      compressionMethod :: BinaryCache.CompressionMethod,
      -- | The compression level to use.
      compressionLevel :: Int,
      -- | The chunk size to use.
      chunkSize :: ChunkSize,
      -- | The number of chunks to upload concurrently.
      numConcurrentChunks :: Int,
      -- | Whether to mit the deriver from the narinfo.
      omitDeriver :: Bool
    }
  
  defaultWithXzipCompressor :: (MonadIO m) => (ConduitT ByteString ByteString m () -> b) -> b
  defaultWithXzipCompressor = ($ Lzma.compress (Just 2))
  
  defaultWithXzipCompressorWithLevel :: (MonadIO m) => Int -> (ConduitT ByteString ByteString m () -> b) -> b
  defaultWithXzipCompressorWithLevel l = ($ Lzma.compress (Just l))
  
  defaultWithZstdCompressor :: (MonadIO m) => (ConduitT ByteString ByteString m () -> b) -> b
  defaultWithZstdCompressor = ($ Zstd.compress 8)
  
  defaultWithZstdCompressorWithLevel :: (MonadIO m) => Int -> (ConduitT ByteString ByteString m () -> b) -> b
  defaultWithZstdCompressorWithLevel l = ($ Zstd.compress l)
  
  pushSingleStorePath ::
    (MonadMask m, MonadUnliftIO m) =>
    -- | details for pushing to cache
    PushParams m r ->
    -- | store path
    StorePath ->
    -- | r is determined by the 'PushStrategy'
    m r
  pushSingleStorePath pushParams storePath = retryAll $ \retrystatus -> do
    storeHash <- liftIO $ Store.getStorePathHash storePath
    let strategy = pushParamsStrategy pushParams storePath
    res <- liftIO $ narinfoExists pushParams storeHash
    case res of
      Right NoContent -> onAlreadyPresent strategy -- we're done as store path is already in the cache
      Left err
        | isErr err status404 -> uploadStorePath pushParams storePath retrystatus
        | isErr err status401 -> on401 strategy err
        | otherwise -> onError strategy err
  
  narinfoExists :: PushParams m r -> ByteString -> IO (Either ClientError NoContent)
  narinfoExists pushParams storeHash = do
    let cacheName = pushParamsName pushParams
        authToken = getCacheAuthToken (pushParamsSecret pushParams)
    retryHttp $
      (`runClientM` pushParamsClientEnv pushParams) $
        API.narinfoHead
          cachixClient
          authToken
          cacheName
          (NarInfoHash.NarInfoHash (decodeUtf8With lenientDecode storeHash))
  
  getCacheAuthToken :: PushSecret -> Token
  getCacheAuthToken (PushToken token) = token
  getCacheAuthToken (PushSigningKey token _) = token
  
  uploadStorePath ::
    (MonadUnliftIO m, MonadCatch m) =>
    -- | details for pushing to cache
    PushParams m r ->
    StorePath ->
    RetryStatus ->
    -- | r is determined by the 'PushStrategy'
    m r
  uploadStorePath pushParams storePath retrystatus = do
    let store = pushParamsStore pushParams
        strategy = pushParamsStrategy pushParams storePath
  
    -- TODO: storePathText is redundant. Use storePath directly.
    storePathText <- liftIO $ Store.storePathToPath store storePath
  
    -- This should be a noop because storePathText came from a StorePath
    normalized <- liftIO $ Store.followLinksToStorePath store storePathText
    pathInfo <- newPathInfoFromStorePath store normalized
    let narSize = fromIntegral (piNarSize pathInfo)
  
    onAttempt strategy retrystatus narSize
  
    eresult <-
      runConduitRes $
        streamNarIO narEffectsIO (toS storePathText) Data.Conduit.yield
          .| streamUploadNar pushParams storePath narSize retrystatus
  
    case eresult of
      Left e -> onError strategy e
      Right (uploadResult, uploadNarDetails) -> do
        nic <- newNarInfoCreate pushParams storePath pathInfo uploadNarDetails
        completeNarUpload pushParams uploadResult nic
        onDone strategy
  
  -- | Push an entire closure
  --
  -- Note: 'onAlreadyPresent' will be called less often in the future.
  pushClosure ::
    (MonadMask m, MonadUnliftIO m) =>
    -- | Traverse paths, responsible for bounding parallel processing of paths
    --
    -- For example: @'mapConcurrentlyBounded' 4@
    (forall a b. (a -> m b) -> [a] -> m [b]) ->
    PushParams m r ->
    -- | Initial store paths
    [StorePath] ->
    -- | Every @r@ per store path of the entire closure of store paths
    m [r]
  pushClosure traversal pushParams inputStorePaths = do
    (allPaths, missingPaths) <- getMissingPathsForClosure pushParams inputStorePaths
    paths <- pushOnClosureAttempt pushParams allPaths missingPaths
    flip traversal paths $ \storePath ->
      retryAll $ uploadStorePath pushParams storePath
  
  completeNarUpload ::
    (MonadUnliftIO m) =>
    PushParams n r ->
    Push.S3.UploadMultipartResult ->
    Api.NarInfoCreate ->
    m ()
  completeNarUpload pushParams Push.S3.UploadMultipartResult {..} nic = do
    let cacheName = pushParamsName pushParams
        authToken = getCacheAuthToken (pushParamsSecret pushParams)
        clientEnv = pushParamsClientEnv pushParams
  
    -- Complete the multipart upload and upload the narinfo
    let completeMultipartUploadRequest =
          API.completeNarUpload cachixClient authToken cacheName urNarId urUploadId $
            Multipart.CompletedMultipartUpload
              { Multipart.parts = urParts,
                Multipart.narInfoCreate = nic
              }
  
    liftIO $ void $ retryHttp $ withClientM completeMultipartUploadRequest clientEnv escalate
  
  data UploadNarDetails = UploadNarDetails
    { undNarSize :: Integer,
      undNarHash :: Text,
      undFileSize :: Integer,
      undFileHash :: Text
    }
    deriving stock (Eq, Show)
  
  -- | A simplified type for ValidPathInfo.
  -- Can be constructed either from an existing NarInfo, or by querying the Nix store.
  -- Only includes the common fields between remote and local path infos.
  data PathInfo = PathInfo
    { piPath :: FilePath,
      piNarHash :: Text,
      piNarSize :: Integer,
      piDeriver :: Maybe Text,
      piReferences :: Set FilePath
    }
    deriving stock (Eq, Show)
  
  -- | Create a 'PathInfo' for a store path by querying the Nix store.
  newPathInfoFromStorePath :: (MonadIO m) => Store -> StorePath -> m PathInfo
  newPathInfoFromStorePath store storePath = liftIO $ do
    pathInfo <- Store.queryPathInfo store storePath
  
    path <- Store.storePathToPath store storePath
    narHash <- Store.validPathInfoNarHash32 pathInfo
    let narSize = fromIntegral $ Store.validPathInfoNarSize pathInfo
    deriver <-
      Store.validPathInfoDeriver store pathInfo
        >>= mapM Store.getStorePathBaseName
    references <-
      Store.validPathInfoReferences store pathInfo
        >>= mapM Store.getStorePathBaseName
  
    return $
      PathInfo
        { piPath = toS path,
          piNarHash = decodeUtf8 narHash,
          piNarSize = narSize,
          piDeriver = fmap toS deriver,
          piReferences = Set.fromList (fmap toS references)
        }
  
  -- | Create a 'PathInfo' for a store path from an existing NarInfo.
  newPathInfoFromNarInfo :: (Applicative m) => NarInfo.SimpleNarInfo -> m PathInfo
  newPathInfoFromNarInfo narInfo =
    pure $
      PathInfo
        { piPath = NarInfo.storePath narInfo,
          piNarHash = NarInfo.narHash narInfo,
          piNarSize = NarInfo.narSize narInfo,
          piDeriver = NarInfo.deriver narInfo,
          piReferences = NarInfo.references narInfo
        }
  
  -- | A conduit that compresses and streams a NAR to a cache.
  streamUploadNar ::
    (MonadUnliftIO m) =>
    PushParams m r ->
    StorePath ->
    Int64 ->
    RetryStatus ->
    ConduitT ByteString Void (ResourceT m) (Either ClientError (Push.S3.UploadMultipartResult, UploadNarDetails))
  streamUploadNar pushParams storePath storePathSize retrystatus = do
    let cacheName = pushParamsName pushParams
        authToken = getCacheAuthToken (pushParamsSecret pushParams)
        clientEnv = pushParamsClientEnv pushParams
  
    let strategy = pushParamsStrategy pushParams storePath
        withCompressor = case compressionMethod strategy of
          BinaryCache.XZ -> defaultWithXzipCompressorWithLevel (compressionLevel strategy)
          BinaryCache.ZSTD -> defaultWithZstdCompressorWithLevel (compressionLevel strategy)
  
    let uploadOptions =
          Push.S3.UploadMultipartOptions
            { Push.S3.numConcurrentChunks = numConcurrentChunks strategy,
              Push.S3.chunkSize = chunkSize strategy,
              Push.S3.compressionMethod = compressionMethod strategy
            }
  
    narSizeRef <- liftIO $ newIORef 0
    fileSizeRef <- liftIO $ newIORef 0
    narHashRef <- liftIO $ newIORef ("" :: ByteString)
    fileHashRef <- liftIO $ newIORef ("" :: ByteString)
  
    result <- withCompressor $ \compressor ->
      awaitForever Data.Conduit.yield
        .| passthroughSizeSink narSizeRef
        .| passthroughHashSink narHashRef
        .| onUncompressedNARStream strategy retrystatus storePathSize
        .| compressor
        .| passthroughSizeSink fileSizeRef
        .| passthroughHashSinkB16 fileHashRef
        .| Push.S3.uploadMultipart clientEnv authToken cacheName uploadOptions
  
    for result $ \uploadResult -> liftIO $ do
      narSize <- readIORef narSizeRef
      narHash <- ("sha256:" <>) . System.Nix.Base32.encode <$> readIORef narHashRef
      fileSize <- readIORef fileSizeRef
      fileHash <- readIORef fileHashRef
  
      let uploadNarDetails =
            UploadNarDetails
              { undNarSize = narSize,
                undNarHash = narHash,
                undFileSize = fileSize,
                undFileHash = toS fileHash
              }
  
      return (uploadResult, uploadNarDetails)
  
  -- | A conduit that streams an existing NAR to a cache.
  -- Used to copy NARs between caches, e.g. S3 -> Cachix.
  streamCopy ::
    (MonadUnliftIO m) =>
    PushParams m r ->
    StorePath ->
    Int64 ->
    RetryStatus ->
    BinaryCache.CompressionMethod ->
    ConduitT ByteString Void (ResourceT m) (Either ClientError (Push.S3.UploadMultipartResult, UploadNarDetails))
  streamCopy pushParams storePath claimedFileSize retrystatus compressionMethod = do
    let cacheName = pushParamsName pushParams
        authToken = getCacheAuthToken (pushParamsSecret pushParams)
        clientEnv = pushParamsClientEnv pushParams
        strategy = pushParamsStrategy pushParams storePath
  
    let uploadOptions =
          Push.S3.UploadMultipartOptions
            { Push.S3.numConcurrentChunks = numConcurrentChunks strategy,
              Push.S3.chunkSize = chunkSize strategy,
              -- TODO: why is this not part of the strategy?
              Push.S3.compressionMethod = compressionMethod
            }
  
    fileSizeRef <- liftIO $ newIORef 0
    fileHashRef <- liftIO $ newIORef ("" :: ByteString)
  
    result <-
      awaitForever Data.Conduit.yield
        .| onUncompressedNARStream strategy retrystatus claimedFileSize
        .| passthroughSizeSink fileSizeRef
        .| passthroughHashSinkB16 fileHashRef
        .| Push.S3.uploadMultipart clientEnv authToken cacheName uploadOptions
  
    for result $ \uploadResult -> liftIO $ do
      fileSize <- readIORef fileSizeRef
      fileHash <- readIORef fileHashRef
  
      let uploadNarDetails =
            UploadNarDetails
              { undNarSize = 0,
                undNarHash = "",
                undFileSize = fileSize,
                undFileHash = toS fileHash
              }
  
      return (uploadResult, uploadNarDetails)
  
  newNarInfoCreate ::
    (MonadIO m, MonadCatch m) =>
    PushParams n r ->
    StorePath ->
    PathInfo ->
    UploadNarDetails ->
    m Api.NarInfoCreate
  newNarInfoCreate pushParams storePath pathInfo UploadNarDetails {..} = do
    let store = pushParamsStore pushParams
    let strategy = pushParamsStrategy pushParams storePath
    storeDir <- Store.storeDir store
    storePathText <- liftIO $ toS <$> Store.storePathToPath store storePath
  
    when (undNarHash /= piNarHash pathInfo) $
      throwM $
        NarHashMismatch
          [iTrim|
  ${storePathText}: the computed NAR hash doesn't match the hash returned by the Nix store.
  
  Expected: ${undNarHash}
  Got:      ${piNarHash pathInfo}
  
  1. Try repairing the Nix store:
  
     sudo nix-store --verify --repair --check-contents
  
  2. If the issue persists, report it to https://github.com/cachix/cachix/issues
          |]
  
    let deriver =
          fromMaybe "unknown-deriver" $
            if omitDeriver strategy
              then Nothing
              else piDeriver pathInfo
  
    let references = fmap toS $ Set.toList $ piReferences pathInfo
    let fpReferences = fmap (\fp -> toS storeDir <> "/" <> fp) references
  
    let (storeHash, storeSuffix) = splitStorePath storePathText
    let fp = fingerprint storePathText undNarHash undNarSize fpReferences
        sig = case pushParamsSecret pushParams of
          PushToken _ -> Nothing
          PushSigningKey _ signKey -> Just $ toS $ B64.encode $ unSignature $ dsign (signingSecretKey signKey) fp
  
    let nic =
          Api.NarInfoCreate
            { Api.cStoreHash = storeHash,
              Api.cStoreSuffix = storeSuffix,
              Api.cNarHash = undNarHash,
              Api.cNarSize = undNarSize,
              Api.cFileSize = undFileSize,
              Api.cFileHash = undFileHash,
              Api.cReferences = references,
              Api.cDeriver = deriver,
              Api.cSig = sig
            }
  
    escalate $ Api.isNarInfoCreateValid nic
  
    return nic
  
  getMissingPathsForClosure :: (MonadIO m, MonadMask m) => PushParams n r -> [StorePath] -> m ([StorePath], [StorePath])
  getMissingPathsForClosure pushParams inputPaths = do
    let store = pushParamsStore pushParams
    paths <- computeClosure store inputPaths
    queryNarInfoBulk pushParams paths
  
  -- | Get the transitive closure of dependencies
  computeClosure :: (MonadIO m) => Store -> [StorePath] -> m [StorePath]
  computeClosure store inputPaths =
    liftIO $ do
      inputSet <- Std.Set.fromListFP inputPaths
      closure <- Store.computeFSClosure store Store.defaultClosureParams inputSet
      Std.Set.toListFP closure
  
  queryNarInfoBulk :: (MonadIO m, MonadMask m) => PushParams n r -> [StorePath] -> m ([StorePath], [StorePath])
  queryNarInfoBulk pushParams paths = do
    let clientEnv = pushParamsClientEnv pushParams
  
    pathsAndHashes <- liftIO $
      for paths $ \path -> do
        hash' <- decodeUtf8With lenientDecode <$> Store.getStorePathHash path
        pure (hash', path)
    let hashes = fmap fst pathsAndHashes
    let processBatch hashesChunk = retryHttp $ liftIO $ do
          result <-
            API.narinfoBulk
              cachixClient
              (getCacheAuthToken (pushParamsSecret pushParams))
              (pushParamsName pushParams)
              hashesChunk
              `runClientM` clientEnv
          escalate result
    -- 32 bytes of each hash for 1MB limit
    missingHashesList <- concat <$> mapM processBatch (chunksOf 32768 hashes)
  
    let missingHashes = Set.fromList missingHashesList
    let missing = map snd $ filter (\(hash', _path) -> Set.member hash' missingHashes) pathsAndHashes
    return (paths, missing)
  
  mapConcurrentlyBounded :: (Traversable t) => Int -> (a -> IO b) -> t a -> IO (t b)
  mapConcurrentlyBounded bound action items = do
    qs <- QSem.newQSem bound
    let wrapped x = bracket_ (QSem.waitQSem qs) (QSem.signalQSem qs) (action x)
    mapConcurrently wrapped items
  
  -------------------
  -- Private terms --
  splitStorePath :: Text -> (Text, Text)
  splitStorePath storePath =
    (T.take 32 (T.drop 11 storePath), T.drop 44 storePath)
  