{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE ApplicativeDo #-}

module Glean.Write.Options
  ( sendQueueOptions
  , sendAndRebaseQueueOptions
  , writerOptions
  ) where

import Control.Exception
import qualified Data.ByteString as B
import Data.Default (def)
import qualified Options.Applicative as O
import System.IO
import Text.Printf

import Thrift.Protocol.Compact
import Util.Log

import Glean.Write.Async (WriterSettings(..), WriterEvent(..))
import Glean.Write.SendQueue (SendQueueSettings(..))
import Glean.Write.SendAndRebaseQueue (SendAndRebaseQueueSettings(..))
import Glean.Util.RetryChannelException


sendQueueOptions :: O.Parser SendQueueSettings
sendQueueOptions = do
  sendQueueThreads <- O.option O.auto $
    O.long "sender-threads"
    <> O.metavar "N"
    <> O.value 1
    <> O.showDefault
    <> O.help "number of concurrent sender threads"
  sendQueueMaxMemory <- O.option O.auto $
    O.long "max-send-queue-size"
    <> O.metavar "N"
    <> O.value 2000000000
    <> O.showDefault
    <> O.help "maximum size of send queue (in bytes)"
  sendQueueMaxBatches <- O.option O.auto $
    O.long "max-send-queue-batches"
    <> O.metavar "N"
    <> O.value 1024
    <> O.showDefault
    <> O.help "maximum number of batches in send queue"
  sendQueueRetry <- fmap toRetryPolicy $ O.option O.auto $
    O.long "send-retries"
    <> O.metavar "N"
    <> O.value 20
    <> O.showDefault
    <> O.help "number of retries"
  return def
    { sendQueueThreads
    , sendQueueMaxMemory
    , sendQueueMaxBatches
    , sendQueueRetry
    }
  where
  toRetryPolicy n = defaultRetryPolicy { maxRetries = n }


sendAndRebaseQueueOptions :: O.Parser SendAndRebaseQueueSettings
sendAndRebaseQueueOptions = do
  sendAndRebaseQueueSendQueueSettings <- sendQueueOptions
  sendAndRebaseQueueFactCacheSize <- O.option O.auto $
    O.long "fact-cache-size"
    <> O.metavar "N"
    <> O.value 2000000000
    <> O.showDefault
    <> O.help "size of the cache for sent facts"
  sendAndRebaseQueueFactBufferSize <- O.option O.auto $
    O.long "fact-buffer-size"
    <> O.metavar "N"
    <> O.value (200*1000*1000)
    <> O.showDefault
    <> O.help "size of the buffer for unsent facts per sender"
  return SendAndRebaseQueueSettings
    { sendAndRebaseQueueSendQueueSettings
    , sendAndRebaseQueueFactCacheSize
    , sendAndRebaseQueueFactBufferSize
    , sendAndRebaseQueueSenders =
        sendQueueThreads sendAndRebaseQueueSendQueueSettings
    , sendAndRebaseQueueAllowRemoteReferences = False
    , sendAndRebaseQueueStats = Nothing
    }


writerOptions :: O.Parser WriterSettings
writerOptions = do
  writerMaxSize <- O.option O.auto $
    O.long "max-batch-size"
    <> O.metavar "N"
    <> O.value 30000000
    <> O.showDefault
    <> O.help "maximum size of Thrift batch (in bytes)"
  writerDump <- O.optional $ O.strOption $
    O.long "dump-batches"
    <> O.metavar "DIR"
    <> O.help ("also dump all batches in binary format to DIR. " <>
      "The files can be used with 'glean write --file-format=binary'")
  return def
    { writerMaxSize = writerMaxSize
    , writerLog = logBatch writerDump
    }
  where
    logBatch dump (WriterPushing batch)
      | Just dir <- dump = do
        let serialized = serializeCompact batch
        bracket
          (openTempFile dir "batch.bin")
          (hClose . snd)
          $ \(f,h) -> do
            logInfo $ printf "dumping batch (%d bytes) to %s"
              (B.length serialized) f
            B.hPut h serialized
    logBatch _ _ = return ()