{-
  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 ConstraintKinds #-}
{-# LANGUAGE ApplicativeDo #-}

module Glean.Glass.Utils (
  -- Search utilities
    fetchData
  , fetchDataRecursive
  , searchWithLimit
  , searchWithTimeLimit
  , searchReposWithLimit

  , searchPredicateWithLimit
  , searchRecursiveWithLimit

  -- File utilities
  , pathFragments
  , joinFragments

  -- Tuple utilities
  , fst4

  -- List utils
  , takeFairN
  , splitOnAny

  -- Str utils
  , splitString

  -- Exceptions
  , eThrow

  -- Types
  , QueryType

  -- compat stuff
  , liftMaybe
  , maybeT

  ,translateDocumentSymbolListXResult) where

import Control.Monad.Catch ( throwM )
import Data.Maybe (fromMaybe)
import Data.Tuple.Extra ( snd3 )
import Data.Text ( Text, length )
import qualified Data.Text as Text
import Data.Typeable ( Typeable )
import System.FilePath ( splitDirectories, joinPath )
import qualified Data.List as List

import Control.Monad.Extra ( MonadPlus(mzero) )
import Control.Monad.Trans ( MonadTrans(..) )

import Glean ( recursive, limit, limitTime, search, getFirstResult, Predicate )
import Glean.Angle as Angle ( query, Angle, Type )
import Util.Text (slice)

import qualified Glean.Haxl.Repos as Glean
import Glean.Haxl.Repos (RepoHaxl, ReposHaxl)

import qualified Glean.Glass.Types as Types
import Glean.Glass.Types (
    mAXIMUM_SYMBOLS_QUERY_LIMIT,
    mAXIMUM_QUERY_TIME_LIMIT,
    ServerException(..), Path (Path), RepoName
  )

type QueryType q =
  ( Typeable q
  , Show q
  , Type q
  )

--
-- | Evaluate an Angle data query and return the first result or Nothing
-- No more than 1 result will be returned from the server.
--
fetchData :: (QueryType a) => Angle a -> RepoHaxl u w (Maybe a)
fetchData = getFirstResult . query

-- | Fetch exactly 0 or 1 results, recusively
fetchDataRecursive :: (QueryType a) => Angle a -> RepoHaxl u w (Maybe a)
fetchDataRecursive = getFirstResult . recursive . query

-- | Run a non-recursive data query with optional limit on search results.
-- Without an explicit limit the global limits apply.
-- This uses the global time limit of 30s per query (2x the expected thrift
-- limit)
searchWithLimit :: QueryType q => Maybe Int -> Angle q -> RepoHaxl u w [q]
searchWithLimit mlimit = searchWithTimeLimit mlimit
  (fromIntegral mAXIMUM_QUERY_TIME_LIMIT)

-- | Like searchWithLimit but enforce a time bound on the query too.
-- Time budget in milliseconds
searchWithTimeLimit
  :: QueryType q => Maybe Int -> Int -> Angle q -> RepoHaxl u w [q]
searchWithTimeLimit mlimit time =
    fmap fst <$> search . limitTime time . limit item . Angle.query
  where
    item = fromMaybe (fromIntegral mAXIMUM_SYMBOLS_QUERY_LIMIT) mlimit

-- | Run a non-recursive predicate only query with optional limit
searchPredicateWithLimit
  :: (Predicate q, QueryType q) => Maybe Int -> Angle q -> RepoHaxl u w [q]
searchPredicateWithLimit = searchWithLimit

-- | Run a recursive data query with optional limit on search results
-- If not limit set, MAXIMUM_SYMBOLS_QUERY_LIMIT is enforced
--
-- Returns a status bit indicating truncated results
--
searchRecursiveWithLimit
  :: QueryType q => Maybe Int -> Angle q -> RepoHaxl u w ([q], Bool)
searchRecursiveWithLimit mlimit =
    search . recursive . limitTime time . limit item . Angle.query
  where
    item = fromMaybe (fromIntegral mAXIMUM_SYMBOLS_QUERY_LIMIT) mlimit
    time = fromIntegral mAXIMUM_QUERY_TIME_LIMIT

-- | Run a non-recursive data query with optional limit over multiple repos
searchReposWithLimit
  :: QueryType q
  => Maybe Int
  -> Angle q
  -> (q -> RepoHaxl u w a)
  -> ReposHaxl u w [a]
searchReposWithLimit limit angle act = do
  results <- Glean.queryAllRepos $ do
    res <- searchWithLimit limit angle
    mapM act res -- we would like this to be concurrent
  return $ maybe id take limit results

-- | Split a filepath into a list of directory components
pathFragments :: Text -> [Text]
pathFragments = map Text.pack . splitDirectories . Text.unpack

joinFragments :: [Text] -> Text
joinFragments = Text.pack . joinPath . map Text.unpack

-- | Take elements in row,column ordering, so that we select
-- evenly up to N from the set of input lists
--
-- > takeFairN 3 [[1,2],[3],[4,5]] == [1,3,4]
--
takeFairN :: Int -> [[a]] -> [a]
takeFairN _ [] = []
takeFairN n [vs] = take n vs
takeFairN n xs = take n (concat (List.transpose xs))

--
-- Splitv string on substring, with left to right precedence of patterns
--
splitOnAny :: [Text] -> Text -> [Text]
splitOnAny pats src =
  List.foldl' (\acc p -> concatMap (Text.splitOn p) acc) [src] pats

subString :: Int -> Int -> Text -> Text
subString start len = slice start len

-- Split a type string along reference spans. Annotate extracted fragments.
-- (annotation, start, length) if spans are inconsistent, don't fragment
splitString :: Text -> [(ann, Int, Int)] -> [(Text, Maybe ann)]
splitString s xrefs =
  let res = reverse $ splitStringAux 0 (List.sortOn snd3 xrefs) []
      check = Text.concat $ fmap fst res
  in if check == s then res else [(s, Nothing)]
  where
    n = Data.Text.length s
    splitStringAux pos xrefs res = case xrefs of
        [] | pos == n -> res
           | otherwise -> (subString pos (n - pos) s, Nothing) : res
        (ann, start, length) : xrefs' ->
          if pos == start then
            splitStringAux
              (pos + length)
              xrefs'
              ((subString pos length s, Just ann) : res)
          else
            splitStringAux
              start
              xrefs
              ((subString pos (start - pos) s, Nothing) : res)

-- From Control.Monad.Extra in newer versions
liftMaybe :: (MonadPlus m) => Maybe a -> m a
liftMaybe = maybe mzero return

maybeT :: (MonadTrans t, Monad m, MonadPlus (t m)) => m (Maybe b) -> t m b
maybeT act = lift act >>= liftMaybe

eThrow :: Either Text a -> RepoHaxl u w a
eThrow (Right x) = pure x
eThrow (Left err) = throwM $ ServerException err

fst4 :: (a,b,c,d) -> a
fst4 (x,_,_,_) = x


-- replace locationRange and revision in DocumentSymbolListXResult
--  oldRepo -> newRepo
--  path -> prefixPath <> path
--  revision -> newRev (if provided)
translateDocumentSymbolListXResult
 :: RepoName
 -> RepoName
 -> Text
 -> Maybe Types.Revision
 -> Types.DocumentSymbolListXResult
 -> Types.DocumentSymbolListXResult
translateDocumentSymbolListXResult oldRepo newRepo prefixPath newRev result =
  let defaultRev = Types.documentSymbolListXResult_revision result in
  let references' =
        translateRef <$> Types.documentSymbolListXResult_references result in
  result{
        Types.documentSymbolListXResult_revision = fromMaybe defaultRev newRev,
        Types.documentSymbolListXResult_references = references'}
  where
    translateRef ref =
      let t = translateLocationRange $ Types.referenceRangeSymbolX_target ref in
      ref{Types.referenceRangeSymbolX_target = t}
    translateLocationRange ref@(Types.LocationRange repo path range) =
      if repo == oldRepo then
          Types.LocationRange  newRepo (translatePath path) range
        else ref
    translatePath (Path path) = Path (prefixPath <> path)