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

-- | Generate "Glean.Schema"
module Glean.Schema.Gen.OCaml
  ( genSchemaOCaml
  ) where

import qualified Data.HashMap.Strict as HashMap
import Data.Text (Text)
import Data.List (partition)
import qualified Data.Text as Text
import Control.Monad.State

import Glean.Schema.Gen.Utils
    ( NameSpaces,
      NamePolicy(predNames, typeNames),
      sortDeclsByNamespace,
      newline,
      cap1,
      mkNamePolicy,
      addNamespaceDependencies,
      Oncall, buckOncallAnnotation )
import Glean.Angle.Types
    ( PredicateRef,
      TypeRef,
      Version,
      PredicateDef_(predicateDefRef, predicateDefKeyType,
                    predicateDefValueType),
      TypeDef_(typeDefRef, typeDefType),
      FieldDef_(fieldDefType, fieldDefName),
      Type_(..) )
import Glean.Schema.Types
    ( ResolvedPredicateDef, ResolvedTypeDef, ResolvedType )
import Data.Char (isUpper, toLower)

generated :: Text
generated = Text.concat
  [ "regenerate: buck2 run fbcode//glean/schema/gen:gen-schema  "
  , "-- --ocaml fbcode/hphp/hack/src/typing/write_symbol_info/schema "
  , "--dir DEST_DIR"
  ]

dune :: [Text] -> Text
dune modules = Text.unlines
  [ "; \x40generated"
  , "; " <> generated
  ,  "(library"
  , " (name glean_schema)"
  , " (wrapped false)"
  , "(modules"
  , " " <> Text.intercalate "\n " modules <> ")"
  , "(libraries"
  , " utils_core"
  , " core"
  , " base64)"
  , "(preprocess"
  , " (pps ppx_deriving.std)))"
  ]

-- disable warning for "unused module imports" and "unused rec".
-- Not all generated files use it.
ocamlHeaderGen :: [Text] -> Text
ocamlHeaderGen imports = Text.unlines $
  [ "(*"
  , " * 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."
  , " *)"
  , ""
  , "(* \x40generated"
  , "   " <> generated <> " *)"
  , ""
  , "[@@@warning \"-33-39\"]"
  ] ++ (("open " <>) <$> imports) ++ [ "" ]

ocamlHeader :: Text
ocamlHeader =  ocamlHeaderGen ["Hh_json", "Core"]

ocamlHeader' :: Text
ocamlHeader' =  ocamlHeaderGen ["Hh_json"]

-- TODO use "Core.Map"
factIdMl :: Text
factIdMl = ocamlHeader' <> Text.unlines
  [ ""
  , "type t = int [@@deriving ord]"
  , ""
  , "let next ="
  , "  let x = ref 1 in"
  , "  fun () ->"
  , "    let r = !x in"
  , "    x := !x + 1;"
  , "    r"
  , ""
  , "let to_json_number i = JSON_Number (string_of_int i)"
  , ""
  , "module Map = Map.Make (struct"
  , "  type t = int"
 , ""
 , "  let compare = Int.compare"
 , "end)"
 , ""
  ]

utilMl :: Text
utilMl = ocamlHeader <> Text.unlines
  [ ""
  , "let key v = JSON_Object [(\"key\", v)]"
  , ""
  , "let id fact_id = JSON_Object [(\"id\", Fact_id.to_json_number fact_id)]"
  , ""
  ]

targetsHeader :: Maybe Oncall -> Text
targetsHeader oncall = Text.unlines
  [ "# \x40generated"
  , "# " <> generated
  , "load(\"@fbcode_macros//build_defs:ocaml_library.bzl\", \"ocaml_library\")"
  , buckOncallAnnotation oncall
  , ""
  ]

targetsLib :: Text -> [Text] -> [Text] -> Text
targetsLib name srcs deps = Text.unlines
  [ "ocaml_library("
  , "  name = \"" <> name <> "\","
  , "  srcs = ["
  , Text.intercalate "\n" (buckItem <$> srcs)
  , "  ],"
  , "  deps = ["
  , Text.intercalate "\n" (buckItem <$> deps)
  , "  ],"
  , "  preprocess = ["
  , "    \"ppx_deriving.show\","
  , "    \"ppx_deriving.eq\","
  , "    \"ppx_deriving.ord\","
  , "    \"ppx_deriving.enum\""
  , "  ],"
  , "  external_deps = ["
  , "    (\"supercaml\", None, \"base64\"),"
  , "    (\"supercaml\", None, \"core\")"
  , "  ]"
  , ")"
  , ""
  ]
  where buckItem x = "    \"" <> x <> "\","

moduleDecl :: Int -> Text
moduleDecl 0 = "module rec "
moduleDecl _ = "and "

type GenTypeName = Text
type GenFunName = Text
type GenVars = Text
type GenCode = Text
type GenType = Text

-- all the text generated from an angle type/pred
-- we pick pieces of this to generate the OCaml modules
-- signature and implementation
type GenEntity = (GenTypeName, GenFunName, GenType, GenVars, GenCode)

displayList :: Text -> Text -> Text -> [Text] -> Text
displayList first second sep texts =
    Text.intercalate sep $ zipWith (<>) prefixes texts
  where
     prefixes = first : repeat second

typeDecl :: GenEntity -> Text
typeDecl (name, _fun_name, type_, _var, _code ) =
  mconcat [name, " = ", type_, "\n  [@@deriving ord]"]

toJsonDecl :: GenEntity -> Text
toJsonDecl (name, fun_name, _type, _var, _code ) =
  mconcat ["  val ", fun_name, ": ", name, " -> json"]

toJsonImpl :: GenEntity -> Text
toJsonImpl (_name, fun_name, _type, var, code ) =
  mconcat [fun_name, " ", var, " = ", code]

genModuleType :: Int -> Text -> [GenEntity] -> Text
genModuleType decl name types = Text.unlines
    [ moduleDecl decl <> name <> ": sig"
    , displayList "  type " "  and " "\n" (typeDecl <$> types)
    , ""
    , "  val to_json : t -> json"
    , "end = struct"
    , displayList "  type " "  and " "\n" (typeDecl <$> types)
    , ""
    , displayList "  let rec " "  and " "\n" (toJsonImpl <$> types)
    , "end"
    ]

genModulePred :: Int -> Text -> [GenEntity] -> Text
genModulePred decl name types = Text.unlines
  [ moduleDecl decl <> name <> ": sig"
  , "  type t ="
  , "    | Id of Fact_id.t"
  , "    | Key of key"
  , "  [@@deriving ord]"
  , ""
  , displayList "  and " "  and " "\n" (typeDecl <$> types)
  , ""
  , "  val to_json: t -> json"
  , ""
  , Text.intercalate "\n" (toJsonDecl <$> types)
  , ""
  , "end = struct"
  , "  type t ="
  , "    | Id of Fact_id.t"
  , "    | Key of key"
  , "  [@@deriving ord]"
  , ""
  , displayList "  and " "  and " "\n" (typeDecl <$> types)
  , ""
  , "  let rec to_json = function"
  , "    | Id f -> Util.id f"
  , "    | Key t -> Util.key (to_json_key t)"
  , ""
  , displayList "  and " "  and " "\n" (toJsonImpl <$> types)
  , "end"
  ]

-- Extra state used in generation function to store generated
-- code/type for anonymous datatypes we can't translate directly
-- in OCaml
type GenAnonTypes = [(GenTypeName, GenType)]
type GenAnonTypesCode = [(GenTypeName, GenVars, GenCode)]

data FieldKind = Sum | Record

-- | Generate the type representing a predicate or type
genOCamlType
  :: NameSpaces -> NamePolicy -> ResolvedType -> State GenAnonTypes GenType
genOCamlType ns namePolicy t = case t of
  ByteTy -> return "char"
  NatTy -> return "int"
  StringTy -> return "string"
  ArrayTy ByteTy -> return "string list"
  ArrayTy ty -> do
    t <- genOCamlType ns namePolicy ty
    return $ t <> " list"
  RecordTy [] -> return "unit"
  RecordTy fields -> do
    f <- mapM (genField Record) fields
    return $ "{\n" <> Text.unlines f <> "  }"
  SumTy fields -> do
    f <- mapM (genField Sum) fields
    return $ "\n" <> Text.intercalate "\n" f
  SetTy ty -> do
    t <- genOCamlType ns namePolicy ty
    return $ t <> " list" -- Suboptimal but sets in OCaml is such a pain
  PredicateTy _ pred -> return $ predToModule ns pred namePolicy <> ".t"
  NamedTy _ tref -> return $ typeToModule ns tref namePolicy <> ".t"
  MaybeTy ty -> do
    t <- genOCamlType ns namePolicy ty
    return $ t <> " option"
  EnumeratedTy names ->
    return $ "\n" <> Text.unlines (("    | " <>) . nameToConstructor <$> names)
  BooleanTy -> return "bool"
  TyVar{} -> error "genOCamlType: TyVar"
  HasTy{} -> error "genOCamlType: HasTy"
  HasKey{} -> error "genOCamlType: HasKey"
  ElementsOf{} -> error "genOCamlType: ElementsOf"
  where
    genField fieldKind field = do
      ty <- genOCamlTypeFromField field ns namePolicy (fieldDefType field)
      case fieldKind of
        Record -> return $ "    " <> fieldToVar field <> ": " <> ty <> ";"
        Sum -> return $ "     | " <> fieldToConstructor field <> " of " <> ty

genOCamlTypeFromField
  :: FieldDef_ s PredicateRef TypeRef
  -> NameSpaces
  -> NamePolicy
  -> ResolvedType
  -> State GenAnonTypes GenType
genOCamlTypeFromField field ns namePolicy t = do
  let typeName = fieldToTypeName field
  type_ <- genOCamlType ns namePolicy t
  st <- get
  case t of
    RecordTy (_ : _) -> do
      put $ (typeName, type_) : st
      return typeName
    SumTy _ -> do
      put $ (typeName, type_) : st
      return typeName
    EnumeratedTy _ -> do
      put $ (typeName, type_) : st
      return typeName
    _ -> return type_

genOCamlToJson
  :: GenVars -> NameSpaces -> NamePolicy -> ResolvedType
  -> State GenAnonTypesCode (GenVars, GenCode)
genOCamlToJson var ns namePolicy t = case t of
  ByteTy ->
    return (var, "JSON_Number (string_of_int (int_of_char " <> var <> "))")
  NatTy -> return (var, "JSON_Number (string_of_int " <> var <> ")")
  StringTy -> return (var, "JSON_String " <> var)
  ArrayTy ByteTy ->
    return (var, "JSON_String (List.map ~f:Base64.encode_string " <> var
    <> "|> String.concat ~sep:\"\")")
  ArrayTy ty -> do
    (_, code) <- genOCamlToJson "x" ns namePolicy ty
    return (var, "JSON_Array (List.map ~f:(fun x -> " <> code  <> ") "
       <> var <> ")")
  RecordTy [] ->
    return ("_", "JSON_Object []")
  RecordTy fields -> do
    let fieldNames = fieldToVar <$> fields
        vars = "{" <> Text.intercalate "; " fieldNames <> "}"
        fieldToVarPair field = do
          let camlVar = fieldToVar field
              jsonKey = fieldToJSONKey field
              type_ = fieldDefType field
          (_, camlType) <-
            genOCamlToJsonFromField field camlVar ns namePolicy type_
          let isOpt = case type_ of
                MaybeTy _ -> True
                _ -> False
              camlPair = "(\"" <> jsonKey <> "\", " <> camlType <> ")"
          return (camlVar, camlPair, isOpt)
    fields_ <- mapM fieldToVarPair fields
    let (optFields, reqFields) = partition (\(_, _, isOpt) -> isOpt) fields_
        reqFieldsDefs = Text.unlines (
          let toRegString  (_, pair, _) = "      " <> pair <> ";" in
          [ "    let fields = [" ] <>
          (toRegString <$> reqFields) <>
          [ "    ] in" ])
        optFieldsDefs =
          let toOptFieldDef (var, pair, _) = Text.unlines
                [ "    let fields =",
                  "      match " <> var <> " with",
                  "      | None -> fields",
                  "      | Some " <> var <> " -> " <> pair <>  " :: fields in" ]
          in
          Text.concat $ toOptFieldDef <$> optFields
    return (vars, "\n" <> reqFieldsDefs <> optFieldsDefs
     <> "    JSON_Object fields\n")
  SumTy fields -> do
    let typeSumField field = do
          let var = fieldToVar field
          (_, genType) <-
            genOCamlToJsonFromField field var ns namePolicy (fieldDefType field)
          let key = fieldToJSONKey field
          let constr = fieldToConstructor field
          return $ Text.concat [ "     | ", constr, " ", var,
                        " -> JSON_Object [(\"", key, "\", ", genType, ")]"]
    fields <- mapM typeSumField fields
    return ("", "function\n" <> Text.unlines fields)
  SetTy ty -> do
    (_, code) <- genOCamlToJson "x" ns namePolicy ty
    return (var, "JSON_Array (List.map ~f:(fun x -> " <> code  <> ") "
       <> var <> ")")
  PredicateTy _ pred ->
    let moduleName = predToModule ns pred namePolicy in
        return (var, moduleName <> ".to_json " <> var)
  NamedTy _ tref ->
    let moduleName = typeToModule ns tref namePolicy in
        return (var, moduleName <> ".to_json " <> var)
  MaybeTy ty -> do
    (_, genTy) <- genOCamlToJson "x" ns namePolicy ty
    return (var, "Option.map ~f:(fun x -> " <> genTy <> ") " <> var)
  EnumeratedTy names ->
    let enumNames = zipWith (\idx val -> (idx, val)) [(0::Int)..] names in
    let enumCase (num, name) =
          Text.concat
            ["     | ", nameToConstructor name,
             " -> JSON_Number (string_of_int ", Text.pack $ show num, ")"] in
    let handleEnumCases = enumCase <$> enumNames in
        return ("", "function\n" <> Text.unlines handleEnumCases)
  BooleanTy -> return (var, "JSON_Bool " <> var)
  TyVar{} -> error "genOCamlToJson: TyVar"
  HasTy{} -> error "genOCamlToJson: HasTy"
  HasKey{} -> error "genOCamlToJson: HasKey"
  ElementsOf{} -> error "genOCamlToJson: ElementsOf"

genOCamlToJsonFromField
  :: FieldDef_ s PredicateRef TypeRef -> GenVars -> NameSpaces -> NamePolicy
  -> ResolvedType -> State GenAnonTypesCode (GenVars, GenCode)
genOCamlToJsonFromField field var ns namePolicy t = do
  let typeName = fieldToTypeName field
      res = (var, typeName <> "_to_json " <> var)
  st <- get
  (var0, code) <- genOCamlToJson var ns namePolicy t
  case t of
   RecordTy [] -> do
     return (var, "(ignore " <>  var <> "; JSON_Object [])")
   RecordTy _ -> do
     put $ (typeName, var0, code) : st
     return res
   SumTy _ -> do
     put $ (typeName, var0, code) : st
     return res
   EnumeratedTy _ -> do
     put $ (typeName, var0, code) : st
     return res
   MaybeTy ty -> genOCamlToJsonFromField field var ns namePolicy ty
   _ -> return (var0, code)


refToModule :: NameSpaces -> Maybe (NameSpaces, Text) -> Text
refToModule curNs ref =
  case ref of
    Just (ns :: NameSpaces, x) ->
      let toplevelMod = if ns == curNs then
             ""
           else
             cap1 $ Text.intercalate "_" ns <> "." in
       toplevelMod <> cap1 x
    _ -> error $ "OCamlName: " ++ show ref

typeToModule :: NameSpaces -> TypeRef -> NamePolicy -> Text
typeToModule curNs ref namePolicy =
  let map = typeNames namePolicy in
  refToModule curNs (HashMap.lookup ref map)

predToModule :: NameSpaces -> PredicateRef -> NamePolicy -> Text
predToModule curNs ref namePolicy =
  let map = predNames namePolicy in
  refToModule curNs (HashMap.lookup ref map)

-- used for anonymous types, we reuse the field name
-- e.g. angle: type t = { a : { ...} }
--      ocaml: type a = ...
-- TODO: handle naming conflict, use path of fields
fieldToTypeName :: FieldDef_ s PredicateRef TypeRef -> GenTypeName
fieldToTypeName = fieldDefName

nameToConstructor :: Text -> Text
nameToConstructor = cap1

fieldToJSONKey :: FieldDef_ s PredicateRef TypeRef -> Text
fieldToJSONKey field = fieldDefName field

fieldToConstructor :: FieldDef_ s PredicateRef TypeRef -> Text
fieldToConstructor field = cap1 (fieldDefName field)

fieldToVar :: FieldDef_ s PredicateRef TypeRef -> Text
fieldToVar field =
  from_ (camelToUnderscore (fieldDefName field))
  where
    from_ name
      | Text.unpack name `elem` keywords = name <> "_"
      | otherwise = name
      where
        keywords = ["begin", "end", "method", "type", "module", "to", "val",
                    "object"]

    camelToUnderscore = Text.pack . go . Text.unpack
      where
        go [] = []
        go (x:xs)
          | isUpper x = '_' : toLower x : go xs
          | otherwise = x : go xs

nsToTarget :: NameSpaces -> Text
nsToTarget n = Text.intercalate "_" n

nsToFileName :: NameSpaces -> Text
nsToFileName n = nsToTarget n <> ".ml"

genSchema
 :: NameSpaces
  -> [ResolvedPredicateDef]
  -> [ResolvedTypeDef]
  -> NamePolicy
  -> Text
genSchema namespaces preds types namePolicy =
      Text.intercalate newline [ocamlHeader, predModules, typeModules]
    where
      predModules =
        Text.unlines [genModuleP num namespaces namePolicy pred |
          (num, pred) <- zipWith (\idx val -> (idx, val)) [0..]  preds]
      typeModules =
        let startValue = length preds in
        Text.unlines [genModuleT num namespaces namePolicy type_ |
          (num, type_) <- zipWith (\idx val -> (idx, val)) [startValue..] types]
      genModuleP (num :: Int) ns namePolicy pred =
        let ref = predicateDefRef pred
            key = predicateDefKeyType pred
            value = predicateDefValueType pred
            name = predToModule ns ref namePolicy
        in case value of
              RecordTy [] ->
                genModulePred num name
                  (typeToGenEntity "key" "to_json_key" ns namePolicy key)
              _ ->
                genModulePred num name
                  (typeToGenEntity "key" "to_json_key" ns namePolicy key ++
                   typeToGenEntity "value" "to_json_value" ns namePolicy value)
      genModuleT (num :: Int) ns namePolicy type_ =
          genModuleType num name
            (typeToGenEntity "t" "to_json" ns namePolicy key)
        where
          ref = typeDefRef type_
          key = typeDefType type_
          name = typeToModule ns ref namePolicy

      typeToGenEntity
        :: GenTypeName -> GenFunName -> NameSpaces -> NamePolicy -> ResolvedType
        -> [GenEntity]
      typeToGenEntity typeName funName ns namePolicy x  =
        let ((var, code), otherCode) =
              runState (genOCamlToJson "x" ns namePolicy x) []
            (kTy, otherTypes) =
              runState (genOCamlType ns namePolicy x) []
            others = [(name, name <> "_to_json", kTy, var, code) |
              (name, kTy) <- otherTypes,
              (name', var, code) <- otherCode, name == name']
        in
          (typeName, funName, kTy, var, code) : others

genTargets :: [(NameSpaces, [NameSpaces])] -> Maybe Oncall -> Text
genTargets deps oncall =
  targetsHeader oncall
  <> targetsLib "common" ["util.ml", "fact_id.ml"] [jsonDep]
  <> Text.concat (genTarget <$> deps)
  where
  jsonDep = "//hphp/hack/src/utils/hh_json:hh_json"
  genTarget (ns, imports) =
    let
      extra = (":" <>) . nsToTarget <$> imports
      name = nsToTarget ns
      file = nsToFileName ns
      targetDeps = [ ":common", jsonDep ] ++ extra in
    targetsLib name [file] targetDeps

-- entry point
-- generate all files, possibly restricted to schemas in [NameSpaces]
genSchemaOCaml
  :: Maybe [NameSpaces]
  -> Version
  -> [ResolvedPredicateDef]
  -> [ResolvedTypeDef]
  -> Maybe Oncall
  -> [(FilePath,Text)]
genSchemaOCaml onlySchemas _version preddefs typedefs oncall =
  ( "dune", dune modules) :
  ( "TARGETS", genTargets deps oncall) :
  ( "fact_id.ml", factIdMl) :
  ( "util.ml", utilMl) :
  [ (Text.unpack (nsToFileName namespaces),
    genSchema namespaces preds types namePolicy)
    | (namespaces, (_, preds, types)) <- schemas
  ]
  where
    filter = case onlySchemas of
      Nothing -> \_ _ -> True
      Just schemas -> (\x _ -> elem x schemas)
    declsPerNamespace = HashMap.filterWithKey filter (
      addNamespaceDependencies $ sortDeclsByNamespace preddefs typedefs)
    schemas = HashMap.toList declsPerNamespace
    deps = (\(x, (y, _, _)) -> (x, y)) <$> schemas
    namePolicy = mkNamePolicy preddefs typedefs
    modules' = nsToTarget . fst <$> deps
    modules = ["util", "fact_id"] ++ modules'