content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE DuplicateRecordFields #-}
  
  module Cachix.Deploy.Activate where
  
  import Cachix.API.WebSocketSubprotocol qualified as WSS
  import Cachix.Client.InstallationMode qualified as InstallationMode
  import Cachix.Client.NetRc qualified as NetRc
  import Cachix.Client.URI (URI)
  import Cachix.Client.URI qualified as URI
  import Cachix.Deploy.Log qualified as Log
  import Cachix.Types.BinaryCache qualified as BinaryCache
  import Cachix.Types.Permission (Permission (..))
  import Data.Aeson qualified as Aeson
  #if MIN_VERSION_aeson(2,0,0)
  import qualified Data.Aeson.KeyMap as HM
  #else
  import qualified Data.HashMap.Strict as HM
  #endif
  import Data.Conduit.Combinators qualified as Conduit
  import Data.Conduit.Process qualified as Conduit
  import Data.Vector qualified as Vector
  import Protolude hiding (log, toS)
  import Protolude.Conv (toS)
  import Servant.Auth.Client (Token (..))
  import System.Directory qualified as Directory
  import System.FilePath ((</>))
  import System.IO.Temp (withSystemTempDirectory)
  import System.Process
  import Prelude (String)
  
  data Status
    = Success
    | Failure FailureReason
    | Rollback FailureReason
    deriving (Show, Exception)
  
  data FailureReason
    = NetworkTestFailure
    | RollbackScriptExitFailure
    | RollbackScriptUnexpectedError IOException
    | ShellCommandFailure {command :: String, exitCode :: Int}
    | UnexpectedError SomeException
    deriving (Show)
  
  instance Exception FailureReason where
    displayException = \case
      NetworkTestFailure -> "Cannot connect back to Cachix Deploy after activating the new deployment"
      RollbackScriptExitFailure -> "The rollback script returned a non-zero exit code"
      RollbackScriptUnexpectedError e -> "Cannot run rollback script: " <> displayException e
      ShellCommandFailure {command, exitCode} ->
        toS $
          unwords
            [ "Failed to run " <> toS command,
              show exitCode
            ]
      UnexpectedError e ->
        toS $
          unwords
            [ "The deployment failed with an unexpected error:",
              toS (displayException e)
            ]
  
  downloadStorePaths ::
    -- | Logging context
    Log.LogStream ->
    -- | Deployment details
    WSS.DeploymentDetails ->
    -- | Binary cache args
    [String] ->
    IO ()
  downloadStorePaths logStream deploymentDetails cacheArgs = do
    -- Download the store path from the binary cache
    -- TODO: add GC root so it's preserved for the next command
    runShell logStream "nix-store" (["-r", toS storePath] <> cacheArgs)
  
    -- Download the rollback script, if provided
    for_ (WSS.rollbackScript deploymentDetails) $ \script ->
      runShell logStream "nix-store" (["-r", toS script] <> cacheArgs)
    where
      storePath = WSS.storePath deploymentDetails
  
  type RollbackAction = IO ()
  
  -- | Activate the profile and return a rollback action if available.
  activate ::
    -- | Logging context
    Log.LogStream ->
    -- | Profile name
    Text ->
    -- | Store path to activate
    FilePath ->
    -- | Returns a rollback action if available
    IO (Maybe RollbackAction)
  activate logStream profileName storePath = do
    (profilePath, activationScripts) <- getActivationScript profileName storePath
    previousProfilePath <- toStorePath profilePath
  
    -- Activate the configuration
    -- TODO: Check with Domen whether we can exit early here
    forM_ activationScripts $ uncurry (runShell logStream)
  
    pure $ Just rollback <*> previousProfilePath
    where
      toStorePath profilePath = do
        profileExists <- Directory.doesPathExist profilePath
        if profileExists
          then Just <$> Directory.canonicalizePath profilePath
          else pure Nothing
  
      -- We can't use '--rollback' because it just selects the next generation
      -- down from our deployment, which is not necessarily the generation that
      -- was previously active.
      rollback = void . activate logStream profileName
  
  type Command = (String, [String])
  
  getActivationScript :: Text -> FilePath -> IO (FilePath, [Command])
  getActivationScript profile storePath = do
    let checkPath p = Directory.doesPathExist $ toS storePath </> p
    isNixOS <- checkPath "nixos-version"
    isNixDarwin <- checkPath "darwin-version"
    isNixDarwinJSON <- checkPath "darwin-version.json"
    isHomeManager <- checkPath "hm-version"
    user <- InstallationMode.getUser
    let systemProfileDir = "/nix/var/nix/profiles"
    let perUserProfileDir = systemProfileDir </> "per-user" </> toS user
    let mkProfilePath profileBaseDir defaultProfile =
          profileBaseDir </> if profile == "" then defaultProfile else toS profile
    -- Sets the new profile. This is needed for NixOS and nix-darwin. The Home
    -- Manager activation script does the profile setting by itself.
    -- TODO: document what happens if the wrong user is used for the agent
    let setNewProfile profilePath =
          ("nix-env", ["-p", profilePath, "--set", storePath])
    return $ case (isNixOS, isNixDarwin || isNixDarwinJSON, isHomeManager) of
      (True, _, _) ->
        let profilePath = mkProfilePath systemProfileDir "system"
         in ( profilePath,
              [ setNewProfile profilePath,
                (toS storePath </> "bin/switch-to-configuration", ["switch"])
              ]
            )
      (_, True, _) ->
        -- https://github.com/LnL7/nix-darwin/blob/master/pkgs/nix-tools/darwin-rebuild.sh
        let profilePath = mkProfilePath systemProfileDir "system-profiles/system"
         in ( profilePath,
              [ ("mkdir", ["-p", "-m", "0755", "/nix/var/nix/profiles/system-profiles"]),
                setNewProfile profilePath,
                (toS storePath </> "sw/bin/darwin-rebuild", ["activate"])
              ]
            )
      (_, _, True) ->
        ( mkProfilePath perUserProfileDir "home-manager",
          [(toS storePath </> "activate", [])]
        )
      (_, _, _) ->
        let profilePath = mkProfilePath systemProfileDir "system"
         in (profilePath, [setNewProfile profilePath])
  
  -- TODO: send errors as well
  -- TODO: fix either/maybe types
  getClosureSize :: [String] -> Text -> IO (Either Text (Maybe Int64))
  getClosureSize cacheArgs storePath = do
    (exitCode, pathInfoJSON, nixError) <- readProcessWithExitCode "nix" (cacheArgs <> ["--extra-experimental-features", "nix-command", "path-info", "-S", "--json", toS storePath]) ""
    case exitCode of
      ExitFailure _ -> pure $ Left (toS nixError)
      ExitSuccess -> pure $ Right $ Aeson.decode (toS pathInfoJSON) >>= extractClosureSize
  
  extractClosureSize :: Aeson.Value -> Maybe Int64
  extractClosureSize (Aeson.Array vector) = case Vector.toList vector of
    [Aeson.Object obj] -> case HM.lookup "closureSize" obj of
      Just (Aeson.Number num) -> Just $ floor num
      _ -> Nothing
    _ -> Nothing
  extractClosureSize _ = Nothing
  
  -- TODO: don't create tmpfile for public caches
  withCacheArgs :: URI -> WSS.AgentInformation -> Text -> ([String] -> IO a) -> IO a
  withCacheArgs host agentInfo agentToken m =
    withSystemTempDirectory "netrc" $ \dir -> do
      let filepath = dir </> "netrc"
      args <- case WSS.cache agentInfo of
        Just cache -> do
          let cacheName = WSS.cacheName cache
          let cacheURI = URI.appendSubdomain cacheName host
          -- TODO: ugh
          let bc =
                BinaryCache.BinaryCache
                  { BinaryCache.name = "",
                    BinaryCache.uri = URI.serialize cacheURI,
                    BinaryCache.publicSigningKeys = [],
                    BinaryCache.isPublic = WSS.isPublic cache,
                    BinaryCache.githubUsername = "",
                    BinaryCache.permission = Read,
                    BinaryCache.preferredCompressionMethod = BinaryCache.XZ
                  }
          NetRc.add (Token (toS agentToken)) [bc] filepath
          return $ cachesArgs cache <> ["--option", "netrc-file", filepath]
        Nothing ->
          return []
      m args
    where
      cachesArgs :: WSS.Cache -> [String]
      cachesArgs cache =
        let cacheName = WSS.cacheName cache
            cacheURI = URI.appendSubdomain cacheName host
            hostname = (URI.hostBS . URI.getHostname) cacheURI
            officialCache = "cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY="
            substituters = ["--option", "extra-substituters", URI.serialize cacheURI]
            noNegativeCaching = ["--option", "narinfo-cache-negative-ttl", "0"]
            sigs = ["--option", "trusted-public-keys", officialCache <> " " <> toS hostname <> "-1:" <> toS (WSS.publicKey cache)]
         in substituters ++ sigs ++ noNegativeCaching
  
  runShell :: Log.LogStream -> FilePath -> [String] -> IO ()
  runShell logStream cmd args = runShellWithExitCode logStream cmd args >>= handleError
    where
      handleError (ExitFailure exitCode) = throwIO $ ShellCommandFailure {command = cmd, exitCode}
      handleError ExitSuccess = pure ()
  
  runShellWithExitCode :: Log.LogStream -> FilePath -> [String] -> IO ExitCode
  runShellWithExitCode logStream cmd args = do
    Log.streamLine logStream $ "$ " <> toS cmd <> " " <> toS (unwords $ fmap toS args)
    (exitCode, _, _) <- Conduit.sourceProcessWithStreams (proc cmd args) Conduit.sinkNull logStream logStream
    pure exitCode
  