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

{-# OPTIONS_GHC -Wno-name-shadowing #-}
module Haxl.DataSource.Glean
  ( get
  , getRec
  , getOfId
  , getRecOfId
  , getKey
  , getKeyRec
  , getKeyOfId
  , getKeyRecOfId
  , haxlRepo
  , search
  , search_
  , searchWith
  , count
  , withRepo
  , HasRepo(..)
  ) where

import Data.Hashable
import Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
import Data.Typeable

import Haxl.Core hiding (Env)

import Glean.Backend.Types
import Glean.Query.Thrift
import Glean.Types
import Glean.Typed as Typed

-- User Interface --------

class HasRepo u where
  getRepoFromUserEnv :: u -> Repo

instance HasRepo Repo where
  getRepoFromUserEnv = id

haxlRepo :: HasRepo u => GenHaxl u w Repo
haxlRepo = do
  userEnv <- env userEnv
  return $ getRepoFromUserEnv userEnv

withRepo :: HasRepo u => u -> GenHaxl u w a -> GenHaxl u w a
withRepo userEnv action = do
  coreEnv <- env id
  withEnv coreEnv{userEnv=userEnv} action

get, getRec
  :: (Typeable p, Show p, Predicate p, HasRepo u)
  => p
  -> GenHaxl u w p

get p = haxlRepo >>= \repo -> dataFetch $ Get (getId p) False repo
getRec p = haxlRepo >>= \repo -> dataFetch $ Get (getId p) True repo

getOfId, getRecOfId
  :: (Typeable p, Show p, Predicate p, HasRepo u)
  => IdOf p
  -> GenHaxl u w p

getOfId p = haxlRepo >>= \repo -> dataFetch $ Get p False repo
getRecOfId p = haxlRepo >>= \repo -> dataFetch $ Get p True repo

getKey, getKeyRec
  :: ( Typeable p, Typeable (KeyType p)
     , Show p, Show (KeyType p)
     , Predicate p, HasRepo u )
  => p
  -> GenHaxl u w (KeyType p)

getKey p = haxlRepo >>= \repo -> dataFetch $ GetKey (getId p) False repo
getKeyRec p = haxlRepo >>= \repo -> dataFetch $ GetKey (getId p) True repo

getKeyOfId, getKeyRecOfId
  :: ( Typeable p, Typeable (KeyType p)
     , Show p, Show (KeyType p)
     , Predicate p, HasRepo u )
  => IdOf p
  -> GenHaxl u w (KeyType p)

getKeyOfId id = haxlRepo >>= \repo -> dataFetch $ GetKey id False repo
getKeyRecOfId id = haxlRepo >>= \repo -> dataFetch $ GetKey id True repo

-- | Perform a query using Glean. Returns the results and a 'Bool'
-- indicating whether the results were truncated (either by the server
-- or by an explicit 'limit' applied to the query).
search
  :: (Typeable q, Show q, HasRepo u)
  => Query q
  -> GenHaxl u w ([q], Bool)
search q = haxlRepo >>= \repo -> do
  (a, b) <- dataFetch $ mkQueryReq repo q False
  return (fromAppendList a, b)

-- | Like 'search', but returns results only. Always returns all the
-- results, streaming results from the server if necessary.
search_
  :: (Typeable q, Show q, HasRepo u)
  => Query q
  -> GenHaxl u w [q]
search_ = fmap fromAppendList . searchWith

searchWith
  :: (Typeable r, Show r, Typeable q, Show q, HasRepo u, QueryResult q r)
  => Query q
  -> GenHaxl u w r
searchWith q = haxlRepo >>= \repo ->
  fmap fst $ dataFetch $ mkQueryReq repo q True

newtype UniqueResults a = UniqueResults { fromUniqueResults :: HashSet a }
  deriving (Semigroup, Monoid, Show)

instance (Hashable q, Eq q) => QueryResult q (UniqueResults q) where
  fromResults = UniqueResults . HashSet.fromList

-- | Count the unique results of a query.
count
  :: forall q u w.
     (Typeable q, Show q, Hashable q, Eq q, HasRepo u)
  => Query q
  -> GenHaxl u w Int
count q = haxlRepo >>= \repo -> do
  (UniqueResults (s :: HashSet q), _) <- dataFetch $ mkQueryReq repo q True
  return (HashSet.size s)

-- -----------------------------------------------------------------------------

-- | smart constructor to ensure the query result encoding is set to binary
mkQueryReq
  :: forall q r . (Show q, Typeable q, QueryResult q r)
  => Repo
  -> Query q
  -> Bool
  -> GleanQuery (r, Bool)
mkQueryReq repo q b = QueryReq q repo b