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

module Glean.Glass.Annotations
  ( getAnnotationsForEntity

  -- * testing
  , prettyHackAnnotation
  ) where


import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Map as Map
import Data.Maybe (catMaybes)
import Control.Monad (forM )

import qualified Glean
import Glean.Glass.Base ( GleanPath(..) )
import Glean.Angle as Angle
import qualified Glean.Haxl.Repos as Glean
import qualified Glean.Schema.Codemarkup.Types as Code
import qualified Glean.Schema.Code.Types as Code
import qualified Glean.Schema.CodeCxx.Types as Cxx1
import qualified Glean.Schema.CodeHack.Types as Hack
import qualified Glean.Schema.CodePython.Types as Python
import qualified Glean.Schema.CodeJava.Types as Java
import qualified Glean.Schema.Cxx1.Types as Cxx1
import qualified Glean.Schema.Hack.Types as Hack
import qualified Glean.Schema.Python.Types as Python
import qualified Glean.Schema.JavaAlpha.Types as Java
import qualified Glean.Schema.JavakotlinAlpha.Types as JavaKotlin

import Glean.Glass.SymbolId (entityToAngle, toSymbolId)
import Glean.Glass.SymbolId.Java ( flattenPath )
import qualified Glean.Glass.Types as Glass
import Glean.Glass.Utils ( searchRecursiveWithLimit )
import Glean.Glass.Path ( fromGleanPath )

-- | For Hack, Java and Python, the annotations are a single array fact.
-- For C++ they're sprinkled across each function decl, forcing us to search
max_annotations_limit :: Int
max_annotations_limit = 10

getAnnotationsForEntity
  :: Glass.RepoName
  -> Code.Entity
  -> Glean.RepoHaxl u w (Either Text (Maybe [Glass.Annotation]))
getAnnotationsForEntity repo entity = fetch repo `mapM` entityToAngle entity

-- Get all the annotations for a given entity. This may include
-- computing symbolId for these annotations.
fetch
  :: Glass.RepoName -> Angle Code.Entity
  -> Glean.RepoHaxl u w (Maybe [Glass.Annotation])
fetch repo ent = do
  anns <- queryAnnotations ent
  annsWithSyms <- forM anns (annotationsToSymbols repo)
  getAnnotations annsWithSyms

-- | Entity to its annotation facts.
queryAnnotations :: Angle Code.Entity -> Glean.RepoHaxl u w [Code.Annotations]
queryAnnotations entity =
  fmap fst <$> searchRecursiveWithLimit (Just max_annotations_limit) $
    entityAnnotations entity

entityAnnotations :: Angle Code.Entity -> Angle Code.Annotations
entityAnnotations ent =
  var $ \anns -> anns `where_` [
    wild .= Angle.predicate @Code.EntityToAnnotations (
        rec $
          field @"entity" ent $
          field @"annotations" anns
        end)
  ]

--
-- Per-language unpacking of annotations
--

-- Maps an Annotations (e.g. a list of annotation, to their possible
-- corresponding symbolId). Codemarkup doesn't expose a single Annotation
-- type, so we use zipped list. Currently, the symbolId are only used for
-- Hack. For the other languages, we use an empty list of symbolIds.
newtype AnnotationsSymbolId =
  AnnotationsSymbolId (Code.Annotations, [Maybe Glass.SymbolId])

annotationsToSymbols
  :: Glass.RepoName
  -> Code.Annotations
  -> Glean.RepoHaxl u w AnnotationsSymbolId
annotationsToSymbols repo annotations = case annotations of
  Code.Annotations_hack (Hack.Annotations_attributes attrs) -> do
    attributeToDecl <- queryAttributeToDecl attrs
    decls <- forM attributeToDecl (declarationsToSymbolId repo)
    let mapAnnotSym = Map.fromList decls
        syms = (`Map.lookup` mapAnnotSym) <$> attrs
    return $ AnnotationsSymbolId (annotations, syms)
  _ -> return $ AnnotationsSymbolId (annotations, [])

-- Generate a query of the form
--   hack.AttributeToDeclaration.6 {attribute = X0}
--   where X0 = [$388098 : hack.UserAttribute.6,
--               $41135,
--               $139800678][..]"]
-- From AttributeToDeclaration facts, we get the info
-- required to build an Annotation SymbolId
queryAttributeToDecl
  :: [Hack.UserAttribute]
  -> Glean.RepoHaxl u w [Hack.AttributeToDeclaration]
queryAttributeToDecl [] = return []
queryAttributeToDecl attrs = fmap fst <$> searchRecursiveWithLimit
    (Just (length attrs)) $
  Angle.predicate @Hack.AttributeToDeclaration $
    rec $
      field @"attribute" (asPredicate (elementsOf factIds))
    end
  where
    factIds = factIdsArray (map Glean.getId attrs)

-- | Unpack the decl and file fields of the annotation lookup and
-- construct a symbol id
declarationsToSymbolId
  :: Glass.RepoName
  -> Hack.AttributeToDeclaration
  -> Glean.RepoHaxl u w (Hack.UserAttribute, Glass.SymbolId)
declarationsToSymbolId repo attrDecl = do
  Hack.AttributeToDeclaration_key{..} <- Glean.keyOf attrDecl
  path <- Glean.keyOf attributeToDeclaration_key_file
  let userAttribute = attributeToDeclaration_key_attribute
  let decl = attributeToDeclaration_key_declaration
      entity = Code.Entity_hack (Hack.Entity_decl decl) -- do this in Angle
  sym <- toSymbolId (fromGleanPath repo (GleanPath path)) entity
  return (userAttribute, sym)

--
-- The type of annotations is language specific. We unwrap here in the client
-- this also lets us do a bit of language-specific processing
--
-- We should think about encoding this in the schema Glean-side
--
class HasAnnotations a where
  getAnnotations :: a -> Glean.RepoHaxl u w (Maybe [Glass.Annotation])

instance HasAnnotations a => HasAnnotations [a] where
  getAnnotations [] = pure Nothing
  getAnnotations anns = do
    result <- mapM getAnnotations anns
    return $ Just (concat (catMaybes result))
    -- we could drop empty results here, and not emit the attr at all?

instance HasAnnotations AnnotationsSymbolId where
  getAnnotations (AnnotationsSymbolId (ann, syms)) = case (ann, syms) of
    (Code.Annotations_cxx ann, _) -> getAnnotations ann
    (Code.Annotations_hack ann, syms) -> getAnnotations (ann, syms)
    (Code.Annotations_python anns, _) -> getAnnotations anns
    (Code.Annotations_java anns, _) -> getAnnotations anns
    (Code.Annotations_fbthrift{}, _) -> pure Nothing -- Not yet supported
    (Code.Annotations_EMPTY, _) -> pure Nothing

instance HasAnnotations Cxx1.Annotations where
  getAnnotations (Cxx1.Annotations_attributes anns) = getAnnotations anns
  getAnnotations Cxx1.Annotations_EMPTY = pure Nothing

instance HasAnnotations Cxx1.Attribute where
  getAnnotations Cxx1.Attribute{attribute_key=Just key} = return $ Just
    [ Glass.Annotation
        { annotation_source = key
        , annotation_name = key
        , annotation_symbol = Nothing
        }
    ]
  getAnnotations _ = pure Nothing

instance HasAnnotations Python.Annotations where
  getAnnotations (Python.Annotations_decorators anns) = getAnnotations anns
  getAnnotations Python.Annotations_EMPTY = pure Nothing

instance HasAnnotations Python.Decorator where
  getAnnotations key = return $ Just $ pure Glass.Annotation {
          annotation_source = key,
          annotation_symbol = Nothing,
          annotation_name = key
        }

instance HasAnnotations Java.Annotations where
  getAnnotations (Java.Annotations_annotations anns) = getAnnotations anns
  getAnnotations Java.Annotations_EMPTY = pure Nothing

instance HasAnnotations Java.Annotation where
  getAnnotations ann = do
    Java.Annotation_key mname _ctor _constant _span <- Glean.keyOf ann
    JavaKotlin.QName_key qName mcontext <- Glean.keyOf mname
    nameStr <- Glean.keyOf qName
    context <- Glean.keyOf mcontext
    toks <- flattenPath context
    let annKey = Text.intercalate "." (nameStr : reverse toks)

    return $ Just $ pure Glass.Annotation {
          annotation_source = nameStr,
          annotation_symbol = Nothing, -- annotations don't have definitions yet
          annotation_name = annKey
        }

instance HasAnnotations (Hack.Annotations, [Maybe Glass.SymbolId]) where
  getAnnotations (Hack.Annotations_attributes anns, syms) =
    getAnnotations (anns, syms)
  getAnnotations (Hack.Annotations_EMPTY, _) = pure Nothing

instance HasAnnotations ([Hack.UserAttribute], [Maybe Glass.SymbolId]) where
  getAnnotations (attrs, syms) = getAnnotations $ zip attrs syms

instance HasAnnotations (Hack.UserAttribute, Maybe Glass.SymbolId) where
  getAnnotations
    (Hack.UserAttribute
      {
        userAttribute_key = Just Hack.UserAttribute_key
          {
            userAttribute_key_name=Hack.Name{name_key=Just name}
          , userAttribute_key_parameters=args
          }
      }, symbol) = return $ Just
          [ Glass.Annotation
              { annotation_source = prettyHackAnnotation name args
              , annotation_name = name
              , annotation_symbol = symbol
              }
          ]
  getAnnotations _ = pure Nothing

prettyHackAnnotation :: Text -> [Text] -> Text
prettyHackAnnotation name [] = name
prettyHackAnnotation name args = name <> parens (commas (map quotes args))
  where
    parens x = "(" <> x <> ")"
    commas = Text.intercalate ", "
    quotes arg = "\"" <> arg <> "\""