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

-- |
-- Small domain-specific language for building Angle queries
-- programmatically.
--
{-# LANGUAGE TypeApplications, TypeOperators, AllowAmbiguousTypes #-}
{-# OPTIONS_GHC -Wno-star-is-type #-}
module Glean.Query.Angle
  ( Angle
  , AngleStatement
  , var
  , Type
  , vars
  , predicate
  , where_
  , not_
  , (.=)
  , stmt
  , (.|)
  , or_
  , (.->)
  , query
  , nat
  , byte
  , enum
  , string
  , stringPrefix
  , byteArray
  , array
  , tuple
  , wild
  , never
  , field
  , end
  , rec
  , alt
  , asPredicate
  , sig
  , factId
  , factIds
  , factIdsArray
  , elementsOf
  , arrayPrefix
  , unit
  , if_
  , true
  , false
  , bool
  , just
  , nothing
  , (.+)
  , zipp
  , conc
  , display
  , RecordFields
  , SumFields
  , TFields(..)
  , HasFields
  , HasField
  , AngleEnum(..)
  , AngleVars
  , new
  , old
  ) where

import Control.Monad.State.Strict
import Data.ByteString (ByteString)
import Data.Proxy
import Data.String
import qualified Data.List.Extra as List (nubOrd)
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NonEmpty (fromList, toList)
import Data.Text (Text)
import Compat.Prettyprinter hiding ((<>))
import Compat.Prettyprinter.Render.Text
import qualified Data.Text as Text
import qualified Data.Text.Encoding as Text
import Data.Word
import GHC.TypeLits hiding (Nat)
import TextShow

import Glean.Angle.Types hiding
  (SourcePat, SourceStatement, SourceQuery, Field, Type)
import qualified Glean.Angle.Types as Angle
import qualified Glean.Display as Display
import Glean.Query.Thrift.Internal as Thrift
import Glean.Typed hiding (end)
import Glean.Types (Nat, Byte)

data SpanAngleDSL = DSL
  deriving (Show, Eq)

instance IsSrcSpan SpanAngleDSL where
  type Loc SpanAngleDSL = SpanAngleDSL
  startLoc _ = DSL
  endLoc _ = DSL
  mkSpan _ _ = DSL

instance Pretty SpanAngleDSL where
  pretty DSL = "angle DSL"

type SourcePat = SourcePat' SpanAngleDSL ()
type SourceStatement = SourceStatement' SpanAngleDSL ()
type SourceQuery = SourceQuery' SpanAngleDSL ()

newtype Angle t = Angle { gen :: State Int SourcePat }

newtype AngleStatement =
  AngleStatement { genStmt :: State Int SourceStatement }

-- | Render to angle text, to aid logging and debugging
display :: Angle t -> Text
display = render . build

render :: SourceQuery -> Text
render q = renderStrict (layoutCompact (Display.displayDefault q))

build :: Angle t -> SourceQuery
build (Angle m) =
  case evalState m 0 of
    NestedQuery _ q -> q
    t -> Angle.SourceQuery (Just t) [] Unordered

-- | Build a query. It can returns facts of a predicate:
--
-- >    query $ predicate @Pp.Define (field @"macro" "NULL" end)
--
-- Or arbitrary values:
--
-- >    query $ var $ \n ->
-- >      n `where_` [
-- >        stmt $ predicate @Hack.MethodDeclaration $
-- >          rec $
-- >            field @"name" "foo" $
-- >            field @"container"
-- >              (rec $
-- >                 field @"class_"
-- >                    (rec (field @"name" n end))
-- >                 end)
-- >      ]
query :: (Type t) => Angle t -> Query t
query = Thrift.angleData . display . sig
  -- adding a type signature ensures that our type matches the type
  -- Glean infers. Otherwise these could diverge, leading to
  -- deserialization errors or just wrong data.

class AngleVars f r where
  -- | Use `vars` to batch up a series of nested `var` calls:
  --
  -- > let foo :: Angle a -> Angle b -> Angle c -> Angle d -> Angle e
  -- > vars foo == var $ \a -> var $ \b -> var $ \c -> var $ \d -> foo a b c d
  vars :: f -> Angle r

instance (a ~ b) => AngleVars (Angle a) b where
  vars = id

instance (Type a, AngleVars s r) => AngleVars (Angle a -> s) r where
  vars f = var $ \ a -> vars (f a)

-- | Introduce an Angle variable
var :: forall a b . Type a => (Angle a -> Angle b) -> Angle b
var f = Angle $ do
  x <- get
  modify (+1)
  let v = Angle (pure (Variable DSL ("X" <> showt x)))
  ty <- gen (sig v)
  r <- gen (f v)
  let stmt = Angle.SourceStatement (Wildcard DSL) ty
  case r of
    Angle.NestedQuery DSL (Angle.SourceQuery q stmts ord) ->
      return (Angle.NestedQuery DSL (Angle.SourceQuery q (stmt : stmts) ord))
    _ ->
      return (Angle.NestedQuery DSL (Angle.SourceQuery (Just r) [stmt] Ordered))

predicate :: forall p . Predicate p => Angle (KeyType p) -> Angle p
predicate (Angle pat) = Angle $ do
  p <- pat
  return $ App DSL (Variable DSL (predicateRef (Proxy @p))) [p]

-- | Build a query of the form `X where statements`
where_ :: Angle t -> [AngleStatement] -> Angle t
where_ t stmts = Angle $
  nest <$> gen t <*> mapM genStmt stmts
  where
  -- merge adjacent where-clauses
  nest (Angle.NestedQuery DSL (Angle.SourceQuery (Just q) stmts1 ord)) stmts2 =
    Angle.NestedQuery DSL $
      Angle.SourceQuery (Just q) (stmts1 <> stmts2) ord
  nest q stmts =
    Angle.NestedQuery DSL (Angle.SourceQuery (Just q) stmts Ordered)

not_ :: [AngleStatement] -> AngleStatement
not_ stmts = unit' .= Angle (Negation DSL <$> gen t)
  where
    t = unit' `where_` stmts
    unit' = sig unit

-- | Build a statement, `A = B`
(.=) :: Angle t -> Angle t -> AngleStatement
l .= r = AngleStatement $ Angle.SourceStatement <$> gen l <*> gen r

-- | Build a statement from a plain Angle expression, ignoring the result.
-- `stmt A` is equivalent in Angle to the statement `_ = A`. For example
--
-- > var $ \x ->
-- >   x `where_` [
-- >      stmt $ predicate @Src.File x
-- >   ]
--
stmt :: Angle t -> AngleStatement
stmt = (wild .=)

infix 1 .=
infixr 2 `or_`
infixr 3 .|

class ToExpr a where
  type Result a
  toExpr :: a -> Angle (Result a)

instance ToExpr (Angle a) where
  type Result (Angle a) = a
  toExpr = id

instance ToExpr AngleStatement where
  type Result AngleStatement = ()
  toExpr stmt = sig unit `where_` [stmt]

instance ToExpr [AngleStatement] where
  type Result [AngleStatement] = ()
  toExpr stmts = sig unit `where_` stmts

-- | Build an or-pattern, `A | B`
(.|) :: (Result a ~ Result b, ToExpr a, ToExpr b) => a -> b -> Angle (Result a)
a .| b = Angle $ OrPattern DSL <$> gen (toExpr a) <*> gen (toExpr b)

-- | Build an or-pattern between statements (deprecated: just use '.|' instead)
or_ :: [AngleStatement] -> [AngleStatement] -> Angle ()
or_ = (.|)

-- | Build a key-value pattern, `A -> B`
(.->) :: Angle a -> Angle b -> Angle c
a .-> b = Angle $ KeyValue DSL <$> gen a <*> gen b

instance IsString (Angle Text) where
  fromString s = Angle (pure (String DSL (Text.pack s)))

-- | A string expression. Note that literal strings may be used as Angle
-- expressions directly if the @OverloadedStrings@ extension is
-- enabled.
string :: Text -> Angle Text
string t = Angle (pure (String DSL t))

nat :: Word64 -> Angle Nat
nat w = Angle (pure (Nat DSL w))

byte :: Word8 -> Angle Byte
byte w = Angle (pure (Nat DSL (fromIntegral w)))

-- | HACK: A typeclass to represent enums generated from Angle. This is used
-- to splice values of these Angle-generated Thrift enum values back into Angle
-- queries.
--
-- Ideally implementations of these should be codegenerated or we should come
-- up with a better way of embedding enum values in this DSL.
class AngleEnum a where
  type AngleEnumTy a
  enumName :: a -> Text

enum :: AngleEnum a => a -> Angle (AngleEnumTy a)
enum e = Angle (pure (Enum DSL (enumName e)))

stringPrefix :: Text -> Angle Text
stringPrefix t = Angle (pure (StringPrefix DSL t))

-- | A query for a literal value of type [byte]
byteArray :: ByteString -> Angle ByteString
byteArray b = Angle (pure (String DSL (Text.decodeUtf8 b)))

-- | Build an array expression
array :: [Angle t] -> Angle [t]
array xs = Angle $ Array DSL <$> mapM gen xs

-- | Build a tuple
tuple :: AngleTuple a => a -> Angle (AngleTupleTy a)
tuple = fromTuple

type SourceField = Angle.Field SpanAngleDSL () SourceRef SourceRef

-- | Match a record. Zero or more of the fields may be matched.
--
-- >   rec $
-- >     field @"abc" x $
-- >     field @"def" y $
-- >   end
--
rec :: HasFields f (RecordFields t) => Fields f -> Angle t
rec fs = Angle $ Struct DSL <$> go fs []
  where
  go :: forall f . Fields f -> [SourceField] -> State Int [SourceField]
  go f acc = case f of
    NoFields -> return acc
    Field l v rest  -> do
      v' <- gen v
      go rest (Angle.Field l v' : acc)

-- ToDo: we will also probably want a way to cast a record to a
-- subtype, and to emit a type signature.

-- | Match an alternative of a sum type
--
-- >   alt @"abc" value .| alt @"def" value
--
alt
  :: forall (l :: Symbol) v r .
     (KnownSymbol l, HasField l v (SumFields r)) =>
     Angle v
  -> Angle r
alt val = Angle $ do
  v <- gen val
  return $ Struct DSL
    [Angle.Field (Text.pack (symbolVal (Proxy @l))) v]

data Fields :: TFields -> * where
  NoFields :: Fields 'TNoFields
  Field :: Text -> Angle v -> Fields fs -> Fields ('TField l v fs)

field
  :: forall (l :: Symbol) v fs . KnownSymbol l =>
     Angle v
  -> Fields fs
  -> Fields ('TField l v fs)
field val rest = Field (Text.pack (symbolVal (Proxy @l))) val rest

end :: Fields 'TNoFields
end = NoFields

wild :: Angle t
wild = Angle $ pure $ Wildcard DSL

never :: Angle t
never = Angle $ pure (Variable DSL "never")

-- | Use this when you want a variable to match a nested predicate
-- rather than its key.
--
-- >  var $ \(n :: Angle Hack.Name) ->
-- >    n `where_` [
-- >      stmt $ predicate @Hack.MethodDeclaration $
-- >        rec $
-- >          field @"name" (asPredicate n) $
-- >          field @"container"
-- >             (rec $
-- >                field @"class_"
-- >                   (rec (field @"name" "MyClass" end))
-- >              end) $
-- >         end
-- >    ]
--
asPredicate :: Angle p -> Angle (KeyType p)
asPredicate (Angle a) = Angle a

-- | Sometimes the Angle typechecker needs a type signature.
-- This adds a type signature for any Angle type.
sig :: forall p. Type p => Angle p -> Angle p
sig a = Angle $ do
  ta <- gen a
  return $ TypeSignature DSL ta $ sourceType (Proxy @p)

class AngleTuple a where
  type AngleTupleTy a
  fromTuple :: a -> Angle (AngleTupleTy a)

instance AngleTuple (Angle a, Angle b) where
  type AngleTupleTy (Angle a, Angle b) = (a,b)
  fromTuple (a, b) = Angle $ do
    ta <- gen a
    tb <- gen b
    return $ Tuple DSL [ta,tb]

instance AngleTuple (Angle a, Angle b, Angle c) where
  type AngleTupleTy (Angle a, Angle b, Angle c) = (a,b,c)
  fromTuple (a,b,c) = Angle $ do
    ta <- gen a
    tb <- gen b
    tc <- gen c
    return $ Tuple DSL [ta,tb,tc]

instance AngleTuple (Angle a, Angle b, Angle c, Angle d) where
  type AngleTupleTy (Angle a, Angle b, Angle c, Angle d) = (a,b,c,d)
  fromTuple (a,b,c,d) = Angle $ do
    ta <- gen a
    tb <- gen b
    tc <- gen c
    td <- gen d
    return $ Tuple DSL [ta,tb,tc,td]

instance AngleTuple (Angle a, Angle b, Angle c, Angle d, Angle e) where
  type AngleTupleTy (Angle a, Angle b, Angle c, Angle d, Angle e) = (a,b,c,d,e)
  fromTuple (a,b,c,d,e) = Angle $ do
    ta <- gen a
    tb <- gen b
    tc <- gen c
    td <- gen d
    te <- gen e
    return $ Tuple DSL [ta,tb,tc,td,te]

instance AngleTuple (Angle a, Angle b, Angle c, Angle d, Angle e, Angle f) where
  type AngleTupleTy (Angle a, Angle b, Angle c, Angle d, Angle e, Angle f) =
    (a,b,c,d,e,f)
  fromTuple (a,b,c,d,e,f) = Angle $ do
    ta <- gen a
    tb <- gen b
    tc <- gen c
    td <- gen d
    te <- gen e
    tf <- gen f
    return $ Tuple DSL [ta,tb,tc,td,te,tf]

-- | Build a fact id with an explicit type @$123 : Type.2@
factId :: forall p. Predicate p => IdOf p -> Angle p
factId = sig . factId_

-- | Build a fact id without an explicit type, @$ 123@
factId_ :: forall p. IdOf p -> Angle p
factId_ = Angle
  . return
  . FactId DSL Nothing
  . fromIntegral
  . fromFid
  . idOf

-- | Convert @factIds (1 :| [2,3]) :: NonEmpty (IdOf p)@ into
-- @($1 : p.v) ++ $2 ++ $3@ for building angle
-- queries, where @v@ is the version of predicate @p@.
-- Uses disjuction, and nubOrds the input list
factIds :: forall p. Predicate p => NonEmpty (IdOf p) -> Angle p
factIds xs = sig $ foldr1 (.|) (fmap factId_ (nubOrd' xs))
  where
    -- Our version of package extra lacks Data.List.NonEmpty.Extra
    nubOrd' :: Ord a => NonEmpty a -> NonEmpty a
    nubOrd' = NonEmpty.fromList . List.nubOrd . NonEmpty.toList

-- | Similar to `factIds` but generate an array of ids with the signature of the
-- first element only. This doesn't nubOrd the input set
factIdsArray :: forall p. Predicate p => [IdOf p] -> Angle [p]
factIdsArray [] = array []
factIdsArray (x:xs) = array (sig (factId_ x) : map factId_ xs)

elementsOf :: Angle [x] -> Angle x
elementsOf listOfX = Angle $ do
  xs <- gen listOfX
  return (ElementsOfArray DSL xs)

arrayPrefix :: NonEmpty (Angle x) -> Angle x
arrayPrefix pats = Angle $ do
  pats <- traverse gen pats
  return (ArrayPrefix DSL pats)

unit :: Angle ()
unit = Angle $ pure (Struct DSL  [])

if_ :: Angle cond -> Angle a -> Angle a -> Angle a
if_ cond t e = Angle $ IfPattern DSL <$> gen cond <*> gen t <*> gen e

true :: Angle Bool
true = Angle $ pure (Variable DSL "true")

false :: Angle Bool
false = Angle $ pure (Variable DSL "false")

-- | Lift a boolean literal into Angle
bool :: Bool -> Angle Bool
bool True = true
bool False = false

just :: Angle a -> Angle (Maybe a)
just x = Angle $ do
  tx <- gen x
  pure (Struct DSL [Angle.Field "just" tx])

nothing :: Angle (Maybe a)
nothing = Angle $ pure (Variable DSL "nothing")

{-
  TODO:
  | KeyValue (SourcePat_ v t) (SourcePat_ v t)
-}

-- Primitives

(.+) :: (ToExpr a, ToExpr b, Result a ~ Nat, Result b ~ Nat)
     => a -> b -> Angle Nat
a .+ b =
  Angle $ App DSL (Variable DSL "prim.addNat") <$>
  sequence [gen $ toExpr a, gen $ toExpr b]

zipp :: Angle [a] -> Angle [b] -> Angle [(a,b)]
zipp a b =
  Angle $ App DSL (Variable DSL "prim.zip") <$>
  sequence [gen a, gen b]

conc :: Angle [a] -> Angle [a] -> Angle [a]
conc a b =
  Angle $ App DSL (Variable DSL "prim.concat") <$>
  sequence [gen a, gen b]

-- -----------------------------------------------------------------------------
-- Special queries

-- | Query only new facts of the given predicate, i.e. facts in the
-- topmost DB in the stack.
--
-- (This is a special-purpose query tool for building incremental fact
-- derivers)
new :: Predicate p => Angle p -> Angle p
new = specialPredicateQuery "new"

-- | Query only old facts of the given predicate, i.e. facts from all
-- but the topmost DB in the stack.
--
-- (This is a special-purpose query tool for building incremental fact
-- derivers)
old :: Predicate p => Angle p -> Angle p
old = specialPredicateQuery "old"

specialPredicateQuery :: Predicate p => Text -> Angle p -> Angle p
specialPredicateQuery suffix (Angle pat) = Angle $ do
  p <- pat
  case p of
    App DSL (Variable DSL p) args ->
      return $ App DSL (Variable DSL (p <> "#" <> suffix)) args
    _ -> error ": not a predicate"

-- -----------------------------------------------------------------------------
-- Type checking for records and sum types

-- | Maps a Thrift-generated record type to its fields
type family RecordFields t :: TFields

-- | Maps a Thrift-generated union type to its fields
type family SumFields t :: TFields

-- | @HasFields f t@ asserts that all the fields @f@ are present in
-- @t@, in any order.
class HasFields (f :: TFields) (t :: TFields)

instance HasFields 'TNoFields f
instance (HasField l v f, HasFields r f) => HasFields ('TField l v r) f

-- | @HasField l v f@ asserts that field @l@ is present in @f@ with type @v@
class HasField (l :: Symbol) v (f :: TFields)

instance {-# OVERLAPPING #-} (v ~ v') => HasField l v ('TField l v' f)
instance HasField l v f => HasField l v ('TField l' v' f)

data TFields where
  TNoFields :: TFields
  TField :: Symbol -> v -> TFields -> TFields