{-
  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.GenHaskell
  ( genHsCode
  , writeHsCode
  , writeModule, showModule
  , ThriftModule(..)
  , commonPragmas
  ) where

import Data.List
import Data.List.Extra
import Data.Text (Text)
import Language.Haskell.Exts hiding (parse, Decl, name, app)
import System.Directory
import System.FilePath
import Text.Printf
import qualified Data.Set as Set
import qualified Data.Text as Text
import qualified Language.Haskell.Exts.Syntax as HS

import Thrift.Compiler.GenClient
import Thrift.Compiler.GenConst
import Thrift.Compiler.GenEnum
import Thrift.Compiler.GenFunction
import Thrift.Compiler.GenService
import Thrift.Compiler.GenStruct
import Thrift.Compiler.GenTypedef
import Thrift.Compiler.GenUnion
import Thrift.Compiler.GenUtils

import Thrift.Compiler.Options
import Thrift.Compiler.Plugins.Haskell
import Thrift.Compiler.Types as Thrift hiding (noLoc)

data InstancesFile a = InstancesFile
  { ifPragmas :: [ModulePragma a]
  , ifImports :: [ImportDecl a]
  , ifDecls   :: [HS.Decl a]
  }

data ThriftModule = ThriftModule
  { tmPath :: FilePath
  , tmContents :: String
  , tmModuleName :: String
  }

writeHsCode :: Options Haskell -> Program Haskell Thrift.Loc -> IO [FilePath]
writeHsCode opts prog =
  -- Write the Generated Files
  mapM writeModule =<< genHsCode opts prog

genHsCode :: Options Haskell -> Program Haskell Thrift.Loc -> IO [ThriftModule]
genHsCode Options{..} prog@Program{..} = do
  let
    progHSPath = Text.unpack $ Text.replace "." "/" progHSName
    HsOpts{..} = optsLangSpecific
    dir = progOutPath </> hsoptsGenPrefix </> progHSPath
  relativeDir <- makeRelativeToCurrentDirectory dir

  let
    (typesModuleName, typesModuleBase) =
      genTypesModule prog hsoptsExtensions hsoptsExtraHasFields
    -- Get instances file if it exists
    typesModuleBody = case progInstances of
      Just (Module _ _ pragmas imports decls) ->
        showModuleWithInstances (relativeDir </> "Types.hs") typesModuleBase $
        InstancesFile pragmas imports decls
      _ -> showThriftModule typesModuleBase

  return $
    ThriftModule
      (dir </> "Types.hs")
      typesModuleBody
      typesModuleName :
    concat
      [ [ ThriftModule
          (dir </> Text.unpack serviceResolvedName </> "Client.hs")
          (showThriftModule clientModule)
          clientModuleName
        , ThriftModule
          (dir </> Text.unpack serviceResolvedName </> "Service.hs")
          (showThriftModule serviceModule)
          serviceModuleName
        ]
      | D_Service s@Service{..} <- progDecls
      , let
          (clientModuleName, clientModule) = genClientModule prog s
          (serviceModuleName, serviceModule) = genServiceModule prog s
      ]

genTypesModule
  :: Program Haskell Thrift.Loc -> [Text] -> Bool -> (String, Module ())
genTypesModule prog@Program{..} extensions extraHasFields =
  genModule prog "Types" pragmas
    (concat exports)
    (map genImportModule $ Set.toList imports)
    (concat decls)
  where
    pragmas = commonPragmas (options progEnv) ++
              map (LanguagePragma () . (:[]) . textToName)
                (("GeneralizedNewtypeDeriving" : extensions) ++
                (if extraHasFields then hasFieldsExtensions else []))

    hasFieldsExtensions =
      [ "DataKinds"
      , "FlexibleInstances"
      , "MultiParamTypeClasses"
      ]

    (decls, imports, exports) = foldr genDecl ([], baseImports, []) progDecls
    baseImports = Set.fromList $
      [ QImport "Thrift.CodegenTypesOnly" "Thrift"
      , QImport "Prelude" "Prelude"
      ] ++
      map importFromInclude progIncludes

    genDecl decl (ds, is, es) = (d : ds, Set.union i is, e : es)
      where
        (d, i, e) = case decl of
          D_Struct s ->
            ( genStructDecl extraHasFields s
            , genStructImports s
            , [structExport s]
            )
          D_Union u ->
            (genUnionDecl u, genUnionImports u, [unionExport u])
          D_Typedef t ->
            (genTypedefDecl t True, genTypedefImports t, [tdefExport t])
          D_Enum en -> (genEnumDecl en, genEnumImports, enumExport en)
          D_Const c -> (genConstDecl c, genConstImports c, [constExport c])
          -- Services are not included in this module
          D_Service{} -> mempty
          -- Interactions are not included in this module
          D_Interaction{} -> mempty

        tdefExport Typedef{..} = case tdTag of
          IsNewtype -> newtypeExport tdResolvedName
          IsTypedef -> EAbs () (NoNamespace ()) $ unqualSym tdResolvedName
        newtypeExport name =
          EThingWith () (NoWildcard ()) (unqualSym name)
          [ ConName () (textToName name)
          , VarName () (textToName $ "un" <> name)
          ]
        structExport :: HS Struct -> ExportSpec ()
        structExport Struct{..} =
          EThingWith () (NoWildcard ()) (unqualSym structResolvedName) $
          ConName () (textToName structResolvedName) :
          [ VarName () (textToName fieldResolvedName)
          | Field{..} <- structMembers ]
        unionExport :: HS Union -> ExportSpec ()
        unionExport Union{..} =
          EThingWith () (NoWildcard ()) (unqualSym unionResolvedName) $
          (case unionHasEmpty of
             HasEmpty -> (ConName () (textToName unionEmptyName) :)
             NonEmpty -> id)
          [ ConName () (textToName altResolvedName) | UnionAlt{..} <- unionAlts ]
        enumExport :: HS Thrift.Enum -> [ExportSpec ()]
        enumExport Enum{..} = case enumFlavour of
          PseudoEnum{} ->
            newtypeExport enumResolvedName :
            [ EVar () $ unqualSym evResolvedName
            | EnumValue{..} <- enumConstants
            ]
          SumTypeEnum{..} ->
            [ EThingWith () (NoWildcard ()) (unqualSym enumResolvedName) $
              [ ConName () (textToName evResolvedName)
              | EnumValue{..} <- enumConstants
              ] ++
              [ ConName () (textToName $ enumResolvedName <> "__UNKNOWN")
              | not enumNoUnknown
              ]
            ]
        constExport = EVar () . unqualSym . constResolvedName

genClientModule :: Program Haskell Thrift.Loc -> HS Service -> (String, Module ())
genClientModule prog@Program{..} service@Service{..} =
  genModule prog (serviceResolvedName <> ".Client") pragmas exports imports decls
  where
    theseFunctions = getServiceFunctions service
    pragmas = commonPragmas opts ++
              map (LanguagePragma () . (:[]) . textToName)
              (["FlexibleContexts", "TypeFamilies", "TypeOperators"] ++
               hsoptsExtensions optsLangSpecific)
    exports = EAbs () (NoNamespace ()) (unqualSym serviceResolvedName) :
              concatMap
              (\Function{..} ->
                map (EVar () . unqualSym . ($ funResolvedName)) $
                [ id, (<> "IO"), ("send_" <>), ("_build_" <>) ] ++
                (if funIsOneWay then [] else [ ("recv_" <>), ("_parse_"<>) ]))
              theseFunctions
    imports = map genImportModule $ Set.toList $ Set.unions $
              Set.singleton (QImport "Thrift.Codegen" "Thrift") :
              Set.singleton (TypesImport progHSName) :
              Set.fromList (map importFromInclude progIncludes) :
              genClientImports progHSName service :
              map genFunctionImports theseFunctions
    decls   = genClientDecls service ++
              concatMap (genFunctionDecls service) theseFunctions
    opts@Options{..} = options progEnv

genServiceModule :: Program Haskell Thrift.Loc -> HS Service -> (String, Module ())
genServiceModule prog@Program{..} service@Service{..} =
  genModule prog (serviceResolvedName <> ".Service") pragmas exports imports decls
  where
    pragmas = commonPragmas (options progEnv) ++
              [ LanguagePragma () [textToName "GADTs"]
              ]
    exports = genServiceExports service
    imports = map genImportModule $ Set.toList $ Set.unions $
              Set.singleton (QImport "Thrift.Codegen" "Thrift") :
              Set.singleton (QImport (progHSName <> ".Types") "Types") :
              Set.fromList (map importFromInclude progIncludes) :
              [genServiceImports progHSName service]
    decls   = genServiceDecls service

commonPragmas :: Options Haskell -> [ModulePragma ()]
commonPragmas Options{ optsLangSpecific = HsOpts{..} } =
  (if hsoptsDupNames
   then (LanguagePragma () [textToName "DuplicateRecordFields"] :)
   else id)
  [ LanguagePragma () [textToName "OverloadedStrings"]
  , LanguagePragma () [textToName "BangPatterns"]
  , OptionsPragma () (Just GHC) "-fno-warn-unused-imports"
  , OptionsPragma () (Just GHC) "-fno-warn-overlapping-patterns"
  , OptionsPragma () (Just GHC) "-fno-warn-incomplete-patterns"
  , OptionsPragma () (Just GHC) "-fno-warn-incomplete-uni-patterns"
  , OptionsPragma () (Just GHC) "-fno-warn-incomplete-record-updates"
  ]

genModule
  :: Program Haskell Thrift.Loc
  -> Text
  -> [ModulePragma ()]
  -> [ExportSpec ()]
  -> [ImportDecl ()]
  -> [HS.Decl ()]
  -> (String, Module ())
genModule Program{..} moduleName pragmas exports imports decls =
  ( fullName
  , Module ()
      (Just
       (ModuleHead () (ModuleName () fullName)
        Nothing
        (Just $ ExportSpecList () exports)))
      pragmas
      imports
      decls
  )
  where
    fullName = Text.unpack $ progHSName <> "." <> moduleName

writeModule :: ThriftModule -> IO FilePath
writeModule ThriftModule{..} = do
  let (dir, _fname) = breakOnEnd "/" tmPath
  createDirectoryIfMissing True dir
  writeFile tmPath tmContents
  return tmPath

showThriftModule :: Module () -> String
showThriftModule = showModule autogenComment

showModule :: String -> Module () -> String
showModule header = (header ++) . prettyPrintWithMode baseMode . (noLoc <$)

showModuleWithInstances
  :: FilePath
  -> Module ()
  -> InstancesFile SrcSpanInfo
  -> String
showModuleWithInstances path baseModule InstancesFile{..} = fileBody
  where
    -- Put everything together
    fileBody = foldl1 catWithPragma
               [ fileHeader, fileImports, fileDecls ]
    catWithPragma a b = a ++ linePragma (length (lines a)) ++ b
    linePragma n = printf "{-# LINE %d \"%s\" #-}\n" (n + 2) path

    fileHeader = autogenComment ++ extraPragmas
    fileImports = unlines imports ++ extraImports
    fileDecls   = unlines decls ++ extraDecls

    (imports, decls) = splitAt lastImport baseFile
    baseFile = lines $ prettyPrintWithMode baseMode $ noLoc <$ baseModule

    lenBase = length baseFile
    lastImport = maybe lenBase (lenBase -) $
                 findIndex (isPrefixOf "import") $ reverse baseFile
    -- Pretty Print the extras
    extraPragmas = annot ifPragmas
    extraImports = annot ifImports
    extraDecls   = annot ifDecls
    annot xs = unlines $ map pp xs

    pp :: (HS.Annotated s, Pretty (s SrcSpanInfo))
       => s SrcSpanInfo -> String
    pp d =
      printf "{-# LINE %d \"%s\" #-}\n" srcSpanStartLine srcSpanFilename
       ++ prettyPrintWithMode baseMode d
      where
        SrcSpanInfo SrcSpan{..} _ = ann d

autogenComment :: String
autogenComment = unlines
  [ "-----------------------------------------------------------------"
  , "-- Autogenerated by Thrift"
  , "--"
  , "-- DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING"
  , "--  @" ++ "generated"
  , "-----------------------------------------------------------------"
  ]

baseMode :: PPHsMode
baseMode = defaultMode
  { classIndent   = 2
  , doIndent      = 3
  , multiIfIndent = 3
  , caseIndent    = 2
  , letIndent     = 2
  , whereIndent   = 2
  , onsideIndent  = 2
  , spacing       = True
  , layout        = PPOffsideRule
  }