{-# LANGUAGE DuplicateRecordFields #-}

module Main
  ( main,
  )
where

import Cachix.API.Error (escalateAs)
import Cachix.API.WebSocketSubprotocol qualified as AgentInformation (AgentInformation (..))
import Cachix.API.WebSocketSubprotocol qualified as DeploymentDetails (DeploymentDetails (..))
import Cachix.API.WebSocketSubprotocol qualified as WSS
import Cachix.Client.URI qualified as URI
import Cachix.Deploy.Activate qualified as Activate
import Cachix.Deploy.Agent qualified as Agent
import Cachix.Deploy.Deployment (Deployment (..))
import Cachix.Deploy.Lock qualified as Lock
import Cachix.Deploy.Log qualified as Log
import Cachix.Deploy.Websocket qualified as WebSocket
import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM.TMQueue qualified as TMQueue
import Control.Exception.Safe qualified as Safe
import Data.Aeson qualified as Aeson
import Data.Conduit.TQueue qualified as Conduit
import Data.Time.Clock (getCurrentTime)
import Data.UUID qualified as UUID
import Data.UUID.V4 qualified as UUID
import GHC.IO.Encoding
import Katip qualified as K
import Network.WebSockets qualified as WS
import Protolude hiding (toS)
import Protolude.Conv
import System.IO (BufferMode (..), hSetBuffering)

lockFilenameFrom :: Text -> FilePath
lockFilenameFrom agentName = "deployment-" <> toS agentName

-- | Activate the new deployment.
--
-- If the target profile is already locked by another deployment, exit
-- immediately and rely on the backend to reschedule.
main :: IO ()
main = do
  setLocaleEncoding utf8
  setFileSystemEncoding utf8
  hSetBuffering stdout LineBuffering
  hSetBuffering stderr LineBuffering

  deployment@Deployment
    { agentName,
      agentToken,
      host,
      logOptions,
      deploymentDetails
    } <-
    escalateAs (FatalError . toS) . Aeson.eitherDecode . toS =<< getContents

  let deploymentID = DeploymentDetails.id deploymentDetails
  let headers = WebSocket.createHeaders agentName agentToken
  let port = fromMaybe (URI.Port 80) (URI.getPortFor (URI.getScheme host))
  let logWebsocketOptions =
        WebSocket.Options
          { WebSocket.host = URI.getHostname host,
            WebSocket.port = port,
            WebSocket.path = "/api/v1/deploy/log/" <> UUID.toText deploymentID,
            WebSocket.useSSL = URI.requiresSSL (URI.getScheme host),
            WebSocket.headers = headers,
            WebSocket.identifier = Agent.agentIdentifier agentName
          }
  let serviceWebsocketOptions =
        WebSocket.Options
          { WebSocket.host = URI.getHostname host,
            WebSocket.port = port,
            WebSocket.path = "/ws-deployment",
            WebSocket.useSSL = URI.requiresSSL (URI.getScheme host),
            WebSocket.headers = headers,
            WebSocket.identifier = Agent.agentIdentifier agentName
          }

  lockFilePath <- Lock.newLockFilePath (lockFilenameFrom agentName)

  Log.withLog logOptions $ \withLog ->
    void . Lock.withTryLock lockFilePath $ do
      -- Open a connection to logging stream
      (logQueue, loggingThread) <- runLogStream withLog logWebsocketOptions

      -- Open a connection to Cachix and block until it's ready.
      service <- WebSocket.new withLog serviceWebsocketOptions
      shutdownService <- Agent.connectToService service

      deploy withLog deployment service (Conduit.sinkTMQueue logQueue)
        `finally` do
          withLog $ K.logLocM K.DebugS "Cleaning up websocket connections"
          atomically $ TMQueue.closeTMQueue logQueue
          shutdownService
          Async.wait loggingThread

-- | Run the deployment commands
deploy ::
  -- | Logging context
  Log.WithLog ->
  -- | Deployment information passed from the agent
  Deployment ->
  Agent.ServiceWebSocket ->
  -- | Logging Websocket connection
  Log.LogStream ->
  IO ()
deploy withLog Deployment {..} service logStream = do
  withLog $ K.logLocM K.InfoS $ K.ls $ "Deploying #" <> deploymentIndex <> ": " <> storePath

  activationStatus <- handleAsActivationStatus $
    Activate.withCacheArgs host agentInformation agentToken $ \cacheArgs -> do
      startDeployment Nothing

      Activate.downloadStorePaths logStream deploymentDetails cacheArgs

      -- Read the closure size and report
      --
      -- TODO: query the remote store to get the size before downloading (and
      -- possibly running out of disk space)
      closureSize <- fromRight Nothing <$> Activate.getClosureSize cacheArgs storePath
      when (isJust closureSize) $ startDeployment closureSize

      rollbackAction <- Activate.activate logStream profileName (toS storePath)

      -- Run tests on the new deployment
      testResults <- handleAsFailureReason $ do
        -- Run a basic network test against the backend
        pong <- WebSocket.waitForPong 10 service
        when (isNothing pong) $ throwIO Activate.NetworkTestFailure

        -- Run the optional rollback script
        for (WSS.rollbackScript deploymentDetails) $ \rollbackScript -> do
          Log.streamLine logStream "Running rollback script."
          rollbackScriptResult <- Safe.tryIO $ Activate.runShellWithExitCode logStream (toS rollbackScript) []
          case rollbackScriptResult of
            Right ExitSuccess -> pure ()
            Right (ExitFailure _) -> throwIO Activate.RollbackScriptExitFailure
            Left e -> throwIO (Activate.RollbackScriptUnexpectedError e)

      -- Roll back if any of the tests have failed
      case testResults of
        Right _ -> pure Activate.Success
        Left testErrors ->
          case rollbackAction of
            Just rollback -> do
              Log.streamLine logStream "Deployment failed, rolling back ..."
              rollback
              throwIO (Activate.Rollback testErrors)
            Nothing -> do
              Log.streamLine logStream "Skipping rollback as this is the first deployment."
              throwIO (Activate.Failure testErrors)

  case activationStatus of
    Activate.Failure e -> logDeploymentFailed e
    Activate.Rollback e -> logDeploymentFailed e
    Activate.Success -> do
      -- NOTE: the activate command uses this message to detect the end of the log
      Log.streamLine logStream "Successfully activated the deployment."
      withLog $ K.logLocM K.InfoS $ K.ls $ "Deployment #" <> deploymentIndex <> " finished."

  endDeployment activationStatus
  where
    storePath = DeploymentDetails.storePath deploymentDetails
    deploymentID = DeploymentDetails.id deploymentDetails
    deploymentIndex = show (DeploymentDetails.index deploymentDetails)

    handleAsActivationStatus :: IO Activate.Status -> IO Activate.Status
    handleAsActivationStatus action =
      action
        `Safe.catches` [ Safe.Handler (\(e :: Activate.Status) -> pure e),
                         Safe.Handler (\(e :: SomeException) -> pure $ Activate.Failure (Activate.UnexpectedError e))
                       ]

    handleAsFailureReason :: IO a -> IO (Either Activate.FailureReason a)
    handleAsFailureReason action =
      fmap Right action
        `Safe.catches` [ Safe.Handler (\(e :: Activate.FailureReason) -> pure (Left e)),
                         Safe.Handler (\(e :: SomeException) -> pure $ Left (Activate.UnexpectedError e))
                       ]

    logDeploymentFailed e = do
      Log.streamLine logStream $
        toS $
          unwords
            -- NOTE: the activate command uses this message to detect the end of the log
            [ "Failed to activate the deployment.",
              toS $ displayException e
            ]
      withLog $ K.logLocM K.InfoS $ K.ls $ "Deploying #" <> deploymentIndex <> " failed."

    startDeployment :: Maybe Int64 -> IO ()
    startDeployment closureSize = do
      now <- getCurrentTime
      msg <-
        createMessage $
          WSS.DeploymentStarted
            { WSS.id = deploymentID,
              WSS.time = now,
              WSS.closureSize = closureSize
            }
      WebSocket.send service (WebSocket.DataMessage msg)

    endDeployment :: Activate.Status -> IO ()
    endDeployment status = do
      let hasSucceeded =
            case status of
              Activate.Success -> True
              _ -> False
      now <- getCurrentTime
      msg <-
        createMessage $
          WSS.DeploymentFinished
            { WSS.id = deploymentID,
              WSS.time = now,
              WSS.hasSucceeded = hasSucceeded
            }
      WebSocket.send service (WebSocket.DataMessage msg)

    createMessage :: WSS.AgentCommand -> IO (WSS.Message WSS.AgentCommand)
    createMessage command = do
      uuid <- UUID.nextRandom
      return $
        WSS.Message
          { WSS.method = method,
            WSS.command = command,
            WSS.id = uuid,
            WSS.agent = Just $ AgentInformation.id agentInformation
          }
      where
        -- TODO: move to WSS
        method = case command of
          WSS.DeploymentStarted {} -> "DeploymentStarted"
          WSS.DeploymentFinished {} -> "DeploymentFinished"

-- Log

-- TODO: prepend katip-like format to each line
-- TODO: Re-use the WebSocket module here (without ping?)
runLogStream ::
  Log.WithLog ->
  WebSocket.Options ->
  -- | Returns a queue for writing messages and the thread handle
  IO (TMQueue.TMQueue ByteString, Async.Async ())
runLogStream withLog options = do
  queue <- TMQueue.newTMQueueIO
  thread <- Async.async $
    WebSocket.reconnectWithLog withLog $
      WebSocket.runClientWith options WS.defaultConnectionOptions $ \connection ->
        Log.streamLog withLog connection queue
          `finally` WebSocket.waitForGracefulShutdown connection
  return (queue, thread)