{-
  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 #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
-- | Define internal pieces, please import "Glean.Query.Thrift" instead
module Glean.Query.Thrift.Internal
  ( -- * Types
    Query(..)
    -- * Query combinators
  , angle
  , angleData
  , keys
  , recursive
  , limit
  , limitBytes
  , limitTime
  , store
  , allFacts
  , expanding
    -- * Support
  , reportUserQueryStats
  , showUserQueryStats
  , displayQuery
  , decodeResults
  , justCheck
  , dbgPredHasFacts
  ) where

import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Data.Default
import Data.Dynamic
import Data.Hashable
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.IORef
import qualified Data.Map as Map
import Data.Maybe
import Data.Proxy
import Data.Text as Text
import qualified Data.Text.Encoding as Text
import Compat.Prettyprinter hiding ((<>))
import qualified Data.Vector as Vector
import Text.Printf

import Util.Log

import Glean.Typed as Typed
import Glean.Types as Thrift

type ResultDecoder a =
     IntMap Thrift.Fact               -- ^ serialized nested facts
  -> IORef (IntMap Dynamic)           -- ^ cached deserialized facts
  -> IdOf a
  -> Thrift.Fact
  -> IO a

-- | A query that can be performed by 'runQuery'. Build using 'query'.
data Query a = Type a => Query
  { querySpec :: UserQuery
  }

instance Show (Query a) where
  show (Query spec) = show spec

instance Eq (Query a) where
  Query spec1 == Query spec2 = spec1 == spec2

instance Hashable (Query a) where
  hashWithSalt salt (Query spec) = hashWithSalt salt spec

decodeResults :: UserQueryEncodedResults -> ResultDecoder a -> IO [a]
decodeResults results decoder =
  case results of
    UserQueryEncodedResults_bin UserQueryResultsBin{..} -> do
      cacheRef <- newIORef IntMap.empty
      let serialized = IntMap.fromList
            [ (fromIntegral id,f)
            | (id,f) <- Map.toList userQueryResultsBin_nestedFacts ]
      forM (Map.toList userQueryResultsBin_facts) $ \(fid, fact) -> do
        liftIO $ decoder serialized cacheRef (Typed.IdOf (Fid fid)) fact

    UserQueryEncodedResults_listbin UserQueryResultsListBin{..} -> do
      cacheRef <- newIORef IntMap.empty
      let serialized = IntMap.fromList
            [ (fromIntegral id,f)
            | (id,f) <- Map.toList userQueryResultsListBin_nestedFacts ]
          fids = Vector.toList userQueryResultsListBin_ids
          facts = Vector.toList userQueryResultsListBin_facts
          f fid fact = liftIO $
            decoder serialized cacheRef (Typed.IdOf (Fid fid)) fact
      zipWithM f fids facts

    _other -> throwIO $ ErrorCall
      "runQueryPage: server returned the wrong encoding"

-- | A human-readable form of the Query.
displayQuery :: Query a -> Text
displayQuery (Query UserQuery{..}) = Text.decodeUtf8 userQuery_query

allFacts :: forall q . (Predicate q) => Query q
allFacts = angle $ Text.pack $
  show (pretty (getName (Proxy :: Proxy q))) <> " _"

-- | Build a query for facts using Angle syntax
--
-- Warning: nothing is checking that the type of the query matches the
-- expected type at the callsite. If you get this wrong, undefined
-- behaviour will ensue.
angle :: forall r. (Predicate r) => Text -> Query r
angle q = Query
  { querySpec = angleQuery (getName (Proxy @r)) q
  }

-- | Fetch just the keys of the given fact query
keys :: (Predicate r) => Query r -> Query (KeyType r)
keys (Query spec) = Query { querySpec = spec }

angleQuery :: PredicateRef -> Text -> UserQuery
angleQuery pref query = def
  { userQuery_predicate = predicateRef_name pref
  , userQuery_predicate_version =
    Just (fromIntegral (predicateRef_version pref))
  , userQuery_query = Text.encodeUtf8 query
  , userQuery_options = Just def
    { userQueryOptions_expand_results = False
    , userQueryOptions_syntax = QuerySyntax_ANGLE
    }
  }

-- | A query for arbitrary data, as opposed to 'angle' which queries
-- for facts.
--
-- For example:
--
-- > angleData "{ 3, false }" :: Query (Nat, Bool)
--
-- Warning: nothing is checking that the type of the query matches the
-- expected type at the callsite. If you get this wrong, undefined
-- behaviour will ensue.
angleData :: (Type r) => Text -> Query r
angleData query = Query { querySpec = spec  }
  where
  spec = def
    { userQuery_query = Text.encodeUtf8 query
    , userQuery_options = Just def
      { userQueryOptions_expand_results = False
      , userQueryOptions_syntax = QuerySyntax_ANGLE
      }
    }

-- | Make a query recursive
recursive :: Query a -> Query a
recursive (Query q) = Query q'
  where
  q' = q { userQuery_options = Just (fromMaybe def (userQuery_options q))
    { userQueryOptions_recursive = True } }

-- | Instead of 'recursive', pick individual predicates to expand in
-- the result
--
-- e.g.
--   expanding @My.Predicate1 $ expanding @My.Predicate2 $ query ...
--
expanding :: forall p a . Predicate p => Query a -> Query a
expanding (Query q) = Query q'
  where
  old = fromMaybe def (userQuery_options q)
  PredicateRef name version = getName (Proxy @p)
  ref = Thrift.SourcePredicate name (Just version)
  q' = q { userQuery_options = Just old
    { userQueryOptions_expand_predicates =
       ref : userQueryOptions_expand_predicates old }}

-- | Set a limit on the number of results returned by a query. This controls
-- query result page size when streaming.
limit :: Int -> Query a -> Query a
limit n (Query q) = Query q'
  where
  q' = q { userQuery_options = Just (fromMaybe def (userQuery_options q))
    { userQueryOptions_max_results = Just (fromIntegral n) } }

-- | Set a limit on the size of data returned by a query. This
-- controls query result page size when streaming.
limitBytes :: Int -> Query a -> Query a
limitBytes n (Query q) = Query q'
  where
  q' = q { userQuery_options = Just (fromMaybe def (userQuery_options q))
    { userQueryOptions_max_bytes = Just (fromIntegral n) } }

-- | Set a limit on the time that a query can execute for. When streaming,
-- this imposes a time budget on each query page.
limitTime
  :: Int  -- ^ time budget in milliseconds
  -> Query a -> Query a
limitTime n (Query q) = Query q'
  where
  q' = q { userQuery_options = Just (fromMaybe def (userQuery_options q))
    { userQueryOptions_max_time_ms = Just (fromIntegral n) } }

-- | Store derived facts arising from the query. This is intended for
-- caching derived facts after creating a DB, not for normal usage. If
-- the datbase is read-only, this will cause the query to fail.
store :: Query a -> Query a
store (Query q) = Query q'
  where
  q' = q { userQuery_options = Just (fromMaybe def (userQuery_options q))
    { userQueryOptions_store_derived_facts = True } }

justCheck :: Query a -> Query a
justCheck (Query q) = Query q'
  where
  q' = q { userQuery_options = Just (fromMaybe def (userQuery_options q))
  { userQueryOptions_just_check = True } }

dbgPredHasFacts :: Query a -> Query a
dbgPredHasFacts (Query q) = Query q'
  where
  q' = q { userQuery_options = Just (fromMaybe def (userQuery_options q))
  { userQueryOptions_debug = def {queryDebugOptions_pred_has_facts = True}}}

reportUserQueryStats :: Thrift.UserQueryStats -> IO ()
reportUserQueryStats stats =
  vlog 1 $ showUserQueryStats stats

showUserQueryStats :: Thrift.UserQueryStats -> String
showUserQueryStats Thrift.UserQueryStats{..} =
  printf "%d facts, %.2fms, %ld bytes\n"
    userQueryStats_num_facts
    (realToFrac userQueryStats_elapsed_ns / 1000000 :: Double)
    userQueryStats_allocated_bytes