{-# LANGUAGE LambdaCase #-}

-- |
-- Description: the core of the 'Smuggler2.Plugin'
module Smuggler2.Plugin
  ( plugin,
  )
where

import Avail ( AvailInfo, Avails )
import Control.Monad ( unless )
import Data.Bool ( bool )
import Data.Maybe ( fromMaybe, isNothing )
import Data.Version ( showVersion )
import DynFlags
    ( DynFlags(dumpDir), HasDynFlags(getDynFlags), xopt )
import ErrUtils ( compilationProgressMsg, fatalErrorMsg )
import GHC
    ( GenLocated(L),
      GhcPs,
      HsModule(hsmodExports, hsmodImports),
      ImportDecl(ideclHiding, ideclImplicit),
      LIE,
      LImportDecl,
      Located,
      ms_location,
      ml_hs_file,
      ModSummary(ms_hspp_buf, ms_mod),
      Module(moduleName),
      ParsedSource,
      moduleNameString,
      unLoc )
import GHC.LanguageExtensions ( Extension(Cpp) )
import GHC.IO.Encoding ( setLocaleEncoding, utf8 )
import IOEnv ( MonadIO(liftIO), readMutVar )
import Language.Haskell.GHC.ExactPrint
    ( Anns,
      TransformT,
      addTrailingCommaT,
      exactPrint,
      graftT,
      runTransform,
      setEntryDPT )
import Language.Haskell.GHC.ExactPrint.Types ( DeltaPos(DP) )
import Outputable
    ( Outputable(ppr), neverQualify, printForUser, text, vcat )
import Paths_smuggler2 ( version )
import Plugins
    ( CommandLineOption,
      Plugin(pluginRecompile, typeCheckResultAction),
      defaultPlugin,
      purePlugin )
import RnNames ( ImportDeclUsage, findImportUsage )
import Smuggler2.Anns ( mkLoc, mkParenT )
import Smuggler2.Imports ( getMinimalImports )
import Smuggler2.Exports ( mkExportAnnT )
import Smuggler2.Options
    ( ExportAction(AddExplicitExports, NoExportProcessing,
                   ReplaceExports),
      ImportAction(MinimiseImports, NoImportProcessing),
      Options(exportAction, importAction, newExtension),
      parseCommandLineOptions )
import Smuggler2.Parser ( runParser )
import StringBuffer ( StringBuffer(StringBuffer), lexemeToString )
import System.Directory ( removeFile )
import System.FilePath ( (-<.>), (</>), isExtensionOf, takeExtension )
import System.IO ( IOMode(WriteMode), withFile )
import TcRnExports ( exports_from_avail )
import TcRnTypes
    ( TcGblEnv(tcg_rdr_env, tcg_imports, tcg_mod, tcg_exports,
               tcg_rn_imports, tcg_used_gres, tcg_rn_exports),
      TcM,
      RnM )

-- | 'Plugin' interface to GHC
plugin :: Plugin
plugin =
  defaultPlugin
    { typeCheckResultAction = smugglerPlugin,
      pluginRecompile = purePlugin -- Don't force recompilation.  [Is this the right approach?]
    }

-- | The plugin itself
smugglerPlugin :: [CommandLineOption] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
smugglerPlugin clopts modSummary tcEnv
  -- short circuit, if nothing to do
  | (importAction options == NoImportProcessing)
      && (exportAction options == NoExportProcessing) =
    return tcEnv
  | otherwise = do
    -- Get the imports and their usage
    let imports = tcg_rn_imports tcEnv
    uses <- readMutVar $ tcg_used_gres tcEnv
    let usage = findImportUsage imports uses

    -- This ensures that the source file is not touched if there are no unused
    -- imports, or exports already exist and we are not replacing them
    let noUnusedImports = all (\(_decl, used, unused) -> not (null used) && null unused) usage

    let hasExplicitExports = case tcg_rn_exports tcEnv of
          Nothing -> False -- There is not even a module header
          (Just []) -> False
          (Just _) -> True
    -- ... so short circuit if:
    -- - we are skipping import processing or there are no unused imports, and
    -- - we are skipping export processing or there are explict exports and we are not replacing them
    -- (This is not a complete check; ideally, that the new imp/exports are
    -- different from the existing ones, etc)
    if (importAction options == NoImportProcessing || noUnusedImports)
      && ( exportAction options == NoExportProcessing
             || (hasExplicitExports && exportAction options /= ReplaceExports)
         )
      then return tcEnv
      else do
        dflags <- getDynFlags
        liftIO $ compilationProgressMsg dflags ("smuggler2 " ++ showVersion version)

        -- Dump GHC's view of what the minimal imports are for the current
        -- module, so that they can be annotated when parsed back in
        -- This is needed because there too much information loss between
        -- the parsed and renamed AST to use the latter for reconstituting the
        -- source.  An alternative would be to  "index" each name location with
        -- a SrcSpan to allow the name matchup, and to make the 'ParsedSource' a
        -- 100% representation of the original source (modulo tabs, trailing
        -- whitespace per line).
        let minImpFilePath = mkMinimalImportsPath dflags (ms_mod modSummary)
        printMinimalImports' dflags minImpFilePath usage

        -- Run smuggling only for its side effects
        tcEnv <$ smuggling dflags minImpFilePath
  where
    -- Does all the work
    smuggling :: DynFlags -> FilePath -> RnM ()
    smuggling dflags minImpFilePath = do
      -- The preprocessed source
      let modulePath = case ml_hs_file $ ms_location modSummary of
            Nothing -> error "smuggler2: missing source file location"
            Just loc -> loc

      -- Read files as UTF-8 strings (GHC accepts only ASCII or UTF-8)
      liftIO $ setLocaleEncoding utf8

      -- Get the pre-processed module source code
      let modFileContents = case ms_hspp_buf modSummary of
            -- Not clear under what circumstances this could happen
            Nothing -> error $ "smuggler2: missing source file: " ++ modulePath
            Just contents -> strBufToStr contents

      -- Parse the whole module
      runParser dflags modulePath modFileContents >>= \case
        Left () -> return () -- do nothing if file is invalid Haskell
        Right (annsHsMod, astHsMod) -> do
          -- Read the dumped file of minimal imports
          minImpFileContents <- liftIO $ readFile minImpFilePath

          -- Parse the minimal imports file -- gets the annnotations too
          runParser dflags minImpFilePath minImpFileContents >>= \case
            Left () ->
              liftIO $
                fatalErrorMsg dflags (text $ "smuggler: failed to parse minimal imports from " ++ minImpFilePath)
            Right (annsImpMod, L _ impMod) -> do
              -- The actual minimal imports themselves, as generated by GHC
              let minImports = hsmodImports impMod

              -- What is exported by the module
              exports <-
                if exportAction options == ReplaceExports
                  then exportable
                  else return $ tcg_exports tcEnv -- what is currently exported

              -- Bringing it all together: generate a new ast and annotations for it
              let (astHsMod', (annsHsMod', _locIndex), _log) =
                    runTransform annsHsMod $
                      replaceImports annsImpMod minImports astHsMod
                        >>= addExplicitExports exports

              -- Generate new file extension
              -- Was the CPP used?
              let usedCpp = bool "" "-cpp" (xopt Cpp dflags)
              -- Was the source literate Haskell
              let wasLhs = bool "" "-lhs" (isExtensionOf "lhs" modulePath)
              -- prefix any user-supplied options
              let ext =
                    fromMaybe (takeExtension modulePath) (newExtension options)
                      ++ usedCpp
                      ++ wasLhs

              -- Print the result
              let newContent = exactPrint astHsMod' annsHsMod'
              liftIO $ writeFile (modulePath -<.> ext) newContent

              -- Clean up: delete the GHC-generated imports file
              liftIO $ removeFile minImpFilePath
          where

            -- Generates the things that would be exportabe if there were no
            -- explicit export header, so suitable for replacing one
            exportable :: RnM [AvailInfo]
            exportable = do
              let rdr_env = tcg_rdr_env tcEnv
              let imports = tcg_imports tcEnv --  actually not needed for the Nothing case
              let this_mod = tcg_mod tcEnv
              exports <- exports_from_avail Nothing rdr_env imports this_mod
              return (snd exports)

            --  Replace a target module's imports
            --  See <https://github.com/facebookincubator/retrie/blob/master/Retrie/CPP.hs>
            replaceImports ::
              Monad m =>
              -- | annotations for the replacement imports
              Anns ->
              -- | the replacement imports
              [LImportDecl GhcPs] ->
              -- | target module
              ParsedSource ->
              TransformT m ParsedSource
            replaceImports anns minImports t@(L l m) =
              case importAction options of
                NoImportProcessing -> return t
                _ -> do
                  -- This does all the work
                  -- retrie has a neat `insertImports' function that also
                  -- deduplicates
                  imps <- graftT anns minImports
                  -- nudge down the imports list onto a new line
                  unless (null imps) $ setEntryDPT (head imps) (DP (2, 0))
                  return $ L l m {hsmodImports = imps}

            -- Add explict exports to the target module
            addExplicitExports ::
              Monad m =>
              -- | The list of exports to be added
              Avails ->
              -- | target module
              ParsedSource ->
              TransformT m ParsedSource
            addExplicitExports exports t@(L astLoc hsMod) =
              case exportAction options of
                NoExportProcessing -> return t
                AddExplicitExports ->
                  -- Only add explicit exports if there are none.
                  -- Seems to work even if there is no explict module declaration
                  -- presumably because the annotations that we generate are just
                  -- unused by exactPrint
                  if isNothing currentExplicitExports then result else return t
                ReplaceExports -> result
              where
                currentExplicitExports :: Maybe (Located [LIE GhcPs])
                currentExplicitExports = hsmodExports hsMod
                -- This does all the export replacement work
                result :: Monad m => TransformT m ParsedSource
                result
                  | null exports = return t -- there is nothing exportable
                  | otherwise = do
                    -- Generate the exports list
                    exportsList <- mapM mkExportAnnT exports
                    -- add commas in between and parens around
                    mapM_ addTrailingCommaT (init exportsList)
                    lExportsList <- mkLoc exportsList >>= mkParenT unLoc

                    -- No need to do any graftTing here as we have been modifying the
                    -- annotations in the current transformation
                    return $ L astLoc hsMod {hsmodExports = Just lExportsList}

    -- This version of the GHC function ignores implicit imports, as they
    -- cannot be parsed back in.  (There is an extraneous (implicit))
    -- It also provides for leaving out instance-only imports (eg, Data.List() )
    printMinimalImports' :: DynFlags -> FilePath -> [ImportDeclUsage] -> RnM ()
    printMinimalImports' dflags filename imports_w_usage =
      do
        imports' <- Smuggler2.Imports.getMinimalImports imports_w_usage
        liftIO $
          withFile
            filename
            WriteMode
            ( \h ->
                -- The neverQualify is important.  We are printing Names
                -- but they are in the context of an 'import' decl, and
                -- we never qualify things inside there
                -- E.g.   import Blag( f, b )
                -- not    import Blag( Blag.f, Blag.g )!
                printForUser dflags h neverQualify (vcat (map ppr (filter (letThrough . unLoc) imports')))
            )
      where
        notImplicit :: ImportDecl pass -> Bool
        notImplicit = not . ideclImplicit
        notInstancesOnly :: ImportDecl pass -> Bool
        notInstancesOnly i = case ideclHiding i of
          Just (False, L _ []) -> False
          _ -> True
        keepInstanceOnlyImports :: Bool
        keepInstanceOnlyImports = importAction options /= MinimiseImports
        letThrough :: ImportDecl pass -> Bool
        letThrough i = notImplicit i && (keepInstanceOnlyImports || notInstancesOnly i)

    -- Construct the path into which GHC's version of minimal imports is dumped
    mkMinimalImportsPath :: DynFlags -> Module -> FilePath
    mkMinimalImportsPath dflags this_mod
      | Just d <- dumpDir dflags = d </> basefn
      | otherwise = basefn
      where
        basefn =
          "smuggler2-" ++ moduleNameString (moduleName this_mod) ++ "."
            ++ fromMaybe "smuggler2" (newExtension options)
            ++ ".imports"

    options :: Options
    options = parseCommandLineOptions clopts

    strBufToStr :: StringBuffer -> String
    strBufToStr sb@(StringBuffer _ len _) = lexemeToString sb len