{-
  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 NamedFieldPuns #-}
module Thrift.Compiler.GenJSONLoc
  ( -- * Main generation
    genJSONLoc
  , writeJSONLoc
  , getAstPathLoc
   -- * Utility functions
  , displayAnnotatedType
  , genType
  ) where

import Prelude hiding (Enum)
import Data.Aeson
import Data.Aeson.Encode.Pretty
import Data.ByteString.Builder
import qualified Data.ByteString.Lazy as LBS
import qualified Data.HashMap.Strict as HashMap
import Data.Proxy
import Data.Some
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Lazy as Lazy
import qualified Data.Text.Lazy.Encoding as Lazy
import GHC.TypeLits
import System.Directory
import System.FilePath

import Util.Aeson

import Thrift.Compiler.Typechecker
import Thrift.Compiler.Types as Types
import Thrift.Compiler.Options as Options
import Thrift.Compiler.Plugin
-- import Thrift.Compiler.Types as Thrift hiding (noLoc)

writeJSONLoc
  :: Typecheckable l
  => Program l Loc         -- ^ Top level program ti generate
  -> Maybe [Program l Loc] -- ^ Dependencies if using recursive mode
  -> IO FilePath
writeJSONLoc prog deps = do
  createDirectoryIfMissing True dir
  LBS.writeFile path $ prettyJSON $ genJSONLoc prog deps
  return path
  where
    path = dir </> file
    (dir, file) = getAstPathLoc prog
    prettyJSON = encodePretty' defConfig { confCompare = compare }

getAstPathLoc :: Program l a -> (FilePath, FilePath)
getAstPathLoc Program{..} = (progOutPath, Text.unpack progName ++ ".ast")

genJSONLoc :: Typecheckable l => Program l Loc -> Maybe [Program l Loc] -> Value
genJSONLoc prog Nothing = Object $ genJSONProg prog
genJSONLoc prog (Just deps) = toJSON $ genJSONProg prog : map genJSONProg deps

genJSONProg :: Typecheckable l => Program l Loc -> Object
genJSONProg Program{..} = objectFromList
  [ "name"     .= progHSName
  , "path"     .= progPath
  , "includes" .= map Types.progPath progIncludes
  , "typedefs" .= map genTypedef dTdefs
  , "enums"    .= map genEnum dEnums
  , "consts"   .= map genConst dConsts
  , "structs"  .= map genStruct dStructs
  , "unions"   .= map genUnion dUnions
  , "services" .= map genService dServs
  , "options"  .= genOptions (options progEnv)
  ]
  where
    Decls{..} = partitionDecls progDecls

-- Locs ------------------------------------------------------------------------

tx :: Text -> Text
tx = id

displayLoc :: Loc -> Text
displayLoc = Text.pack . show

displayLocated :: Located Loc -> Text
displayLocated = displayLoc . lLocation

displayTypeLoc :: TypeLoc n Loc -> [Text]
displayTypeLoc x = map displayLocated $ case x of
   Arity0Loc{..} -> [a0Ty]
   Arity1Loc{..} -> [a1Ty, a1OpenBrace, a1CloseBrace]
   Arity2Loc{..} -> [a2Ty, a2OpenBrace, a2Comma, a2CloseBrace]

displaySeparator :: Separator Loc -> Object
displaySeparator (Semicolon loc) = objectFromList
  [ "sep_type" .= tx "Semicolon"
  , "loc" .= displayLocated loc
  ]
displaySeparator (Comma loc) = objectFromList
  [ "sep_type" .= tx "Comma"
  , "loc" .= displayLocated loc
  ]
displaySeparator NoSep = objectFromList
  [ "sep_type" .= tx"NoSep"
  ]

displayAnnValue :: AnnValue -> Object
displayAnnValue (IntAnn i v) = objectFromList
  [ "ann_value_type" .= tx "IntAnn"
  , "i" .= i
  , "v" .= v
  ]
displayAnnValue (TextAnn t q) = objectFromList
  [ "ann_value_type" .=  tx "TextAnn"
  , "t" .= t
  , "q" .= Text.pack (show q)
  ]

displayAnnotation :: Annotation Loc -> Object
displayAnnotation SimpleAnn{..} = objectFromList
  [ "ann_type" .= tx "SimpleAnn"
  , "ann_tag" .= saTag
  , "loc" .= displayLocated saLoc
  , "sep" .= displaySeparator saSep
  ]
displayAnnotation ValueAnn{..} = objectFromList
  [ "ann_type" .= tx "ValueAnn"
  , "tag" .= vaTag
  , "val" .= displayAnnValue vaVal
  , "loc_tag" .= displayLocated vaTagLoc
  , "loc_equal" .= displayLocated vaEqual
  , "loc_val" .= displayLocated vaValLoc
  , "sep" .= displaySeparator vaSep
  ]

displayAnnotations :: Annotations Loc -> Object
displayAnnotations Annotations{..} = objectFromList
  [ "loc_open" .= displayLocated annOpenParen
  , "loc_close" .= displayLocated annCloseParen
  , "loc_ann_list" .= map displayAnnotation annList
  ]

displayAnnotatedType :: forall t. AnnotatedType Loc t -> Object
displayAnnotatedType AnnotatedType{..} = objectFromList
  [ "type" .= (genTType (atType :: Un t) :: Object)
  , "anns" .= maybe Null (Object . displayAnnotations) atAnnotations
  , "loc" .= displayTypeLoc atLoc
  ]

-- Reconstruct -----------------------------------------------------------------

-- The resolved Loc is in mkTypemap, mkSchemaMap, etc in typecheckModule.
-- I expect that the cross-ref are discovered in envLookup and envCtxLookup
-- when run during resolveDecls, but only implicitly.
--
-- Here the correspondance is reconstructed

reconstructXRef
  :: Typecheckable l
  =>  AnnotatedType Loc v
  -> Type l t
  -> [(Loc, Loc)]
reconstructXRef atIn@AnnotatedType{atType, atLoc} rt = case (atType, rt) of
    (TSet at1, TSet rt1) -> reconstructXRef at1 rt1
    (THashSet at1, THashSet rt1) -> reconstructXRef at1 rt1
    (TList at1, TList rt1) -> reconstructXRef at1 rt1
    (TMap ak av, TMap rk rv) ->
      reconstructXRef ak rk ++ reconstructXRef av rv
    (THashMap ak av, THashMap rk rv) ->
      reconstructXRef ak rk ++ reconstructXRef av rv
    (TNamed{}, TTypedef _ _ rtLoc) -> [(getTypeLoc atLoc, rtLoc)]
    (TNamed{}, TNewtype _ _ rtLoc) -> [(getTypeLoc atLoc, rtLoc)]
    (TNamed{}, TStruct _ rtLoc) -> [(getTypeLoc atLoc, rtLoc)]
    (TNamed{}, TException _ rtLoc) -> [(getTypeLoc atLoc, rtLoc)]
    (TNamed{}, TUnion _ rtLoc) -> [(getTypeLoc atLoc, rtLoc)]
    (TNamed{}, TEnum _ rtLoc _) -> [(getTypeLoc atLoc, rtLoc)]
    (_, TSpecial st) -> case backTranslateType st of
      (Some rtSimple, _) -> reconstructXRef atIn rtSimple
    _ -> []

displayXRef
  ::  Typecheckable l
  =>  AnnotatedType Loc v
  -> Type l t
  -> [Value]
displayXRef at rt = map oneXRef (reconstructXRef at rt)

-- | JSON encode a pair of location, fst is usage and snd is definition. Want
-- to hyperlink from usage to destination, and perhaps list all usages of the
-- destination.
oneXRef :: (Loc, Loc) -> Value
oneXRef (aLoc, rLoc) = Object $ objectFromList
  [ "aLoc" .= displayLoc aLoc
  , "rLoc" .= displayLoc rLoc ]

-- | Make this notice when the const value is another const or enum value
-- and reference the defintion. Enrich the Identifier and EnumVal
-- constructors to easily link them.
reconstructXRefConst
  :: Typecheckable l
  => UntypedConst Loc
  -> TypedConst l t
  -> Type l t
  -> [(Loc, Loc)]
reconstructXRefConst UntypedConst{ucLoc} (Identifier _name _rt rLoc) _ =
  [(lLocation ucLoc, rLoc)]
reconstructXRefConst UntypedConst{ucLoc} (WeirdEnumToInt _ _ _ rLoc) _ =
  [(lLocation ucLoc, rLoc)]
reconstructXRefConst UntypedConst{ucLoc} (Literal ev) (TEnum _ _loc _) =
  let EnumVal _name rLocVal = ev in [(lLocation ucLoc, rLocVal)]
reconstructXRefConst _ Literal{} _ = []

-- | Eventually make this notice when the const value is another const
-- and hyperlink that value (Identifier)
displayXRefConst
  :: Typecheckable l
  => UntypedConst Loc
  -> TypedConst l t
  -> Type l t
  -> [Value]
displayXRefConst uc tc ty = map oneXRef (reconstructXRefConst uc tc ty)

-- Typedefs --------------------------------------------------------------------

genTypedef :: Typecheckable l => Typedef 'Resolved l Loc -> Object
genTypedef Typedef{..} = objectFromList
  [ "name" .= tdResolvedName
  , "type" .= genType tdResolvedType
  , "ann_type" .= displayAnnotatedType tdType
  , "newtype" .= case tdTag of { IsNewtype -> True ; IsTypedef -> False }
  , "loc_keyword" .= displayLocated (tdlKeyword tdLoc)
  , "loc_name" .= displayLocated (tdlName tdLoc)
  , "anns" .= maybe Null (Object . displayAnnotations) tdAnns
  , "xref" .= displayXRef tdType tdResolvedType
  ]

-- Enums -----------------------------------------------------------------------

genEnum :: Typecheckable l => Enum 'Resolved l Loc -> Object
genEnum Enum{..} = objectFromList
  [ "name"      .= enumResolvedName
  , "constants" .= map genEnumConst enumConstants
  , "flavour" .= case enumFlavour of
      SumTypeEnum{} -> "sum_type" :: Text
      PseudoEnum{} -> "pseudo"
  , "loc_keyword" .= displayLocated (slKeyword enumLoc)
  , "loc_name" .= displayLocated (slName enumLoc)
  ]

genEnumConst :: EnumValue 'Resolved l Loc -> Object
genEnumConst EnumValue{..} = objectFromList
  [ "name"  .= evResolvedName
  , "value" .= evValue
  , "loc_name" .= displayLocated (evlName evLoc)
  ]

-- Constants -------------------------------------------------------------------

genConst :: Typecheckable l => Const 'Resolved l Loc -> Object
genConst Const{..} = objectFromList
  [ "name"  .= constResolvedName
  , "type"  .= genType constResolvedType
  , "value" .= genConstVal constResolvedType constResolvedVal
  , "ann_type" .= displayAnnotatedType constType
  , "loc_keyword" .= displayLocated (clKeyword constLoc)
  , "loc_name" .= displayLocated (clName constLoc)
  , "xref" .= (displayXRef constType constResolvedType
                ++ displayXRefConst constVal constResolvedVal constResolvedType)
  ]

-- Structs, Exceptions, and Unions ---------------------------------------------

genStruct :: Typecheckable l => Struct 'Resolved l Loc -> Object
genStruct Struct{..} = objectFromList
  [ "name" .= structResolvedName
  , "struct_type" .= case structType of
      StructTy    -> "STRUCT" :: Text
      ExceptionTy -> "EXCEPTION"
  , "fields" .= map genField structMembers
  , "loc_keyword" .= displayLocated (slKeyword structLoc)
  , "loc_name" .= displayLocated (slName structLoc)
  ]

genField :: Typecheckable l => Field u 'Resolved l Loc -> Object
genField Field{..} = objectFromList $
  [ "name"  .= fieldResolvedName
  , "id"    .= fieldId
  , "type"  .= genType fieldResolvedType
  , "xref"  .= displayXRef fieldType fieldResolvedType
  , "loc_name" .= displayLocated (flName fieldLoc)
  ] ++
  (case fieldResolvedVal of
     Nothing -> []
     Just val -> [ "default_value" .= genConstVal fieldResolvedType val ]) ++
  (case fieldTag of
     STRUCT_FIELD -> [ "requiredness" .=
                       case fieldRequiredness of
                         Default    -> "default" :: Text
                         Required{} -> "required"
                         Optional{} -> "optional"
                     ]
     _ -> [])

genUnion :: Typecheckable l => Union 'Resolved l Loc -> Object
genUnion Union{..} = objectFromList
  [ "name"   .= unionResolvedName
  , "fields" .= map genAlt unionAlts
  , "loc_keyword" .= displayLocated (slKeyword unionLoc)
  , "loc_name" .= displayLocated (slName unionLoc)
  ]

genAlt :: Typecheckable l => UnionAlt 'Resolved l Loc -> Object
genAlt UnionAlt{..} = objectFromList
  [ "name" .= altResolvedName
  , "id"   .= altId
  , "type" .= genType altResolvedType
  , "loc_name" .= displayLocated (flName altLoc)
  ]

-- Services and Functions ------------------------------------------------------

genService :: Typecheckable l => Service 'Resolved l Loc -> Object
genService s@Service{..} = objectFromList $
  [ "name"      .= serviceResolvedName
  , "functions" .= map genFunction (getServiceFunctions s)
  , "loc_keyword" .= displayLocated (slKeyword serviceLoc)
  , "loc_name" .= displayLocated (slName serviceLoc)
  ] ++
  (case serviceSuper of
     Nothing -> []
     Just Super{..} -> [ "super" .= genName (fst supResolvedName) ])

genFunction :: Typecheckable l => Function 'Resolved l Loc -> Object
genFunction Function{..} = objectFromList
  [ "name" .= funResolvedName
  , "return_type" .= case funResolvedType of
      Nothing -> simpleType "void"
      Just ty -> withSome ty genType
  , "args"   .= map genField funArgs
  , "throws" .= map genField funExceptions
  , "oneway" .= funIsOneWay
  , "loc_name" .= displayLocated (fnlName funLoc)
  ]


-- Unresolved Types and Constants ----------------------------------------------

type Un t = TType 'Unresolved () Loc t

genTType :: Un t -> Object

-- Base Types
genTType I8  = simpleType "byte"
genTType I16 = simpleType "i16"
genTType I32 = simpleType "i32"
genTType I64 = simpleType "i64"
genTType TFloat  = simpleType "float"
genTType TDouble = simpleType "double"
genTType TBool   = simpleType "bool"
genTType TText   = simpleType "string"
genTType TBytes  = simpleType "binary"

-- Collections
genTType (TSet u)       = collectionTType "set" u
genTType (THashSet u)   = collectionTType "hash_set" u
genTType (TList u)      = collectionTType "list" u
genTType (TMap k v)     = mapTType "map" k v
genTType (THashMap k v) = mapTType "hash_map" k v

-- Named Types

genTType (TNamed n) = simpleName n

collectionTType :: Text -> AnnotatedType Loc t -> Object
collectionTType tyName u = objectFromList
  [ "type" .= tyName
  , "inner_type" .= displayAnnotatedType u
  ]

mapTType :: Text -> AnnotatedType Loc k -> AnnotatedType Loc v -> Object
mapTType tyName k v = objectFromList
  [ "type" .= tyName
  , "key_type" .= displayAnnotatedType k
  , "val_type" .= displayAnnotatedType v
  ]

simpleName :: Text -> Object
simpleName tyName = objectFromList ["name" .= (String tyName)]

-- Types and Constants ---------------------------------------------------------

genType :: Typecheckable l => Type l t -> Object

-- Base Types
genType I8  = simpleType "byte"
genType I16 = simpleType "i16"
genType I32 = simpleType "i32"
genType I64 = simpleType "i64"
genType TFloat  = simpleType "float"
genType TDouble = simpleType "double"
genType TBool   = simpleType "bool"
genType TText   = simpleType "string"
genType TBytes  = simpleType "binary"

-- Collections
genType (TSet u)       = collectionType "set" u
genType (THashSet u)   = collectionType "hash_set" u
genType (TList u)      = collectionType "list" u
genType (TMap k v)     = mapType "map" k v
genType (THashMap k v) = mapType "hash_map" k v

-- Named Types
genType (TStruct name loc)    = namedType "struct" name loc
genType (TException name loc) = namedType "exception" name loc
genType (TUnion name loc)     = namedType "union" name loc
genType (TEnum name loc _)      = namedType "enum" name loc
genType (TTypedef name ty loc) = objectFromList
  [ "type" .= ("typedef" :: Text)
  , "name" .= genName name
  , "inner_type" .= genType ty
  , "loc" .= displayLoc loc
  ]
genType (TNewtype name ty loc) = objectFromList
  [ "type" .= ("newtype" :: Text)
  , "name" .= genName name
  , "inner_type" .= genType ty
  , "loc" .= displayLoc loc
  ]
genType (TSpecial ty) = case backTranslateType ty of
  (Some u, tag) -> genType u <> objectFromList [ "special" .= tag ]

simpleType :: Text -> Object
simpleType tyName = objectFromList ["type" .= (String tyName)]

collectionType :: Typecheckable l => Text -> Type l t -> Object
collectionType tyName u = objectFromList
  [ "type" .= tyName
  , "inner_type" .= genType u
  ]

mapType :: Typecheckable l => Text -> Type l u -> Type l v -> Object
mapType tyName k v = objectFromList
  [ "type" .= tyName
  , "key_type" .= genType k
  , "val_type" .= genType v
  ]

namedType :: Text -> Name -> Loc -> Object
namedType tyName name loc = objectFromList
  [ "type" .= tyName
  , "name" .= genName name
  , "loc" .= displayLoc loc
  ]

genName :: Name -> Object
genName Name{..} = objectFromList $
  [ "name" .= localName resolvedName ] ++
  [ "src" .= m | QName m _ <- [sourceName] ]

genConstVal :: Typecheckable l => Type l t -> TypedConst l t -> Object
genConstVal ty (Literal x) =
  objectFromList [ "literal" .= genLiteral ty x ]
genConstVal _ (Identifier name _ _loc) =
  objectFromList [ "named_constant" .= genName name ]
genConstVal _ (WeirdEnumToInt _ name _ _loc) =
  objectFromList [ "named_constant_enumToInt" .= genName name ]

genLiteral :: Typecheckable l => Type l t -> t -> Object

-- Base Types
genLiteral ty@I8  n = simpleLiteral ty n
genLiteral ty@I16 n = simpleLiteral ty n
genLiteral ty@I32 n = simpleLiteral ty n
genLiteral ty@I64 n =
  -- We need to include the string representation because JSON does not support
  -- 64 bit integers
  simpleLiteral ty n <> objectFromList [ "string" .= show n ]
genLiteral ty@TFloat n =
  simpleLiteral ty n <>
  objectFromList [ "binary" .= toLazyText (floatHexFixed n) ]
genLiteral ty@TDouble n =
  simpleLiteral ty n <>
  objectFromList [ "binary" .= toLazyText (doubleHexFixed n) ]
genLiteral ty@TBool b = simpleLiteral ty b
genLiteral ty@TText s = simpleLiteral ty s
-- Serialized as a hexidecimal string
genLiteral ty@TBytes s = simpleLiteral ty $ toLazyText $ byteStringHex s

-- Collections
genLiteral (TSet u)       (Set xs)     = listLiteral "set" u xs
genLiteral (THashSet u)   (HashSet xs) = listLiteral "hash_set" u xs
genLiteral (TList u)      (List xs)    = listLiteral "list" u xs
genLiteral (TMap k v)     (Map xs)     = mapLiteral "map" k v xs
genLiteral (THashMap k v) (HashMap xs) = mapLiteral "hash_map" k v xs

-- Named Types
genLiteral TStruct{} (Some sval) = genStructVal sval
genLiteral TException{} (Some (EV sval)) = genStructVal sval
genLiteral TUnion{} (Some uval) = genUnionVal uval
genLiteral TEnum{} (EnumVal name _loc) = objectFromList
  [ "type"  .= ("enum" :: Text)
  , "value" .= genName name
  ]
genLiteral (TTypedef _ ty _loc) x = genLiteral ty x
genLiteral (TNewtype _ ty _loc) (New x) = objectFromList
  [ "type"  .= ("newtype" :: Text)
  , "value" .= genLiteral ty x
  ]
genLiteral st@(TSpecial ty) val = case backTranslateLiteral ty val of
  ThisLit u x -> objectFromList
    [ "type"  .= genType st
    , "value" .= genLiteral u x
    ]

simpleLiteral :: (Typecheckable l, ToJSON a) => Type l t -> a -> Object
simpleLiteral ty x = genType ty <> objectFromList [ "value" .= x ]

listLiteral :: Typecheckable l => Text -> Type l t -> [TypedConst l t] -> Object
listLiteral tyName ty xs = objectFromList
  [ "type"  .= tyName
  , "value" .= map (genConstVal ty) xs
  ]

mapLiteral
  :: Typecheckable l
  => Text
  -> Type l k
  -> Type l v
  -> [(TypedConst l k, TypedConst l v)]
  -> Object
mapLiteral tyName kt vt xs = objectFromList
  [ "type"  .= tyName
  , "value" .= map (genPair kt vt) xs
  ]

genPair
  :: Typecheckable l
  => Type l k
  -> Type l v
  -> (TypedConst l k, TypedConst l v)
  -> Value
genPair kt vt (k, v) = Object $ objectFromList
  [ "key" .= genConstVal kt k
  , "val" .= genConstVal vt v
  ]

genStructVal :: Typecheckable l => StructVal l s -> Object
genStructVal s = objectFromList
  [ "type"  .= ("struct" :: Text)
  , "value" .= genStructFields s
  ]

genStructFields :: Typecheckable l => StructVal l s -> [Object]
genStructFields Empty = []
genStructFields (ConsVal proxy ty c s) =
  genFieldVal proxy ty c : genStructFields s
genStructFields (ConsDefault proxy ty s) = objectFromList
  [ "field_name"  .= symbolVal proxy
  , "field_type"  .= genType ty
  , "field_value" .=
    HashMap.singleton ("default" :: Text) Null
  ] :
  genStructFields s
genStructFields (ConsJust proxy ty c s) =
  genFieldVal proxy ty c : genStructFields s
genStructFields (ConsNothing proxy s) = objectFromList
  [ "field_name"  .= symbolVal proxy
  , "field_value" .= Null
  ] :
  genStructFields s

genUnionVal :: Typecheckable l => UnionVal l s -> Object
genUnionVal (UnionVal proxy ty c _) = objectFromList
  -- This isn't technically a thrift type, but we'll use it anyway
  [ "type"  .= ("union" :: Text)
  , "value" .= genFieldVal proxy ty c
  ]

genFieldVal
  :: (Typecheckable l, KnownSymbol s)
  => Proxy s
  -> Type l t
  -> TypedConst l t
  -> Object
genFieldVal proxy ty c = objectFromList
  [ "field_name"  .= symbolVal proxy
  , "field_type"  .= genType ty
  , "field_value" .= genConstVal ty c
  ]

toLazyText :: Builder -> Lazy.Text
toLazyText = Lazy.decodeUtf8 . toLazyByteString

-- Options ---------------------------------------------------------------------

genOptions :: Options.Options l -> Object
genOptions Options.Options{..} = objectFromList
  [ "path" .= optsPath
  , "out_path" .= optsOutPath
  , "include_path" .= optsIncludePath
  , "recursive" .= optsRecursive
  , "genfiles" .= optsThriftMade
  ]