{-
  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 CPP, ConstraintKinds #-}
module Glean.Haxl
  ( runHaxl
  , runHaxlWithWrites
  , initHaxlEnv
  , Haxl
  , haxlRepo
  , get
  , getRec
  , getKey
  , getKeyRec
  , getOfId
  , getRecOfId
  , getKeyOfId
  , getKeyRecOfId
  , keyOf
  , search
  , search_
  , searchWith
  , QueryResult(..)
  , count
  , Query
  , recursive
  , limit
  , trySyncHaxl
  , getFirstResult
  , HaxlWrite
    -- * re-export
  , initGlobalState
  ) where

import Control.Arrow
import Control.Exception (SomeException)
import Data.Maybe (listToMaybe)
import Data.Typeable

import qualified Haxl.Core as Haxl
import Haxl.Core.Monad
import Haxl.DataSource.Glean
import Util.Control.Exception ( isSyncException )

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

#ifdef GLEAN_FACEBOOK
#define MIN_VERSION_haxl(a,b,c) 1
#endif

#if MIN_VERSION_haxl(2,5,0)
type HaxlWrite w = WriteTree w
flatten :: WriteTree w -> [w]
flatten = flattenWT
#else
type HaxlWrite w = w
flatten :: [w] -> [w]
flatten = id
#endif

type Haxl w = GenHaxl Repo (HaxlWrite w)

-- | Initialize for Glean queries
initHaxlEnv :: Backend b => b -> u -> IO (Haxl.Env u (HaxlWrite w))
initHaxlEnv backend e = do
  (state1,state2) <- initGlobalState backend
  let st = Haxl.stateSet state1 $ Haxl.stateSet state2 Haxl.stateEmpty
  Haxl.initEnv st e

runHaxl :: Backend b => b -> u -> GenHaxl u (HaxlWrite w) a -> IO a
runHaxl backend u h = do
  e <- initHaxlEnv backend u
  Haxl.runHaxl e h

runHaxlWithWrites
  :: Backend b
  => b
  -> u
  -> GenHaxl u (HaxlWrite w) a
  -> IO (a, [w])
runHaxlWithWrites backend u h = do
  e <- initHaxlEnv backend u
  second flatten <$> Haxl.runHaxlWithWrites e h

-- | if the Fact has a key, return it, otherwise fetch it with 'getKey'
keyOf
  :: ( Typeable p, Typeable (KeyType p)
     , Show p, Show (KeyType p)
     , Predicate p, HasRepo u )
  => p
  -> GenHaxl u w (KeyType p)
keyOf f = case getFactKey f of
  Nothing -> getKey f
  Just k -> return k

-- | Catch non-asyncronous exceptions inside Haxl, very good for
-- thrift calls that throw interesting exceptions.
trySyncHaxl :: GenHaxl u w b -> GenHaxl u w (Either SomeException b)
trySyncHaxl act = catchIf isSyncException (Right <$> act) (return . Left)

-- | Search for at most 1 result and return it or Nothing
getFirstResult
  :: (Typeable a, Show a, HasRepo u)
  => Query a -> GenHaxl u w (Maybe a)
getFirstResult = fmap (listToMaybe . fst) . search . limit 1