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

module Glean.Haxl.Repos
  ( ReposHaxl
  , RepoHaxl
  , MaybeTReposHaxl
  , MaybeTRepoHaxl
  , QueryRepos
  , queryAllRepos
  , withRepo
  , queryEachRepo
  , runHaxlAllRepos
  , Repos(..)
  ) where

import Control.Monad.Trans.Maybe (MaybeT (..))

import Haxl.Core hiding (Env, runHaxl)
import Haxl.Core.Monad (WriteTree)
import Haxl.DataSource.Glean (HasRepo(..))
import qualified Haxl.DataSource.Glean
import Data.List.NonEmpty

import Glean.Backend.Types
import Glean.Types
import qualified Glean.Haxl as Glean


-- Searching across repos --------

type ReposHaxl u w a = QueryRepos u => GenHaxl u (WriteTree w) a
type RepoHaxl u w a =
  Haxl.DataSource.Glean.HasRepo u => GenHaxl u (WriteTree w) a
type MaybeTReposHaxl u w a = QueryRepos u => MaybeT (GenHaxl u (WriteTree w)) a
type MaybeTRepoHaxl u w a =
  Haxl.DataSource.Glean.HasRepo u => MaybeT (GenHaxl u (WriteTree w)) a

-- This type has instances for both HasRepo and HasRepos.
-- The instance `HasRepo` allows us to use it as the userState
-- but its not exported so the user doesn't accidently run a query over
-- only the first repo. Instead the user should just depend on the
-- QueryRepos constraint
newtype Repos u = Repos u

class HasRepos us where
  headRepoFromUserEnv :: us -> Repo
  tailReposFromUserEnv :: us -> Maybe us
  singleRepo :: Repo -> us

instance HasRepos (NonEmpty Repo) where
  headRepoFromUserEnv (x:|_) = x
  tailReposFromUserEnv (_ :| []) = Nothing
  tailReposFromUserEnv (_ :| x : xs) = Just $ x :| xs
  singleRepo repo = repo :| []

instance HasRepos u => HasRepos (Repos u) where
  headRepoFromUserEnv (Repos xs) = headRepoFromUserEnv xs
  tailReposFromUserEnv (Repos xs) = Repos <$> tailReposFromUserEnv xs
  singleRepo repo = Repos $ singleRepo repo

instance HasRepos u => HasRepo (Repos u) where
  getRepoFromUserEnv (Repos xs) = headRepoFromUserEnv xs

class QueryRepos u where
  queryEachRepo :: (HasRepo u => GenHaxl u w q) -> GenHaxl u w [q]
  withRepo :: Repo -> RepoHaxl u w q -> ReposHaxl u w q

instance (HasRepos u) => QueryRepos (Repos u) where
  queryEachRepo act = do
    head <- act
    userEnv <- env userEnv
    tail <- maybe
      (return [])
      (\r -> Haxl.DataSource.Glean.withRepo r $ queryEachRepo act)
      (tailReposFromUserEnv userEnv)
    return $ head:tail
  withRepo repo act = Haxl.DataSource.Glean.withRepo (singleRepo repo) act

queryAllRepos
  :: RepoHaxl u w [q]
  -> ReposHaxl u w [q]
queryAllRepos act = concat <$> queryEachRepo act

runHaxlAllRepos
  :: (HasRepos u, Backend b)
  => b
  -> u
  -> (forall u'. ReposHaxl u' w a)
  -> IO a
runHaxlAllRepos backend u h = Glean.runHaxl backend (Repos u) h