content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DuplicateRecordFields #-}
  
  module Cachix.Deploy.Agent where
  
  import Cachix.API.WebSocketSubprotocol qualified as WSS
  import Cachix.Client.Config qualified as Config
  import Cachix.Client.URI (URI)
  import Cachix.Client.URI qualified as URI
  import Cachix.Client.Version (versionNumber)
  import Cachix.Deploy.Deployment qualified as Deployment
  import Cachix.Deploy.Lock qualified as Lock
  import Cachix.Deploy.Log qualified as Log
  import Cachix.Deploy.OptionsParser qualified as CLI
  import Cachix.Deploy.StdinProcess qualified as StdinProcess
  import Cachix.Deploy.Websocket qualified as WebSocket
  import Control.Concurrent.Async qualified as Async
  import Control.Concurrent.Extra (once)
  import Control.Concurrent.MVar qualified as MVar
  import Control.Exception.Safe (onException)
  import Control.Exception.Safe qualified as Safe
  import Control.Retry qualified as Retry
  import Data.Aeson qualified as Aeson
  import Data.IORef
  import Data.String (String)
  import Katip qualified as K
  import Paths_cachix (getBinDir)
  import Protolude hiding (onException, toS, (<.>))
  import Protolude.Conv
  import System.Directory qualified as Directory
  import System.Environment (getEnv, lookupEnv)
  import System.FilePath ((<.>), (</>))
  import System.Posix.Files qualified as Posix.Files
  import System.Posix.Process qualified as Posix
  import System.Posix.Signals qualified as Signals
  import System.Posix.Types qualified as Posix
  import System.Posix.User qualified as Posix.User
  import System.Timeout qualified as Timeout
  
  type ServiceWebSocket = WebSocket.WebSocket (WSS.Message WSS.AgentCommand) (WSS.Message WSS.BackendCommand)
  
  data Agent = Agent
    { name :: Text,
      token :: Text,
      profileName :: Text,
      bootstrap :: Bool,
      host :: URI,
      websocket :: ServiceWebSocket,
      agentState :: IORef (Maybe WSS.AgentInformation),
      logOptions :: Log.Options,
      withLog :: Log.WithLog,
      lockFile :: FilePath,
      pid :: Posix.CPid,
      pidFile :: FilePath
    }
  
  mkAgent :: Log.WithLog -> Log.Options -> Maybe FilePath -> Config.CachixOptions -> CLI.AgentOptions -> Text -> IO Agent
  mkAgent withLog logOptions mlockDirectory cachixOptions agentOptions agentToken = do
    agentState <- newIORef Nothing
    pid <- Posix.getProcessID
    lockDirectory <- maybe Lock.getLockDirectory return mlockDirectory
  
    let host = Config.host cachixOptions
        basename = URI.getHostname host
  
        agentName = CLI.name agentOptions
        profileName = fromMaybe "system" (CLI.profile agentOptions)
  
        port =
          host
            & URI.getScheme
            & URI.getPortFor
            & fromMaybe (URI.Port 80)
  
        websocketOptions =
          WebSocket.Options
            { WebSocket.host = basename,
              WebSocket.port = port,
              WebSocket.path = "/ws",
              WebSocket.useSSL = URI.requiresSSL (URI.getScheme host),
              WebSocket.headers = WebSocket.createHeaders agentName agentToken,
              WebSocket.identifier = agentIdentifier agentName
            }
  
    websocket <- WebSocket.new withLog websocketOptions
  
    let lockFilename = "agent-" <> toS agentName
  
    return $
      Agent
        { name = agentName,
          token = agentToken,
          profileName = profileName,
          bootstrap = CLI.bootstrap agentOptions,
          agentState = agentState,
          pid = pid,
          pidFile = lockDirectory </> lockFilename <.> Lock.pidExtension,
          lockFile = lockDirectory </> lockFilename <.> Lock.lockExtension,
          host = host,
          logOptions = logOptions,
          withLog = withLog,
          websocket = websocket
        }
  
  agentIdentifier :: Text -> Text
  agentIdentifier agentName = unwords [agentName, toS versionNumber]
  
  run :: Config.CachixOptions -> CLI.AgentOptions -> IO ()
  run cachixOptions agentOptions =
    Log.withLog logOptions $ \withLog ->
      logExceptions withLog $ do
        checkUserOwnsHome
  
        -- TODO: show a more helpful error if the token is missing
        -- TODO: show a more helpful error when the token isn't valid
        -- TODO: wrap the token in a newtype or use servant's Token
        agentToken <- toS <$> getEnv "CACHIX_AGENT_TOKEN"
  
        agent <- mkAgent withLog logOptions Nothing cachixOptions agentOptions agentToken
  
        withAgentLock agent $ do
          -- Connect to the backend
          channel <- WebSocket.receive (websocket agent)
          shutdownWebsocket <- connectToService (websocket agent)
  
          -- Shutdown the socket on sigint
          installSignalHandlers shutdownWebsocket
  
          WebSocket.readDataMessages channel $ \message ->
            handleCommand agent (WSS.command message)
    where
      verbosity =
        if Config.verbose cachixOptions
          then Log.Verbose
          else Log.Normal
  
      logOptions =
        Log.Options
          { verbosity = verbosity,
            namespace = "agent",
            environment = "production"
          }
  
  handleCommand :: Agent -> WSS.BackendCommand -> IO ()
  handleCommand agent command =
    case command of
      WSS.AgentRegistered agentInformation -> registerAgent agent agentInformation
      WSS.Deployment deploymentDetails -> launchDeployment agent deploymentDetails
  
  logExceptions :: Log.WithLog -> IO () -> IO ()
  logExceptions withLog action =
    action `catches` [agentHandler, exceptionHandler]
    where
      -- Pretty-print any errors thrown by the agent
      agentHandler =
        Handler $ \(e :: AgentError) -> do
          withLog . K.logLocM K.ErrorS . K.ls $ displayException e
          exitFailure
  
      -- Handle any unexcepted exceptions
      exceptionHandler = Handler $ \(e :: SomeException) -> do
        withLog . K.logLocM K.ErrorS . K.ls $
          unlines
            [ "The agent encountered an exception:",
              toS (displayException e)
            ]
        exitFailure
  
  -- | Acquire a lock for this agent. Skip this step if we're bootstrapping the agent.
  withAgentLock :: Agent -> IO () -> IO ()
  withAgentLock agent action =
    if bootstrap agent
      then action
      else tryToAcquireLock 0
    where
      agentName = name agent
  
      tryToAcquireLock :: Int -> IO ()
      tryToAcquireLock attempts = do
        lock <- Lock.withTryLockAndPid (lockFile agent) (pidFile agent) action
        when (isNothing lock) $
          if attempts >= 5
            then throwIO (AgentAlreadyRunning agentName)
            else do
              threadDelay (3 * 1000 * 1000)
              tryToAcquireLock (attempts + 1)
  
  installSignalHandlers :: IO () -> IO ()
  installSignalHandlers shutdown = do
    mainThreadId <- myThreadId
    -- Track Ctrl+C attempts
    interruptRef <- newIORef False
  
    let safeShutdown = do
          -- Timeout after 30 seconds to ensure we don't hang
          void $ Timeout.timeout (30 * 1000 * 1000) shutdown
  
    -- Handle SIGTERM (systemd): always shutdown and exit immediately
    _ <-
      Signals.installHandler
        Signals.sigTERM
        ( Signals.Catch $ do
            safeShutdown
            throwTo mainThreadId ExitSuccess
        )
        Nothing
  
    -- Handle SIGINT (interactive Ctrl+C): first attempt initiates shutdown, second exits
    _ <-
      Signals.installHandler
        Signals.sigINT
        ( Signals.Catch $ do
            isSecondInterrupt <- liftIO $ atomicModifyIORef' interruptRef (True,)
            if isSecondInterrupt
              then throwTo mainThreadId ExitSuccess
              else safeShutdown
        )
        Nothing
  
    return ()
  
  registerAgent :: Agent -> WSS.AgentInformation -> IO ()
  registerAgent Agent {agentState, withLog} agentInformation = do
    withLog $ K.logLocM K.InfoS "Agent registered."
    atomicWriteIORef agentState (Just agentInformation)
  
  launchDeployment :: Agent -> WSS.DeploymentDetails -> IO ()
  launchDeployment agent@Agent {..} deploymentDetails = do
    agentRegistered <- readIORef agentState
  
    case agentRegistered of
      -- TODO: the agent should either not exist before we register or
      -- we should re-register here as a precaution.
      Nothing -> pure ()
      Just agentInformation -> do
        binDir <- toS <$> getBinDir
        exitCode <-
          StdinProcess.readProcess (binDir <> "/.cachix-deployment") [] $
            toS . Aeson.encode $
              Deployment.Deployment
                { Deployment.agentName = name,
                  Deployment.agentToken = token,
                  Deployment.profileName = profileName,
                  Deployment.host = host,
                  Deployment.deploymentDetails = deploymentDetails,
                  Deployment.agentInformation = agentInformation,
                  Deployment.logOptions = logOptions
                }
  
        when
          (bootstrap && exitCode == ExitSuccess)
          (verifyBootstrapSuccess agent)
  
  verifyBootstrapSuccess :: Agent -> IO ()
  verifyBootstrapSuccess agent@(Agent {name, withLog}) = do
    withLog . K.logLocM K.InfoS . K.ls $
      unwords ["Waiting for another agent to take over..."]
  
    magentPid <- waitForAgent (Retry.limitRetries 60 <> Retry.constantDelay 1000) agent
  
    case magentPid of
      Just pid -> do
        withLog . K.logLocM K.InfoS . K.ls $
          unwords ["Found an active agent for", name, "with PID " <> show pid <> ".", "Exiting."]
        exitSuccess
      Nothing -> do
        withLog . K.logLocM K.InfoS . K.ls $
          unwords ["Cannot find an active agent for", name <> ".", "Waiting for more deployments."]
  
  waitForAgent :: Retry.RetryPolicyM IO -> Agent -> IO (Maybe Posix.CPid)
  waitForAgent retryPolicy agent = do
    Retry.retrying
      retryPolicy
      (const $ pure . isNothing)
      (const $ findActiveAgent agent)
  
  -- The PID might be stale in rare cases. Only use this for diagnostics.
  findActiveAgent :: Agent -> IO (Maybe Posix.CPid)
  findActiveAgent Agent {pidFile, lockFile} = do
    Safe.handleAny (const $ pure Nothing) $ do
      lock <- Lock.withTryLock lockFile (pure ())
      -- The lock must be held by another process
      guard (isNothing lock)
  
      -- We should have a PID file
      mpid <- Lock.readPidFile pidFile
      guard (isJust mpid)
  
      return mpid
  
  -- | Asynchronously open and maintain a websocket connection to the backend for
  -- sending deployment progress updates.
  connectToService :: ServiceWebSocket -> IO (IO ())
  connectToService websocket = do
    close <- MVar.newEmptyMVar
  
    thread <- Async.async $
      WebSocket.runConnection websocket $ do
        WebSocket.handleJSONMessages websocket (MVar.readMVar close)
  
    -- Block until the initial connection is established
    void $ MVar.readMVar (WebSocket.connection websocket)
  
    once $ do
      void $ MVar.tryPutMVar close ()
      void $ Timeout.timeout (5 * 1000 * 1000) (Async.wait thread)
  
  -- | Fetch the home directory and verify that the owner matches the current user.
  -- Throws either 'NoHomeFound' or 'UserDoesNotOwnHome'.
  checkUserOwnsHome :: IO ()
  checkUserOwnsHome = do
    home <- Directory.getHomeDirectory `onException` throwIO NoHomeFound
    stat <- Posix.Files.getFileStatus home
    userId <- Posix.User.getEffectiveUserID
  
    when (userId /= Posix.Files.fileOwner stat) $ do
      userName <- Posix.User.userName <$> Posix.User.getUserEntryForID userId
      sudoUser <- lookupEnv "SUDO_USER"
      throwIO $ UserDoesNotOwnHome userName sudoUser home
  
  data AgentError
    = -- | An agent with the same name is already running.
      AgentAlreadyRunning Text
    | -- | No home directory.
      NoHomeFound
    | -- | Safeguard against creating root-owned files in user directories.
      -- This is an issue on macOS, where, by default, sudo does not reset $HOME.
      UserDoesNotOwnHome
        -- | The current user name
        String
        -- | The sudo user name, if any
        (Maybe String)
        -- | The home directory
        FilePath
    deriving (Show)
  
  instance Exception AgentError where
    displayException = \case
      AgentAlreadyRunning agentName -> toS $ unwords ["The agent", agentName, "is already running."]
      NoHomeFound -> "Could not find the user’s home directory. Make sure to set the $HOME variable."
      UserDoesNotOwnHome userName sudoUser home ->
        if isJust sudoUser
          then toS $ unlines [warningMessage, suggestSudoFlagH]
          else toS warningMessage
        where
          warningMessage = "The current user (" <> toS userName <> ") does not own the home directory (" <> toS home <> ")"
          suggestSudoFlagH = "Try running the agent with `sudo -H`."
  