content: {-
    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 TypeApplications, ConstraintKinds, UndecidableInstances #-}
  
  module Glean.Util.ShellPrint
    ( DbVerbosity(..)
    , StatsFormatOpts(..)
    , ShellFormat(..)
    , ShellPrint
    , ShellPrintFormat(..)
    , PrintOpts(..)
    , getTerminalWidth
    , hPutShellPrintLn
    , putShellPrintLn
    , shellFormatOpt
    , shellPrint
    , withFormatOpts
    ) where
  
  import Prelude hiding ((<>))
  
  import Control.Monad
  import Data.Aeson
  import qualified Data.Aeson as J
  import qualified Data.Aeson.Encode.Pretty as J
  import qualified Data.ByteString.Lazy.Char8 as BS
  import Data.Char
  import Data.Int
  import Data.Maybe
  import Data.Text (Text)
  import qualified Data.Text as Text
  import qualified Data.HashMap.Strict as HashMap
  import Data.Time.Clock.POSIX
  import Data.List hiding (span)
  import Data.Void
  import Options.Applicative as O
  import Text.Printf hiding (parseFormat)
  import System.Exit
  import System.IO
  import System.Process
  import System.Timeout
  import Compat.Prettyprinter
  import Prettyprinter.Render.Terminal as Pretty
  import Util.Aeson
  import Util.Control.Exception (catchAll)
  import Util.TimeSec
  import Util.Timing
  
  import Glean.Backend.Types (dbShard)
  import qualified Glean.Types as Thrift
  import Glean.Repo.Text (showRepo)
  
  data ShellPrintFormat
    = TTY
    | PlainText
    | Json
    | CompactJson
    deriving (Eq,Show)
  
  data Context = Context
    { ctxFormat :: ShellPrintFormat
    , ctxNow :: Time
    }
  
  parseFormat :: String -> Maybe ShellPrintFormat
  parseFormat "tty" = Just TTY
  parseFormat "plain" = Just PlainText
  parseFormat "json" = Just Json
  parseFormat "compact-json" = Just CompactJson
  parseFormat _ = Nothing
  
  shellFormatOpt :: Parser (Maybe ShellPrintFormat)
  shellFormatOpt =
    O.optional $ O.option (maybeReader parseFormat) $ mconcat
      [ O.long "format"
      , O.metavar "(tty|plain|json|compact-json)"
      , O.help "Output format"
      ]
  
  data PrintOpts = PrintOpts
    { poFormat :: ShellPrintFormat
    , poNow :: Time
    , poWidth :: Maybe PageWidth
    }
  
  type ShellPrint a = ShellFormat () a
    -- we might prefer this to be "forall o . ShellFormat o a" but that
    -- requires QuantifiedConstraints which isn't available until GHC 8.6.x
    -- and we're keeping GHC 8.4.x compatibility for now.
  
  putShellPrintLn
    :: ShellPrint a
    => Maybe ShellPrintFormat -> a -> IO ()
  putShellPrintLn = hPutShellPrintLn stdout
  
  hPutShellPrintLn
    :: ShellPrint a
    => Handle
    -> Maybe ShellPrintFormat
    -> a
    -> IO ()
  hPutShellPrintLn outh opt x = do
    tty <- hIsTerminalDevice stdout
    now <- utcTimeToPOSIXSeconds <$> getCurrentTime
    width <- fromMaybe 80 <$> getTerminalWidth
    let
      t0 = Time (round now)
      format = fromMaybe (if tty then TTY else PlainText) opt
      opts = PrintOpts
        { poFormat = format
        , poNow = t0
        , poWidth = Just $ AvailablePerLine width 1
        }
    shellPrint outh opts x
  
  shellPrint
    :: forall a. ShellFormat () a
    => Handle -> PrintOpts -> a -> IO ()
  shellPrint handle PrintOpts{..} x =
    Pretty.renderIO handle sds
    where
      sds =
        layoutPretty layout $ doc <> hardline
      doc = case poFormat of
        CompactJson ->
          pretty $
          BS.unpack $
          J.encode $
          shellFormatJson context () x
        Json ->
          pretty $
          BS.unpack $
          J.encodePretty $
          shellFormatJson context () x
        TTY ->
          shellFormatText context () x
        PlainText ->
          unAnnotate $ shellFormatText context () x
      context = Context
        { ctxFormat = poFormat
        , ctxNow = poNow
        }
      layout = LayoutOptions
        { layoutPageWidth = fromMaybe Unbounded poWidth
        }
  
  type Ann = AnsiStyle
  
  -- | Format data of type 'a' with format options of type 'o'
  class ShellFormat o a where
      shellFormatText
        :: Context -> o -> a -> Doc Ann
      shellFormatJson
        :: Context -> o -> a -> Value
  
  data WithFormatOpts o a = WithFormatOpts a o
  instance ShellFormat o a => ShellFormat d (WithFormatOpts o a) where
    shellFormatText ctx _ (WithFormatOpts x opts) = shellFormatText ctx opts x
    shellFormatJson ctx _ (WithFormatOpts x opts) = shellFormatJson ctx opts x
  
  withFormatOpts :: ShellFormat o a => a -> o -> WithFormatOpts o a
  withFormatOpts = WithFormatOpts
  
  instance ShellFormat o Void where
    shellFormatText _ctx _ v = case v of {}
    shellFormatJson _ctx _ v = case v of {}
  
  instance ShellFormat o String where
    shellFormatText _ctx _ s = pretty s
    shellFormatJson _ctx _ s = J.toJSON s
  
  instance ShellFormat o Thrift.DatabaseProperties where
    shellFormatText _ctx _ props = vsep
      [ pretty name <> ":" <+> pretty value
      | (name,value) <- sortOn fst $
          HashMap.toList props
      ]
    shellFormatJson _ctx _ props = J.toJSON props
  
  instance ShellFormat o Thrift.DatabaseStatus where
    shellFormatText _ctx _ status =
      case status of
        Thrift.DatabaseStatus_Complete -> parens "complete"
        Thrift.DatabaseStatus_Available -> parens "available elsewhere"
        Thrift.DatabaseStatus_Finalizing -> parens "finalizing"
        Thrift.DatabaseStatus_Incomplete -> parens "incomplete"
        Thrift.DatabaseStatus_Restoring -> parens "restoring"
        Thrift.DatabaseStatus_Broken -> parens "broken"
        Thrift.DatabaseStatus_Restorable -> parens "restorable"
        Thrift.DatabaseStatus_Missing -> parens "missing deps"
    shellFormatJson _ctx _ status = J.toJSON @String $
        case status of
          Thrift.DatabaseStatus_Complete -> "COMPLETE"
          Thrift.DatabaseStatus_Available-> "AVAILABLE"
          Thrift.DatabaseStatus_Finalizing -> "FINALIZING"
          Thrift.DatabaseStatus_Incomplete -> "INCOMPLETE"
          Thrift.DatabaseStatus_Restoring -> "RESTORING"
          Thrift.DatabaseStatus_Broken -> "BROKEN"
          Thrift.DatabaseStatus_Restorable -> "RESTORABLE"
          Thrift.DatabaseStatus_Missing -> "MISSING"
  
  instance ShellFormat o Thrift.Dependencies where
    shellFormatText _ctx _ dependency =
      case dependency of
        Thrift.Dependencies_stacked Thrift.Stacked{..} ->
          pretty (showRepo $ Thrift.Repo stacked_name stacked_hash)
        Thrift.Dependencies_pruned (Thrift.Pruned baseRepo _ _ _)
          -> pretty (showRepo baseRepo)
    shellFormatJson _ctx _ dependency =
      case dependency of
        Thrift.Dependencies_stacked Thrift.Stacked{..} ->
          J.toJSON (showRepo $ Thrift.Repo stacked_name stacked_hash)
        Thrift.Dependencies_pruned (Thrift.Pruned baseRepo _ _ _)
          -> J.toJSON (showRepo baseRepo)
  
  instance ShellFormat o Thrift.Repo where
    shellFormatText _ctx _ repo = pretty (showRepo repo)
    shellFormatJson _ctx _ repo = J.toJSON (showRepo repo)
  
  statusColour :: Thrift.DatabaseStatus -> Color
  statusColour status = case status of
    Thrift.DatabaseStatus_Complete -> Green
    Thrift.DatabaseStatus_Available -> Green
    Thrift.DatabaseStatus_Finalizing -> Green
    Thrift.DatabaseStatus_Incomplete -> Blue
    Thrift.DatabaseStatus_Restoring -> Black
    Thrift.DatabaseStatus_Broken -> Red
    Thrift.DatabaseStatus_Restorable -> Black
    Thrift.DatabaseStatus_Missing -> Black
  
  statusStyle :: Thrift.DatabaseStatus -> AnsiStyle
  statusStyle = color . statusColour
  
  data DbVerbosity
    = DbSummarise
    | DbDescribe
    deriving (Eq)
  
  instance ShellFormat DbVerbosity Thrift.Database where
    shellFormatText ctx opts db =
      shellFormatText ctx opts (db, []::[(String, Void)])
    shellFormatJson ctx opts db =
      shellFormatJson ctx opts (db, []::[(String, Void)])
  
  instance (ShellFormat DbVerbosity v)
    => ShellFormat DbVerbosity (Thrift.Database, [(String, v)]) where
      shellFormatText ctx@Context{..} opts (db, extras) = nest 2 $ vsep $
        [ annotate (statusStyle status) (shellFormatText ctx () repo)
          <+> shellFormatText ctx () status]
        ++
        [ "Source:" <+> showWhen t
        | Just t <- [Thrift.database_repo_hash_time db]
        ]
        ++
        [ "Created:" <+> showWhen (Thrift.database_created_since_epoch db) ]
        ++
        [ "Completed:" <+> showWhen databaseComplete_time
        | Just Thrift.DatabaseComplete{..} <- [Thrift.database_complete db]
        ]
        ++
        [ "Size:" <+> pretty(showSizeBytes (fromIntegral s))
        | Just Thrift.DatabaseComplete{..} <- [Thrift.database_complete db]
        , Just s <- [databaseComplete_bytes]
        ]
        ++
        [ "Broken:"
            <+> (if Text.null task then emptyDoc else pretty task <> ": ")
            <> pretty reason
        | Just (Thrift.DatabaseBroken task reason) <- [Thrift.database_broken db]
        ]
        ++
        [ pretty key <> ":" <+> shellFormatText ctx opts value
        | (key, value) <- extras]
        ++
        [ "Dependencies:" <+> shellFormatText ctx () dependency
        | Just dependency <- [Thrift.database_dependencies db]
        ]
        ++
        [ "Backup:" <+> pretty loc
        | verbosity == DbDescribe ||
          Thrift.database_status db == Thrift.DatabaseStatus_Restorable
        , Just loc <- [Thrift.database_location db]
        ]
        ++
        [ "Expires in:" <+> pretty (ppTimeSpan timeSpan)
        | verbosity == DbDescribe
        , Just expiresEpochTime <-
            [Thrift.unPosixEpochTime <$> Thrift.database_expire_time db]
        , let expires = Time $ fromIntegral expiresEpochTime
        , let timeSpan = expires `timeDiff` ctxNow
        ]
        ++
        [ "Shard:" <+> pretty (dbShard repo)
        | verbosity == DbDescribe
        ]
        ++
        [ nest 2 $ vsep
          [ "Properties:"
          , shellFormatText ctx () (Thrift.database_properties db)
          ]
        | verbosity == DbDescribe
        ]
        where
          showWhen (Thrift.PosixEpochTime t) =
            pretty (show (posixSecondsToUTCTime (fromIntegral t))) <+>
              parens (pretty (Text.unpack age) <+> "ago")
            where
              age = ppTimeSpanWithGranularity Hour $
                ctxNow `timeDiff` Time (fromIntegral t)
  
          showSizeBytes :: Double -> String
          showSizeBytes b
            | b > 1e9 = printf "%.2f GB" (b / 1e9)
            | otherwise = printf "%.2f MB" (b / 1e6)
  
          status = Thrift.database_status db
          repo = Thrift.database_repo db
          verbosity = opts
      shellFormatJson ctx opts (db, extras) = J.object $
        [ "repo" .= shellFormatJson ctx () repo
        , "status" .= shellFormatJson ctx () status
        , "created" .= jsonTime (Thrift.database_created_since_epoch db)
        , "completed" .= jsonMaybeTime
            (Thrift.databaseComplete_time <$> Thrift.database_complete db)
        , "backup" .= maybe J.Null J.toJSON (Thrift.database_location db)
        , "expires" .= jsonMaybeTime (Thrift.database_expire_time db)
        , "shard" .= J.toJSON (dbShard $ Thrift.database_repo db)
        , "source" .= maybe J.Null jsonTime (Thrift.database_repo_hash_time db)
        , "properties" .= shellFormatJson ctx () (Thrift.database_properties db)
        , "dependencies" .= maybe J.Null J.toJSON
          (Thrift.database_dependencies db)
        , "size" .= maybe J.Null J.toJSON
            (Thrift.databaseComplete_bytes <$> Thrift.database_complete db)
        ] ++
        [ keyFromText (jsonKeyFrom key) .= shellFormatJson ctx opts value
        | (key, value) <- extras]
        where
          status = Thrift.database_status db
          repo = Thrift.database_repo db
          jsonMaybeTime = maybe J.Null jsonTime
          jsonTime (Thrift.PosixEpochTime t) =
            J.toJSON $ posixSecondsToUTCTime (fromIntegral t)
          jsonKeyFrom s = Text.pack $ map f s
            where
              f ' ' = '_'
              f c = toLower c
  
  instance ShellFormat DbVerbosity [Thrift.Database] where
    shellFormatText ctx opts dbs =
      vsep $ concatMap f $ sortOn Thrift.database_created_since_epoch dbs
      where f db = [shellFormatText ctx opts db, pretty ("    "::Text)]
    shellFormatJson ctx v dbs = J.toJSON $ map (shellFormatJson ctx v) dbs
  
  instance (ShellFormat DbVerbosity v)
    => ShellFormat DbVerbosity [(Thrift.Database, [(String, v)])] where
      shellFormatText ctx opts dbs = vsep $ concatMap f $
        sortOn (Thrift.database_created_since_epoch . fst) dbs
        where f x = [shellFormatText ctx opts x, pretty ("    "::Text)]
      shellFormatJson ctx v dbs = J.toJSON $ map (shellFormatJson ctx v) dbs
  
  type PredStatsList =
    [(Either Thrift.Id Thrift.PredicateRef, Thrift.PredicateStats)]
  type PredStatsFilter =
    Either Thrift.Id Thrift.PredicateRef -> Bool
  
  data StatsFormatOpts = StatsFormatOpts
    { showTotal :: Bool
    , sortBySize :: Bool
    }
    deriving (Eq)
  
  instance ShellFormat StatsFormatOpts (PredStatsFilter, PredStatsList) where
      shellFormatText _context opts (filterPred, preds) = vsep $
          [ nest 2 $ vsep
              [ case ref of
                  Left id -> pretty id
                  Right pref -> pretty pref
              , "count:" <+> pretty (Thrift.predicateStats_count stats)
              , "size: " <+> pretty (getSizeInfo
                  (Thrift.predicateStats_size stats) totalSizeBytes)
              ]
          | (ref, stats) <- sort $ filter (filterPred . fst) preds
          ] ++
          if showTotal opts then
            [ ""
            , "Total: " <> pretty totalFacts
                <+> "facts" <+> parens (pretty (showAllocs totalSizeBytes))
            ]
          else
            []
        where
          predicate_count = Thrift.predicateStats_count . snd
          predicate_size = Thrift.predicateStats_size . snd
          totalFacts = foldl' (+) 0 $ map predicate_count preds
          totalSizeBytes =
            foldl' (+) 0 $ map predicate_size preds
          sort = if sortBySize opts then
            sortOn $ negate . predicate_size
          else
            sortOn fst
  
      shellFormatJson _ _ (filterPred, preds) =
        J.toJSON $ filter (filterPred . fst) preds
  
  getSizeInfo :: Int64 -> Int64 -> String
  getSizeInfo bytes total =
    printf "%d (%s) %.4f%%" bytes humanReadableSize percentage_x
      where
        percentage_x :: Double
        percentage_x = 100 * fromIntegral bytes / fromIntegral total
        humanReadableSize = showAllocs bytes
  
  -- | Get the terminal width
  getTerminalWidth :: IO (Maybe Int)
  getTerminalWidth = fmap join $
    -- FIXME: This is a terrible way to get the terminal size but we don't
    -- seem to have any packages which can do this.
    System.Timeout.timeout 100000
      (withCreateProcess
        (proc "stty" ["size"]){std_out = CreatePipe, std_err = CreatePipe}
        (\_ mouth merrh ph -> do
            let outh = fromMaybe (error "outh") mouth
            let errh = fromMaybe (error "errh") merrh
            out <- hGetContents outh
            err <- hGetContents errh
            length out `seq` length err `seq` return ()
            hClose outh
            hClose errh
            ex <- waitForProcess ph
            return $ case ex of
              ExitSuccess
                | [[(_,"")],[(w,"")]] <- map reads $ words out -> Just w
              _ -> Nothing
        ))
    `catchAll` \_ -> return Nothing
  