{-
  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.
-}

module Glean.Query.Thrift
  ( -- * Types
    Query
    -- * Perform query
  , runQuery
  , runQuery_
  , runQueryPage
    -- ** streaming
  , runQueryMapPages_
  , runQueryEach
  , runQueryEachBatch
    -- * Query combinators
  , angle
  , angleData
  , keys
  , recursive
  , limit
  , limitBytes
  , limitTime
  , expanding
  , store
  , allFacts
    -- * Support
  , displayQuery
  ) where

import Control.Concurrent.Async
import Control.Monad
import Data.Default
import Data.Maybe

import Util.Log.Text

import Glean.Typed.Binary
import Glean.Types as Thrift
import Glean.Backend.Types (Backend(..))
import Glean.Query.Thrift.Internal
import Glean.Write.SendBatch

-- | Perform a query, fetching all results
runQuery_
  :: forall q backend . (Backend backend)
  => backend    -- ^ Backend to use to perform the query
  -> Repo       -- ^ Repo to query
  -> Query q    -- ^ The query
  -> IO [q]
runQuery_ backend repo query = do
  r <- runQueryEachBatch backend repo query id $
    \acc page -> return (acc . (page++))
  return (r [])

-- | Perform a query, fetching results up to the limits specified by
-- the query ('limit', 'limitBytes', 'limitTime').
runQuery
  :: forall q backend . (Backend backend)
  => backend    -- ^ Backend to use to perform the query
  -> Repo       -- ^ Repo to query
  -> Query q    -- ^ The query
  -> IO ([q], Bool)
       -- ^ The 'Bool' is 'True' if the query results were truncated
       -- by a limit (either a user-supplied limit, or one imposed by
       -- the query server).
runQuery backend repo query = do
  (io, cont) <- runQueryPage backend repo Nothing query
  results <- io
  return (results, isJust cont)


-- | Perform a query, fetching results up to the limits specified by
-- the query ('limit', 'limitBytes', 'limitTime'), and returning a
-- continuation to fetch more results if a limit was hit.
runQueryPage
  :: forall q backend . (Backend backend)
  => backend
  -> Repo
  -> Maybe UserQueryCont
  -> Query q
  -> IO (IO [q], Maybe UserQueryCont)

runQueryPage be repo cont (Query query) = do
  let
    query' = query
      { userQuery_options = Just
          (fromMaybe def (userQuery_options query))
            { userQueryOptions_continuation = cont }
      , userQuery_encodings = [
          UserQueryEncoding_listbin def,
          UserQueryEncoding_bin def
        ]
      }
  UserQueryResults{..} <- userQuery be repo query'
  mapM_ reportUserQueryStats userQueryResults_stats
  mapM_ (vlog 1) userQueryResults_diagnostics
  mapM_ (waitBatch be) userQueryResults_handle
  let results = decodeResults userQueryResults_results decodeAsFact
  return (results, userQueryResults_continuation)


-- | Perform a query and map an IO function over the results, running the
-- query over multiple pages of results if necessary.
runQueryMapPages_
  :: forall q backend . (Backend backend)
  => backend
  -> Repo
  -> ([q] -> IO ())
  -> Query q
  -> IO ()
runQueryMapPages_ be repo fn query =
  runQueryEachBatch be repo query () (const fn)

-- | Perform a query, by pages, and map a state function over each
-- result in each page. The query is pipelined so that each page of
-- results are decoded and processed while the next page is being
-- fetched.
runQueryEach
  :: forall q backend s . (Backend backend)
  => backend
  -> Repo
  -> Query q
  -> s
  -> (s -> q -> IO s)
  -> IO s
runQueryEach be repo query s fn =
  runQueryEachBatch be repo query s (foldM fn)

-- | Like runQueryEach, but process results in batches.
runQueryEachBatch
  :: forall q backend s . (Backend backend)
  => backend
  -> Repo
  -> Query q
  -> s
  -> (s -> [q] -> IO s)
  -> IO s
runQueryEachBatch be repo query init f = do
  (page, cont) <- runQueryPage be repo Nothing query
  go page cont init
  where
  go page Nothing s = f s =<< page
  go page cont@Just{} s = do
    ((nextPage, nextCont), s2) <- concurrently
      (runQueryPage be repo cont query)
      (f s =<< page)
    go nextPage nextCont s2