{-
  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 #-}
{-# LANGUAGE RecursiveDo #-}
module GleanCLI.Derive (DeriveCommand) where

import Control.Concurrent
import Control.Concurrent.Async (async, wait)
import Control.Exception (bracket_, throwIO)
import Control.Monad (forM, unless)
import Data.Foldable (for_)
import Data.Graph (graphFromEdges, vertices)
import qualified Data.Map as Map
import Data.Text (Text)
import qualified Data.Text as Text
import Options.Applicative

import Util.Log.Text
import Util.OptParse
import Util.Text

import Glean
import Glean.Derive (derivePredicate)
import Glean.Types

import GleanCLI.Types
import GleanCLI.Common
import Text.Printf

data DeriveCommand
  = Derive
      { deriveRepo :: Repo
      , predicates :: [(Text, Maybe ParallelDerivation)]
      , derivePageOptions :: PageOptions
      , deriveMaxConcurrency :: Int
      }

instance Plugin DeriveCommand where
  parseCommand =
    commandParser "derive" (progDesc desc) $ do
      deriveRepo <- dbOpts
      deriveMaxConcurrency <- maxConcurrencyOpt
      derivePageOptions <- pageOpts
      predicates <- many (serial <|> parallel)
      return Derive{..}
    where
    desc = "Concurrently derive and store predicates"
    parallel = option parseParallel
      ( long "parallel"
      <> metavar "PREDICATE,OUTER[,SIZE],QUERY"
      <> help "Derive a predicate in parallel"
      )

    parseParallel = maybeReader $ \s ->
      case Text.splitOn "," (Text.pack s) of
        pred : outer : size : rest | Right n <- textToInt size -> Just
          (pred,
             Just ParallelDerivation
               { parallelDerivation_outer_predicate = outer
               , parallelDerivation_inner_query = Text.intercalate "," rest
               , parallelDerivation_min_batch_size = Just (fromIntegral n) })
        pred : outer : rest -> Just
          (pred,
             Just ParallelDerivation
               { parallelDerivation_outer_predicate = outer
               , parallelDerivation_inner_query = Text.intercalate "," rest
               , parallelDerivation_min_batch_size = Nothing })
        _ -> Nothing

    serial = (,Nothing) <$> strArgument
      ( metavar "PREDICATE"
      <> help "predicates to derive"
      )

  runCommand _ _ backend Derive{..}
    | [(pred, parallel)] <- predicates
    = deriveOne pred parallel
    | otherwise
    = mdo
      -- get the schema from the db
      SchemaInfo{..} <- Glean.getSchemaInfo backend (Just deriveRepo) $
        GetSchemaInfo (SelectSchema_stored Empty) False
      -- get the typechecked predicates from the schema
      let (graph, getNode, _) = graphFromEdges edges
          edges = [ (p, p, pp)
                  | (p, pp) <- Map.toList schemaInfo_derivationDependencies
                  ]
          predicateName pid =
            predicateRef_name (schemaInfo_predicateIds Map.! pid)
          derivations =
            [ (predicateName pred, predicateName <$> deps)
            | (_, pred, deps) <- getNode <$> vertices graph
            ]

      concurrencyAvailable <- newQSem deriveMaxConcurrency

      -- set up the parallel evaluation graph
      let
        waitFor predicate = mapM_ wait (ivars Map.!? predicate)
        -- mapping predicates to evaluation asyncs
        ivars = Map.fromList $ zip (fst <$> derivations) asyncs

      -- sanity check: all the requested derivations must be in the graph
      for_ predicates $ \(pred,_) ->
        unless (pred `elem` map fst derivations) $ do
          throwIO $ userError $ printf
            ("abort: %s is not in the schema or is not derived. This is not" ++
              " supported, try deriving one predicate at a time instead.\n")
            pred

      asyncs <- forM derivations $ \(pred, deps) -> async $ do
        -- wait for my dependencies
        mapM_ waitFor deps

        -- if I am one of the requested predicates, derive me
        for_ (lookup pred predicates) $ \parallel -> do
          bracket_
            (waitQSem concurrencyAvailable)
            (signalQSem concurrencyAvailable) $ do
              logInfo $ "Kicking off: " <> pred
              deriveOne pred parallel
              logInfo $ "Done: " <> pred

      -- evaluate the nodes corresponding to the requested predicates
      mapM_ waitFor (fst <$> predicates)
    where
      deriveOne pred parallel =
        derivePredicate backend deriveRepo
          (Just $ fromIntegral $ pageBytes derivePageOptions)
          (fromIntegral <$> pageFacts derivePageOptions)
          (parseRef pred)
          parallel