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

module Glean.Schema.Gen.Utils
  ( -- * Types
    unitT
  , isByt
  , shouldNameKeyType
    -- * Dependency analysis
  , SomeDecl(..)
  , orderDecls
    -- * Names and namespaces
  , NameSpaces
  , splitDot
  , joinDot
  , dotted
  , underscored
  , sortDeclsByNamespace
    -- * Text utils
  , myUnlines
  , newline
  , cap1
  , low1
    -- * Monad for code generation
  , M
  , Env
  , NamePolicy(..)
  , mkNamePolicy
  , runM
  , typeDef
  , repType
  , NewOrOld(..)
  , nameThisType
  , withPredicateDefHint
  , withRecordFieldHint
  , withUnionFieldHint
  , withTypeDefHint
  , withTopLevelHint
  , pushDefs
  , popDefs
  , predicateName
  , typeName
  , thriftName
  , haskellThriftName
  , addNamespaceDependencies
  , buckOncallAnnotation
  , Oncall
  ) where

import Control.Monad.Reader
import Control.Monad.State
import Data.Graph
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Text (Text)
import qualified Data.Text as Text
import TextShow

import Util.List

import Glean.Angle.Types
import Glean.Schema.Types
import Glean.Schema.Util


-- -----------------------------------------------------------------------------
-- Types

unitT :: ResolvedType
unitT = RecordTy []

isByt :: ResolvedType -> Bool
isByt ByteTy = True
isByt _ = False

-- | Return 'True' if we should generate a new type in Thrift for a
-- predicate key type.  If the type is a base type or is already a
-- typedef, there's no need to generate a new named type.
shouldNameKeyType :: ResolvedType -> Bool
shouldNameKeyType RecordTy{} = True
shouldNameKeyType SumTy{} = True
shouldNameKeyType _ = False

-- -----------------------------------------------------------------------------
-- Dependency analysis of decls

data SomeDecl
  = PredicateDecl ResolvedPredicateDef
  | TypeDecl TypeRef ResolvedType

type Node = (Name,Version)

-- | Sort the declarations to keep C++ Happy.
orderDecls :: [SomeDecl] -> [SomeDecl]
orderDecls decls = map betterNotBeAnyCyclesIn sccs
 where
  betterNotBeAnyCyclesIn (AcyclicSCC one) = one
  betterNotBeAnyCyclesIn (CyclicSCC defs) =
    error $ "cycle: " ++ unwords (map (show . declNode) defs)

  sccs = stronglyConnComp
    [ (decl, declNode decl, outEdges decl) | decl <- decls ]

  declNode :: SomeDecl -> (Name,Version)
  declNode (PredicateDecl PredicateDef{..}) =
    (predicateRef_name predicateDefRef, predicateRef_version predicateDefRef)
  declNode (TypeDecl tref _) =
    (typeRef_name tref, typeRef_version tref)

  outEdges :: SomeDecl -> [Node]
  outEdges d = case d of
    PredicateDecl PredicateDef{..} ->
      outEdgesTs [predicateDefKeyType, predicateDefValueType]
    TypeDecl _ ty ->
      outEdgesT ty

  outEdgesTs = concatMap outEdgesT
  outEdgesFields fields = outEdgesTs [ ty | FieldDef _ ty <- fields ]

  outEdgesT :: ResolvedType -> [Node]
  outEdgesT ByteTy{} = []
  outEdgesT NatTy{} = []
  outEdgesT BooleanTy{} = []
  outEdgesT StringTy{} = []
  outEdgesT (ArrayTy ty) = outEdgesT ty
  outEdgesT (MaybeTy ty) = outEdgesT ty
  outEdgesT (RecordTy fields)  = outEdgesFields fields
  outEdgesT (SumTy fields)  = outEdgesFields fields
  outEdgesT (SetTy ty) = outEdgesT ty
  outEdgesT (NamedTy _ (TypeRef name ver)) = [(name,ver)]
  outEdgesT PredicateTy{} = [] -- See Note [predicate type references]
  outEdgesT EnumeratedTy{} = []
  outEdgesT TyVar{} = error "outEdgesT: TyVar"
  outEdgesT HasTy{} = error "outEdgesT: HasTy"
  outEdgesT HasKey{} = error "outEdgesT: HasKey"
  outEdgesT ElementsOf{} = error "outEdgesT: ElementsOf"

{- Note [predicate type references]

In a struct definition, references to predicates look like this:

  struct SrcLoc {
    Fact<File> file;
    ...
  }

The structure of Fact<P> doesn't depend on the full definition of
P, it's enough to have a "struct P;" declaration, which we have for
all predicates.

In a predicate definition, references to other predicates are like this:

  struct Target : Predicate<Sys::Blob> {..}

And Predicate<> contains types only, so this also doesn't need the
full definition of Sys::Blob.

Therefore, we conclude that predicate references do not constitute a
dependency for the purposes of ordering definitions, and this enables
us to break all cycles involving predicates.  There can of course be
no cycles between typedefs; these will cause an error.
-}

-- -----------------------------------------------------------------------------
-- Names and namespaces

-- Convenience to split a name

joinDot :: (NameSpaces, Text) -> Text
joinDot (ns,t) = Text.intercalate "." (ns ++ [t])

dotted :: NameSpaces -> Text
dotted = Text.intercalate "."

underscored :: NameSpaces -> Text
underscored = Text.intercalate "_"

sortDeclsByNamespace
  :: [ResolvedPredicateDef]
  -> [ResolvedTypeDef]
  -> HashMap NameSpaces ([ResolvedPredicateDef], [ResolvedTypeDef])
sortDeclsByNamespace preds types =
  HashMap.fromListWith combine $
    [ (namespace, ([pred],[]))
    | pred@PredicateDef{..} <- preds
    , let (namespace,_) = splitDot (predicateRef_name predicateDefRef) ] ++
    [ (namespace, ([],[tdef]))
    | tdef@TypeDef{..} <- types
    , let (namespace,_) = splitDot (typeRef_name typeDefRef)  ]
  where
    combine (as,bs) (cs,ds) = (as++cs, bs++ds)

-- -----------------------------------------------------------------------------
-- Text Utils

-- | Better than Text.unlines for codegen purposes
myUnlines :: [Text] -> Text
myUnlines = Text.intercalate newline
          . map Text.stripEnd
          . dropWhile Text.null . dropWhileEnd Text.null
  where
    dropWhileEnd :: (a -> Bool) -> [a] -> [a]
    dropWhileEnd p = foldr (\x xs -> if p x && null xs then [] else x : xs) []

newline :: Text
newline = Text.pack "\n"

cap1 :: Text -> Text
cap1 tIn = if Text.null tIn then tIn else
  let (h, t) = Text.splitAt 1 tIn
  in Text.toUpper h <> t

low1 :: Text -> Text
low1 tIn = if Text.null tIn then tIn else
  let (h, t) = Text.splitAt 1 tIn
  in Text.toLower h <> t

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

-- | Thrift must turn tuples into named structures.  Solve the hard
-- problem of naming inside monad M
type M = ReaderT Env (State ExtraDefs)

type ExtraDefs = [Text]

data Env = Env
  { nameHint :: [Text]
      -- ^ stack of "hints": name components that will be used to
      -- to name an anonymous type from the schema. The hint at
      -- any given point is unique: a type generated from this name
      -- will not clash with anything else.
      --
      -- Note: the 'Text' are sometimes 'Safe' and sometimes not 'Safe'
  , namePolicy :: NamePolicy
  , typeEnv :: HashMap (Name, Version) ResolvedTypeDef
  }

data NamePolicy = NamePolicy
  { predNames :: HashMap PredicateRef (NameSpaces,Text)
  , typeNames :: HashMap TypeRef (NameSpaces,Text)
  }

mkNamePolicy :: [ResolvedPredicateDef] -> [ResolvedTypeDef] -> NamePolicy
mkNamePolicy preds types = NamePolicy{..}
  where
  nameMap = HashMap.fromListWith (++) $
    [ (predicateRef_name predicateDefRef,
       [predicateRef_version predicateDefRef])
    | PredicateDef{..} <- preds ] ++
    [ (typeRef_name typeDefRef, [typeRef_version typeDefRef])
    | TypeDef{..} <- types ]

  -- Figure out how we're going to name each predicate. The most
  -- recent version of P gets to be called P, all the older versions
  -- are called P_v.
  predNames = HashMap.fromList
    [ let
        (ns, name) = splitDot predicateRef_name
      in
      if predicateRef_version == maximum versions
        then (predicateDefRef, (ns, name))
        else (predicateDefRef, (ns, name <> "_" <> showt predicateRef_version))
    | PredicateDef{..} <- preds
    , let PredicateRef{..} = predicateDefRef
    , Just versions <- [HashMap.lookup predicateRef_name nameMap] ]

  -- Similarly for type names.
  typeNames = HashMap.fromList
    [ let
        (ns, name) = splitDot typeRef_name
      in
      if typeRef_version == maximum versions
        then (typeDefRef, (ns, name))
        else (typeDefRef, (ns, name <> "_" <> showt typeRef_version))
    | TypeDef{..} <- types
    , let TypeRef{..} = typeDefRef
    , Just versions <- [HashMap.lookup typeRef_name nameMap] ]


-- | Look up the name to use for a predicate, taking into account versions
predicateName :: Monad m => PredicateRef -> ReaderT Env m (NameSpaces, Text)
predicateName pred = do
  Env{..} <- ask
  case HashMap.lookup pred (predNames namePolicy) of
    Just (ns,x) -> return (ns,x)
    _ -> error $ "predicateThriftName: " ++ show pred

-- | Look up the name to use for a type, taking into account versions
typeName :: Monad m => TypeRef -> ReaderT Env m (NameSpaces, Text)
typeName typeref = do
  Env{..} <- ask
  case HashMap.lookup typeref (typeNames namePolicy) of
    Just (ns,x) -> return (ns,x)
    Nothing ->
      -- it's one we invented, the name should be unique
      return (splitDot (typeRef_name typeref))

-- | The textual name to use for an entity relative to the current namespace
thriftName :: NameSpaces -> (NameSpaces, Text) -> Text
thriftName here (ns, x)
  | here == ns = getSafe $ safeThrift x
  | otherwise = getSafe $ case ns of
    [] -> checkSafe ("thriftName " ++ show (here,(ns,x))) x
    _  -> Safe $ underscored ns <> "." <> getSafe (safeThrift x)

-- | Witness that the name is fully mangled (see 'safeThrift' for the
-- defining smart constructor)
newtype Safe x = Safe { getSafe :: x }

-- | This asserts the input already qualifies as 'Safe' and that 'safeThift'
-- would not transform it.  If 'safeThift' would transform the input then
-- this throws an error.
checkSafe :: String -> Text -> Safe Text
checkSafe msg x =
  let safe = safeThrift x
  in if x == getSafe safe
        then safe
        else error $ msg <> " : " <> show x <> " not really Safe in checkSafe"

safeThrift :: Text -> Safe Text
safeThrift x = Safe x

-- The name of the Thrift type as it appears in Haskell
haskellThriftName :: (NameSpaces, Text) -> Text
haskellThriftName (ns, x) =
  dotted $
    prefix ++
    [ Text.concat (map cap1 ns), "Types", getSafe $ safeThrift x ]
  where
  prefix = ["Glean", "Schema"]

runM
  :: s
  -> NamePolicy
  -> [ResolvedTypeDef]
  -> ReaderT Env (State s) a
  -> (a, s)

runM initState namePolicy types act = (result, finalState)
  where
    (result, finalState) = runState (runReaderT act env) initState
    env = Env
      { namePolicy = namePolicy
      , nameHint = mempty
      , typeEnv = HashMap.fromList
          [ ((typeRef_name typeDefRef, typeRef_version typeDefRef), typedef)
          | typedef@TypeDef{..} <- types ]
      }


-- | Returned by 'nameThisType'
data NewOrOld = New | Old
  deriving (Eq, Ord)

-- Make a name for a type. The returned name is assumed to be in the
-- current namespace.
nameThisType :: Monad m => ResolvedType' s -> ReaderT Env m (NewOrOld, Text)
nameThisType (RecordTy []) = return (Old, "builtin.Unit")
nameThisType _ = do
  Env{..} <- ask
  let
    name = Text.intercalate "_" $ case reverse nameHint of
      (h:hs) -> cap1 h : hs
      [] -> error "nameThisType: empty hint"
  return (New, name)

repType :: Monad m => ResolvedType -> ReaderT Env m (Maybe ResolvedType)
repType (NamedTy _ tr) = do
  maybeTy <- typeDef tr
  case maybeTy of
    Nothing -> return Nothing
    Just ty -> repType ty
repType t = return (Just t)

typeDef :: Monad m => TypeRef -> ReaderT Env m (Maybe ResolvedType)
typeDef TypeRef{..} = do
  te <- typeEnv <$> ask
  case HashMap.lookup (typeRef_name, typeRef_version) te of
    Just TypeDef{..} -> return (Just typeDefType)
    Nothing -> return Nothing -- we don't know


-- | set the 'hint' (in a nested scope) for 'nameThisType' while running 'act'
withHint :: Monad m => Text -> ReaderT Env m a -> ReaderT Env m a
withHint hint = local pushHint
  where pushHint s = s{ nameHint = hint : nameHint s }

-- | For a typedef, we only add this as a hint component if we're at
-- the top level and therefore generating a user-supplied
-- typedef. Otherwise we will get the names of generated typedefs as
-- hints themselves.
withTypeDefHint :: Monad m => Text -> ReaderT Env m a -> ReaderT Env m a
withTypeDefHint hint = local pushHint
  where
  pushHint s = s{ nameHint = if null (nameHint s) then [hint] else nameHint s }

withTopLevelHint :: Monad m => Text -> ReaderT Env m a -> ReaderT Env m a
withTopLevelHint hint = local $ \s -> s { nameHint = [hint] }

withPredicateDefHint :: Monad m => Text -> ReaderT Env m a -> ReaderT Env m a
withPredicateDefHint = withHint

withRecordFieldHint :: Monad m => Text -> ReaderT Env m a -> ReaderT Env m a
withRecordFieldHint = withHint

{- |
For a union field, there's a potential clash to avoid:

> union T {
>    T_x x;
> }
>
> struct T_x {  // generated type
>   ...
> }

We cannot use T_x for the name of the generated type, because it
clashes with the T_x constructor that the hs2 thrift compiler
generates for the field x of T.

So, we add another underscore to the hint for a union field, giving us
T_x_.
-}
withUnionFieldHint :: Monad m => Text -> ReaderT Env m a -> ReaderT Env m a
withUnionFieldHint p = withHint (p <> "_")

pushDefs :: [Text] -> M ()
pushDefs ts = modify (reverse ts ++)

popDefs :: M [Text]
popDefs = state $ \ s -> (reverse s, mempty)


addNamespaceDependencies
  :: HashMap NameSpaces ([ResolvedPredicateDef], [ResolvedTypeDef])
  -> HashMap NameSpaces
    ([NameSpaces], [ResolvedPredicateDef], [ResolvedTypeDef])
addNamespaceDependencies nss =
  HashMap.fromList
    [ (ns, (outEdges ns preds types, preds, types))
    | (ns, (preds, types)) <- HashMap.toList nss
    ]
  where
  outEdges ns preds types =
    uniq $ filter (/= ns) $
      ["builtin"] :  -- builtin is always a dep, so we can get builtin.Unit
      concat (map outEdgesPred preds ++ map outEdgesTypeDef types)

  outEdgesPred :: ResolvedPredicateDef -> [NameSpaces]
  outEdgesPred PredicateDef{..} =
    outEdgesTs [predicateDefKeyType, predicateDefValueType]

  outEdgesTypeDef :: ResolvedTypeDef -> [NameSpaces]
  outEdgesTypeDef TypeDef{..} = outEdgesT typeDefType

  outEdgesTs = concatMap outEdgesT
  outEdgesFields fields = outEdgesTs [ ty | FieldDef _ ty <- fields ]

  outEdgesT :: ResolvedType -> [NameSpaces]
  outEdgesT ByteTy{} = []
  outEdgesT NatTy{} = []
  outEdgesT BooleanTy{} = []
  outEdgesT StringTy{} = []
  outEdgesT (ArrayTy ty) = outEdgesT ty
  outEdgesT (MaybeTy ty) = outEdgesT ty
  outEdgesT (RecordTy fields)  = outEdgesFields fields
  outEdgesT (SumTy fields)  = outEdgesFields fields
  outEdgesT (SetTy ty) = outEdgesT ty
  outEdgesT (NamedTy _ (TypeRef name _)) = [fst (splitDot name)]
  outEdgesT (PredicateTy _ (PredicateRef name _)) = [fst (splitDot name)]
  outEdgesT EnumeratedTy{} = []
  outEdgesT TyVar{} = error "outEdgesT: TyVar"
  outEdgesT HasTy{} = error "outEdgesT: HasTy"
  outEdgesT HasKey{} = error "outEdgesT: HasKey"
  outEdgesT ElementsOf{} = error "outEdgesT: ELementsOf"

type Oncall = Text
buckOncallAnnotation :: Maybe Oncall -> Text
buckOncallAnnotation name = case name of
  Just oncall -> "\noncall(\"" <> oncall <> "\")"
  Nothing -> ""