{-
  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 TypeOperators #-}
module Thrift.Compiler.Plugin
  ( Typecheckable(..), SomeLit(..)
  , qualify, qualifyType
  , getPrefix
  , lowercase, uppercase, toCamel
  , fixLeadingUnderscore
  , toConstructorName
  , getNamespace
  , isNewtype
  , filterHsAnns, getTypeAnns
  ) where

import Data.Bifunctor
import Data.List
import qualified Data.Map.Strict as Map
import Data.Maybe
import Data.Some
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Type.Equality
import Prelude hiding (Enum)

import Thrift.Compiler.Options
import Thrift.Compiler.Parser
import Thrift.Compiler.Typechecker.Monad
import Thrift.Compiler.Types

data SomeLit l = forall t. ThisLit (Type l t) t

-- | Typeclass to define Thrift typecheckable Languages. There are several
-- sections of functions to define here, although many have default
-- implementations
class Monoid (Interface l) => Typecheckable l  where
  -- | Interface that gets generated by the Thrift compiler. This includes all
  -- the symbols exported by the generated code in the target language. It is
  -- used in order to add additional required symbols that get referenced in
  -- splice (hs_include) files. If your language does not support splice files,
  -- then set Interface l = ()
  type Interface l

  -- * Annotation Processing

  -- | Given a resolved Thrift type and a list of annotations, produce some new
  -- transformed type. This is how TSpecial types get generated
  resolveTypeAnnotations
    :: Type l t
    -> [Annotation Loc]
    -> TC l (Some (Type l))
  resolveTypeAnnotations ty _ = pure $ Some ty

  -- | Recursively qualify all of the named types in a SpecialType so that they
  -- can be properly identified in imports
  qualifySpecialType :: (Text, Text) -> SpecialType l t -> SpecialType l t

  -- * Typechecking

  -- | Typecheck literals of special, language specific types
  -- Note: this is *only* for literals, named constants are typechecked by the
  -- main typecheckConst function (even for special types)
  typecheckSpecialConst
    :: SpecialType l t
    -> UntypedConst Loc
    -> TC l (TypedConst l t)

  -- | Type equality for special types
  eqSpecial :: SpecialType l u -> SpecialType l v -> Maybe (u :~: v)

  -- | Generate an interface. See the Interface type family for more info on how
  -- this is used.
  getInterface :: Options l -> ThriftFile a Loc -> Interface l
  getInterface _ _ = mempty

  -- | Get extra Thrift symbols that are referenced in splice files
  getExtraSymbols ::
    Options l -> Interface l -> ThriftFile SpliceFile b -> [Text]
  getExtraSymbols _ _ _ = []

  -- * Renamers

  -- | Rename the thrift module, probably taking namespace into account
  renameModule :: Options l -> ThriftFile a b -> Text
  renameModule _ ThriftFile{..} = thriftName

  renameStruct :: Options l -> Struct s u a -> Text
  renameStruct _ Struct{..} = structName

  renameField
    :: Options l
    -> [Annotation a]
    -> Text
    -> Field u s m a
    -> Text
  renameField _ _ _ Field{..} = fieldName

  renameConst :: Options l -> Text -> Text
  renameConst _ = id

  renameService :: Options l -> Service s u a -> Text
  renameService _ Service{..} = serviceName

  renameFunction :: Options l -> Function s u a -> Text
  renameFunction _ Function{..} = funName

  renameTypedef :: Options l -> Typedef s u a -> Text
  renameTypedef _ Typedef{..} = tdName

  renameEnum :: Options l -> Enum s u a -> Text
  renameEnum _ Enum{..} = enumName

  renameEnumAlt :: Options l -> Enum s u a -> Text -> Text
  renameEnumAlt _ _ name = name

  renameUnion :: Options l -> Union s u a -> Text
  renameUnion _ Union{..} = unionName

  renameUnionAlt :: Options l -> Union s u a -> UnionAlt s u a -> Text
  renameUnionAlt _ _ UnionAlt{..} = altName

  getUnionEmptyName :: Options l -> Union s u a -> Text
  getUnionEmptyName _ Union{..} = unionName <> "_EMPTY"

  -- * Uniqueness options specify whether various renamed symbols need to be
  -- globally unique to the entire Thrift (if False, they are still unique
  -- within their structure)

  -- | Are field names globally unique? If yes, then we will report renamed
  -- field name collisions across structs
  fieldsAreUnique :: Options l -> Bool
  fieldsAreUnique _ = False

  -- | Are union alternatives unique?
  unionAltsAreUnique :: Options l -> Bool
  unionAltsAreUnique _ = False

  -- | Are enum values unique?
  enumAltsAreUnique :: Options l -> Bool
  enumAltsAreUnique _ = False

  -- | Get enum flavour from annotation tags
  enumFlavourTag :: Options l -> Enum s u a -> EnumFlavour
  enumFlavourTag _ _ = SumTypeEnum False

  -- * Back translators
  -- Translate Stuff Back to regular thrift for pretty printing and JSON output

  -- | Translate a special type back to its underlying Thrift type, and its
  -- special name in the target language.
  -- Note that this should be a shallow conversion. Do not backtranslate
  -- recursively
  backTranslateType :: SpecialType l t -> (Some (Type l), Text)

  -- | Translate a literal to its underlying thrift type and representation.
  -- Like @backTranslateType@, the conversion should be shallow
  backTranslateLiteral :: SpecialType l t -> t -> SomeLit l

-- Env Qualifier ---------------------------------------------------------------

-- | Qualify all the named types in an Env
qualify :: Typecheckable l => (Text, Text) -> Env l -> Env l
qualify m env = env
  { typeMap  = qualCtx qualST (typeMap env)
  , schemaMap = qualSchemas (schemaMap env)
  , unionMap = qualSchemas (unionMap env)
  , constMap = qualCtx qualC (constMap env)
  , enumMap  = qualEnums (enumMap env)
  , serviceMap = qualServices (serviceMap env)
  }
  where
    qualCtx
      :: (a -> a)
      -> Context a
      -> Context a
    qualCtx f ctx@Context{..} = ctx { cMap = Map.map f cMap }
    qualST :: Typecheckable l => Some (Type l) -> Some (Type l)
    qualST (Some ty) = Some $ qualifyType m ty
    qualC :: Typecheckable l => (Some (Type l), Name, Loc)
                             -> (Some (Type l), Name, Loc)
    qualC (st, n, loc) = (qualST st, qualName m n, loc)
    qualSchemas
      :: Typecheckable l
      => Map.Map Text (Some (SCHEMA l t)) -> Map.Map Text (Some (SCHEMA l t))
    qualSchemas = Map.map $ \(Some schema) -> Some (qualifySchema m schema)
    qualEnums = Map.map $ \(vs, ns) ->
      (Map.map (first (qualName m)) vs, Map.map (first (qualName m)) ns)
    qualServices = Map.map $ \(a, b, c) -> (qualName m a, b, c)

qualifyType :: Typecheckable l => (Text, Text) -> Type l t -> Type l t
qualifyType _ I8 = I8
qualifyType _ I16 = I16
qualifyType _ I32 = I32
qualifyType _ I64 = I64
qualifyType _ TFloat = TFloat
qualifyType _ TDouble = TDouble
qualifyType _ TBool = TBool
qualifyType _ TBytes = TBytes
qualifyType _ TText = TText
qualifyType m (TSet u)       = TSet $ qualifyType m u
qualifyType m (THashSet u)   = THashSet $ qualifyType m u
qualifyType m (TList u)      = TList $ qualifyType m u
qualifyType m (TMap k v)     = TMap (qualifyType m k) (qualifyType m v)
qualifyType m (THashMap k v) = THashMap (qualifyType m k) (qualifyType m v)
qualifyType m (TTypedef name t loc) =
  TTypedef (qualName m name) (qualifyType m t) loc
qualifyType m (TStruct name loc)    = TStruct (qualName m name) loc
qualifyType m (TException name loc) = TException (qualName m name) loc
qualifyType m (TUnion name loc)     = TUnion (qualName m name) loc
qualifyType m (TEnum name loc nounknown) =
  TEnum (qualName m name) loc nounknown
qualifyType m (TNewtype name t loc) =
  TNewtype (qualName m name) (qualifyType m t) loc
qualifyType m (TSpecial ty) = TSpecial $ qualifySpecialType m ty

qualifySchema :: Typecheckable l => (Text, Text) -> SCHEMA l t s -> SCHEMA l t s
qualifySchema _ SEmpty = SEmpty
qualifySchema m (SField p n t s) =
  SField p n (qualifyType m t) (qualifySchema m s)
qualifySchema m (SReqField p n t s) =
  SReqField p n (qualifyType m t) (qualifySchema m s)
qualifySchema m (SOptField p n t s) =
  SOptField p n (qualifyType m t) (qualifySchema m s)

qualName :: (Text, Text) -> Name -> Name
qualName (tm, rm) Name{..} = Name
  { sourceName = qualName_ tm sourceName
  , resolvedName = qualName_ rm resolvedName
  }

qualName_ :: Text -> Name_ s -> Name_ s
qualName_ m (UName n) = QName m n
qualName_ _ n@QName{} = n

-- Renamer Helpers -------------------------------------------------------------

getPrefix :: [Annotation a] -> Maybe Text
getPrefix anns = listToMaybe
  [ pre
  | ValueAnn{vaVal=TextAnn pre _,..} <- anns
  , vaTag == "hs.prefix"
  ]

lowercase :: Text -> Text
lowercase = uncurry (<>) . first Text.toLower . Text.splitAt 1

uppercase :: Text -> Text
uppercase = uncurry (<>) . first Text.toUpper . Text.splitAt 1

toCamel :: Text -> Text
toCamel = Text.concat . map capitalize . Text.splitOn "_"
  where capitalize = uncurry (<>) . first Text.toUpper . Text.splitAt 1

-- Prepend "TU" for "ThriftUnderscore"
fixLeadingUnderscore :: Text -> Text
fixLeadingUnderscore = uncurry (<>) . first fixUnderscore . Text.splitAt 1
  where
    fixUnderscore s = if s == "_" then "TU__" else s

toConstructorName :: Text -> Text
toConstructorName = fixLeadingUnderscore . uppercase

-- | Select the last namespace header
getNamespace :: Text -> [Header s l a] -> Maybe Text
getNamespace l = foldl' getNS Nothing
  where
    getNS ns HInclude{} = ns
    getNS ns HPackage{} = ns
    getNS ns HNamespace{..}
      | nmLang == l = Just nmName
      | otherwise   = ns

isNewtype :: [Annotation a] -> Bool
isNewtype = any isNewtypeAnn
  where
    isNewtypeAnn SimpleAnn{..} = saTag == "hs.newtype"
    isNewtypeAnn _ = False

filterHsAnns :: [Annotation a] -> [Annotation a]
filterHsAnns = filter $ Text.isPrefixOf "hs." . \case
  SimpleAnn{..} -> saTag
  ValueAnn{..}  -> vaTag

getTypeAnns :: Text -> [Annotation a] -> [(Text, Annotation a)]
getTypeAnns lang anns =
  [ (ty, a) | a@ValueAnn{vaVal=TextAnn ty _,..} <- anns, vaTag == typeTag ]
  where
    typeTag = lang <> ".type"