{-
  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 Thrift.Compiler.GenEnum
  ( genEnumImports
  , genEnumDecl
  ) where

import Prelude hiding (Enum)
import Control.Monad
import Data.List
import Data.List.Extra
import Data.Text (Text)
import Language.Haskell.Exts.Syntax hiding (Type, Decl)
import qualified Data.Set as Set
import qualified Language.Haskell.Exts.Syntax as HS
import TextShow

import Thrift.Compiler.GenConst
import Thrift.Compiler.GenTypedef
import Thrift.Compiler.GenUtils
import Thrift.Compiler.Plugin
import Thrift.Compiler.Plugins.Haskell
import Thrift.Compiler.Types

-- Data Type Declaration -------------------------------------------------------

genEnumImports :: Set.Set Import
genEnumImports = Set.fromList
  [ QImport "Prelude" "Prelude"
  , QImport "Control.Exception" "Exception"
  , QImport "Control.DeepSeq" "DeepSeq"
  , QImport "Data.Aeson" "Aeson"
  , QImport "Data.Default" "Default"
  , QImport "Data.Function" "Function"
  , QImport "Data.HashMap.Strict" "HashMap"
  , QImport "Data.Hashable" "Hashable"
  , QImport "Data.Int" "Int"
  , QImport "GHC.Magic" "GHC"
  , SymImport "Prelude" [ ".", "++", ">", "==" ]
  ]

genEnumDecl :: HS Enum -> [HS.Decl ()]
genEnumDecl Enum{ enumFlavour=PseudoEnum{..},..} =
  genTypedefDecl typedef (not peThriftEnum) ++
  concatMap genConstDecl consts ++
  -- Instances
  if not peThriftEnum
  then mzero
  else
    [ genDefault enumResolvedName enumConstants
    , genShow enumResolvedName enumConstants
    , genThriftEnumInst enumResolvedName enumConstants True True
    ]
  where
    typedef = Typedef
      { tdName = enumName
      , tdResolvedName = enumResolvedName
      , tdTag = IsNewtype
      , tdType = AnnotatedType enumValueType Nothing (Arity0Loc nlc)
      , tdResolvedType = enumValueType
      , tdLoc = TypedefLoc nlc nlc
      , tdAnns = Nothing
      , tdSAnns = []
      }
    consts = flip map enumConstants $ \EnumValue{..} -> Const
      { constName = evName
      , constResolvedName = evResolvedName
      , constType = AnnotatedType (TNamed enumName) Nothing (Arity0Loc nlc)
      , constResolvedType =
        TNewtype (mkName enumName enumResolvedName) enumValueType noLoc
      , constVal =
        UntypedConst nlc $ IntConst (fromIntegral evValue) (showt evValue)
      , constResolvedVal = Literal $ New $ fromIntegral evValue
      , constLoc = ConstLoc nlc nlc nlc NoSep
      , constSAnns = []
      }
genEnumDecl Enum{ enumFlavour=SumTypeEnum{..},..} =
  -- Enum Declaration
  [ DataDecl () (DataType ()) Nothing
    (DHead () $ textToName enumResolvedName)
    -- We generate them in sorted order so that we can derive Bounded correctly
    ((genConstr <$> sortOn evValue enumConstants) ++ [genUnknownConstr | not enumNoUnknown])
    -- Deriving
    (if enumNoUnknown && null enumConstants
     then mzero
     else
       pure $ deriving_ $ map (IRule () Nothing Nothing . IHCon ()) $
       [ qualSym "Prelude" "Eq"
       , qualSym "Prelude" "Show"
       ] ++ [ qualSym "Prelude" "Ord" | canDeriveOrd ])
  ] ++
  -- Instances
  if enumNoUnknown && null enumConstants
  then
    map (genEmptyInstance enumResolvedName)
    -- Using the symbol (==) in the AST is technically wrong, but it
    -- generates correct pretty-printed code and allows us to reuse more code
    [ ("Prelude", "Eq", [ "(==)" ])
    , ("Prelude", "Show", [ "show" ])
    , ("Prelude", "Ord", [ "compare" ])
    , ("Aeson", "ToJSON", [ "toJSON" ])
    , ("Default", "Default", [ "def" ])
    , ("Hashable", "Hashable", [ "hashWithSalt" ])
    , ("DeepSeq", "NFData", [ "rnf" ])
    ] ++
    [genThriftEnumInst enumResolvedName enumConstants enumNoUnknown False]
  else
    [ genToJSON enumResolvedName
    , genNFData enumResolvedName
    , genDefault enumResolvedName enumConstants
    , genHashable enumResolvedName
    , genThriftEnumInst enumResolvedName enumConstants enumNoUnknown False
    ] ++
    [genOrd enumResolvedName | not canDeriveOrd]
  where
    genConstr :: HS EnumValue -> QualConDecl ()
    genConstr EnumValue{..} =
      QualConDecl () Nothing Nothing
                      (ConDecl () (textToName evResolvedName) [])
    -- If the stars align we can derive the Enum instance
    -- This requires the Enum to contain exactly the values [0 .. n-1]
    canDeriveOrd = and $ zipWith (==) [0..] $ sort $ map evValue enumConstants

    -- Use 2 underscores to avoid name collisions.
    genUnknownConstr :: QualConDecl ()
    genUnknownConstr =
      QualConDecl
        ()
        Nothing
        Nothing
        (ConDecl () (textToName $ enumResolvedName <> "__UNKNOWN") [genType (TSpecial HsInt)])

-- Ord Instance ----------------------------------------------------------------

genOrd :: Text -> HS.Decl ()
genOrd name =
  InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
       (IHCon () $ qualSym "Prelude" "Ord")
       (TyCon () $ unqualSym name))
    (Just
     [ InsDecl () $ FunBind ()
       [ Match () (textToName "compare") []
         (UnGuardedRhs () $
          qvar "Function" "on" `app`
          qvar "Prelude" "compare" `app`
          qvar "Thrift" "fromThriftEnum")
         Nothing
       ]
     ])

-- Generate Show Instance --------------------------------------------------

genShow :: Text -> [HS EnumValue] -> HS.Decl ()
genShow name consts =
  InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
       (IHCon () $ qualSym "Prelude" "Show")
       (TyCon () $ unqualSym name))
    (Just
     [ InsDecl () $ FunBind ()
       [ Match () (textToName "showsPrec")
         [ pvar "__d", PApp () (unqualSym name) [pvar "__val"] ]
         (UnGuardedRhs () $
          Case () (qvar "HashMap" "lookup" `app` var "__val" `app` var "__m")
            [ Alt () (PApp () (qualSym "Prelude" "Just") [pvar "__s"])
              (UnGuardedRhs () $ qvar "Prelude" "showString" `app` var "__s")
              Nothing
            , Alt () (PApp () (qualSym "Prelude" "Nothing") [])
              (UnGuardedRhs () $ qvar "Prelude" "showParen" `app`
                infixApp ">" (var "__d") (intLit (10 :: Int)) `app`
                infixApp "."
                  (qvar "Prelude" "showString" `app`
                    stringLit (name <> "__UNKNOWN "))
                  (qvar "Prelude" "showsPrec" `app`
                    intLit (11 :: Int) `app` var "__val"))
              Nothing
            ])
         (Just $ BDecls () [FunBind () [Match () (textToName "__m") []
          (UnGuardedRhs () $ qvar "HashMap" "fromList" `app` pairs)
          Nothing]])
       ]
     ])
  where
    -- https://gitlab.haskell.org/ghc/ghc/-/issues/4505
    chunkLimit = 1000
    pairs
      | length consts <= chunkLimit = pairList consts
      | otherwise = qvar "Prelude" "concat" `app` listE
        (flip map (chunksOf chunkLimit consts) $ \chunk ->
          qvar "GHC" "noinline" `app` qvar "Prelude" "id" `app` pairList chunk)
    pairList cs = listE
      [ Tuple () Boxed [intLit evValue, stringLit $ uppercase evResolvedName]
      | EnumValue{..} <- cs
      ]

-- Aeson Instances -------------------------------------------------------------

genToJSON :: Text -> HS.Decl ()
genToJSON name =
  InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
       (IHCon () $ qualSym "Aeson" "ToJSON")
       (TyCon () $ unqualSym name))
    (Just
     [ InsDecl () $ FunBind ()
       [ Match () (textToName "toJSON") []
         (UnGuardedRhs () $
          qvar "Aeson" "toJSON" `compose` qvar "Thrift" "fromThriftEnum")
         Nothing
       ]
     ])

-- Generate NFData Instance ---------------------------------------------------

genNFData :: Text -> HS.Decl ()
genNFData name =
  InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
       (IHCon () $ qualSym "DeepSeq" "NFData")
       (TyCon () $ unqualSym name))
    (Just
     [ InsDecl () $ FunBind ()
       [ Match () (textToName "rnf")
         [ PApp () (unqualSym arg) [] ]
         (UnGuardedRhs () $
          qvar "Prelude" "seq" `app` var arg `app` unit)
         Nothing
       ]
     ])
    where
      arg = "__" <> name
      unit = Con () (Special () (UnitCon ()))

-- Generate Default Instance --------------------------------------------------

genDefault :: Text -> [HS EnumValue] -> HS.Decl ()
genDefault name consts =
  InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
       (IHCon () $ qualSym "Default" "Default")
       (TyCon () $ unqualSym name))
    (Just
     [ InsDecl () $ FunBind ()
       [ Match () (textToName "def") []
         (UnGuardedRhs () $
          case consts of
            EnumValue{..} : _ -> con evResolvedName
            [] ->
              qvar "Exception" "throw" `app`
              (qvar "Thrift" "ProtocolException" `app`
               stringLit ("def: enum " <> name <> " has no constructors")))
         Nothing
       ]
     ])

-- Generate Hashable Instance -------------------------------------------------

genHashable :: Text -> HS.Decl ()
genHashable name =
  InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
       (IHCon () $ qualSym "Hashable" "Hashable")
       (TyCon () $ unqualSym name))
    (Just
     [ InsDecl () $ FunBind ()
       [ Match () (textToName "hashWithSalt") [ pvar "_salt", pvar "_val" ]
         (UnGuardedRhs () $
          qvar "Hashable" "hashWithSalt" `app` var "_salt" `app`
          (qvar "Thrift" "fromThriftEnum" `app` var "_val"))
         Nothing
       ]
     ])

-- Generate Empty Instance -----------------------------------------------------

genEmptyInstance :: Text -> (Text, Text, [Text]) -> HS.Decl ()
genEmptyInstance name (mname, className, methods) =
  InstDecl () Nothing
    (IRule () Nothing Nothing $
     IHApp ()
       (IHCon () $ qualSym mname className)
       (TyCon () $ unqualSym name))
    (Just $ map
     (\method -> InsDecl () $ FunBind () [genEmptyMethod name method])
     methods)

genEmptyMethod :: Text -> Text -> Match ()
genEmptyMethod name method =
  Match () (textToName method) []
  (UnGuardedRhs () $
   qvar "Exception" "throw" `app`
   (qvar "Thrift" "ProtocolException" `app`
    stringLit
    (mconcat
     [ method, ": Thrift enum '", name, "' is uninhabited"])))
  Nothing

-- Thrift Enum Instance --------------------------------------------------------

genThriftEnumInst :: Text -> [HS EnumValue] -> Bool -> Bool -> HS.Decl ()
genThriftEnumInst ename consts enumNoUnknown pseudoEnum =
  InstDecl () Nothing
    (IRule () Nothing Nothing
      (IHApp ()
       (IHCon () (qualSym "Thrift" "ThriftEnum"))
       (TyCon () (unqualSym ename))))
    (Just $ map (InsDecl () . FunBind ())
       [ genThriftEnumMethod "toThriftEnum" genToEnumMatch
          genToEnumCatchAll genToEnumUnknown genToEnumPseudo
       , genThriftEnumMethod "fromThriftEnum" genFromEnumMatch
          genFromEnumCatchAll genFromEnumUnknown genFromEnumPseudo
       , genAllEnumValues consts
       , if pseudoEnum then [genToEnumEitherPseudo]
         else map genToEnumEitherMatch consts ++ [genToEnumEitherUnknown]
       ]
    )
  where
    genThriftEnumMethod
      method genEnumMatch genEnumCatchAll genEnumUnknown genEnumPseudo
      | pseudoEnum = [genEnumPseudo]
      | enumNoUnknown && null consts = [genEmptyMethod ename method]
      | otherwise = map genEnumMatch consts ++
        [ if enumNoUnknown
          then genEnumCatchAll
          else genEnumUnknown
        ]
    genToEnumMatch :: HS EnumValue -> Match ()
    genToEnumMatch EnumValue{..} =
      Match ()
        (textToName "toThriftEnum")
        [ intP evValue ]
        (UnGuardedRhs () $ Con () $ unqualSym evResolvedName)
        Nothing
    genToEnumCatchAll =
      Match ()
        (textToName "toThriftEnum")
        [ pvar "_val" ]
        (UnGuardedRhs () $
          qvar "Exception" "throw" `app`
          (qvar "Thrift" "ProtocolException" `app`
          infixApp "++"
          (stringLit $ "toThriftEnum: not a valid identifier for enum " <>
                       ename <> ": ")
          (qvar "Prelude" "show" `app` var "_val")))
        Nothing
    genToEnumUnknown =
      Match ()
        (textToName "toThriftEnum")
        [ pvar "val" ]
        (UnGuardedRhs () $ var (ename <> "__UNKNOWN") `app` var "val")
        Nothing
    genToEnumPseudo =
      Match ()
        (textToName "toThriftEnum")
        [ pvar "__val" ]
        (UnGuardedRhs () $ var ename `app`
          (qvar "Prelude" "fromIntegral" `app` var "__val"))
        Nothing
    genFromEnumMatch :: HS EnumValue -> Match ()
    genFromEnumMatch EnumValue{..} =
      Match ()
        (textToName "fromThriftEnum")
        (if pseudoEnum
          then [PApp () (unqualSym ename) [intP evValue]]
          else [PApp () (unqualSym evResolvedName) []])
        (UnGuardedRhs () $ intLit evValue)
        Nothing
    genFromEnumCatchAll =
      Match ()
        (textToName "fromThriftEnum")
        [ pvar "_val" ]
        (UnGuardedRhs () $
          qvar "Exception" "throw" `app`
          (qvar "Thrift" "ProtocolException" `app`
          infixApp "++"
          (stringLit $ "fromThriftEnum: not a valid identifier for enum " <>
                       ename <> ": ")
          (qvar "Prelude" "show" `app` var "_val")))
        Nothing
    genFromEnumUnknown =
      Match ()
        (textToName "fromThriftEnum")
        [PApp () (unqualSym (ename <> "__UNKNOWN")) [pvar "val"]]
        (UnGuardedRhs () $ var "val")
        Nothing
    genFromEnumPseudo =
      Match ()
        (textToName "fromThriftEnum")
        [PApp () (unqualSym ename) [pvar "__val"]]
        (UnGuardedRhs () $ qvar "Prelude" "fromIntegral" `app` var "__val")
        Nothing
    genToEnumEitherMatch :: HS EnumValue -> Match ()
    genToEnumEitherMatch EnumValue{..} =
      Match ()
        (textToName "toThriftEnumEither")
        [ intP evValue ]
        (UnGuardedRhs () $ qvar "Prelude" "Right" `app` var evResolvedName)
        Nothing
    genToEnumEitherUnknown =
      Match ()
        (textToName "toThriftEnumEither")
        [ pvar "val" ]
        (UnGuardedRhs () $ qvar "Prelude" "Left" `app` infixApp "++"
          (stringLit $ "toThriftEnumEither: not a valid identifier for enum " <>
                     ename <> ": ")
          (qvar "Prelude" "show" `app` var "val"))
        Nothing
    genToEnumEitherPseudo =
      Match ()
        (textToName "toThriftEnumEither")
        [ pvar "val" ]
        (UnGuardedRhs () $ If () (qvar "Prelude" "elem"
          `app` var "__val" `app` qvar "Thrift" "allThriftEnumValues")
          (qvar "Prelude" "Right" `app` var "__val")
          (qvar "Prelude" "Left" `app` infixApp "++"
            (stringLit $ "toThriftEnumEither: not a valid identifier for enum " <>
                        ename <> ": ")
            (qvar "Prelude" "show" `app` var "val")))
        (Just $ BDecls () [FunBind () [Match () (textToName "__val") []
          (UnGuardedRhs () $ var ename `app`
            (qvar "Prelude" "fromIntegral" `app` var "val"))
          Nothing]])

genAllEnumValues :: [HS EnumValue] -> [Match ()]
genAllEnumValues cs =
  [ Match
      ()
      (textToName "allThriftEnumValues")
      []
      (UnGuardedRhs () $ listE (genEnumExp <$> sortOn evValue cs))
      Nothing
  ]
  where
    genEnumExp :: HS EnumValue -> Exp ()
    genEnumExp EnumValue{..} = Con () $ unqualSym evResolvedName