{-# LANGUAGE LambdaCase #-}

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

import Avail (AvailInfo, Avails)
import Control.Monad (unless, when)
import Data.Bool (bool)
import Data.List (intersect)
import Data.Maybe (fromMaybe, isJust, isNothing)
import Data.Version (showVersion)
import DynFlags
  ( DynFlags (dumpDir),
    HasDynFlags (getDynFlags),
    setUnsafeGlobalDynFlags,
    xopt,
    xopt_set,
  )
import ErrUtils (compilationProgressMsg, fatalErrorMsg, warningMsg)
import GHC
  ( GenLocated (L),
    GhcPs,
    GhcRn,
    HsModule (hsmodExports, hsmodImports),
    ImportDecl (ideclHiding, ideclImplicit, ideclName),
    LIE,
    LImportDecl,
    Located,
    ModSummary (ms_hspp_buf, ms_mod),
    Module (moduleName),
    ParsedSource,
    ml_hs_file,
    moduleNameString,
    ms_location,
    unLoc,
  )
import GHC.IO.Encoding (setLocaleEncoding, utf8)
import GHC.LanguageExtensions (Extension (Cpp, PatternSynonyms))
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.Exports (mkExportAnnT)
import Smuggler2.Imports (getMinimalImports)
import Smuggler2.Options
  ( ExportAction (AddExplicitExports, NoExportProcessing, ReplaceExports),
    ImportAction (MinimiseImports, NoImportProcessing),
    Options (..),
    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
  ( RnM,
    TcGblEnv (tcg_exports, tcg_imports, tcg_mod, tcg_rdr_env, tcg_rn_exports, tcg_rn_imports, tcg_used_gres),
    TcM,
  )

-- | '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 silently, if nothing to do
  | (importAction options == NoImportProcessing)
      && (exportAction options == NoExportProcessing) =
    return tcEnv
  | otherwise = do
    dflags <- getDynFlags
    liftIO $ setUnsafeGlobalDynFlags dflags -- this seems to be needed for Windows only
    liftIO $ compilationProgressMsg dflags ("smuggler2 " ++ showVersion version)

    -- Get imports  usage
    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.  Assumes
    -- that an open import (ideclHiding Nothing) has unused imports.
    let noUnusedImports =
          all
            ( \(L _ decl, used, unused) ->
                null unused && not (null used) && isJust (ideclHiding decl)
            )
            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 do
        liftIO $
          compilationProgressMsg
            dflags
            ( "smuggler2: nothing to do for module "
                ++ (moduleNameString . moduleName $ ms_mod modSummary)
            )
        return tcEnv
      else do
        when (importAction options == NoImportProcessing && not (null $ leaveOpenImports options))
          . liftIO
          $ warningMsg dflags (text "LeaveOpenModules ignored as NoImportProcessing also specified")
        when (importAction options == NoImportProcessing && not (null $ makeOpenImports options))
          . liftIO
          $ warningMsg dflags (text "MakeOpenModules ignored as NoImportProcessing also specified")

        -- 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.  @ghc-exactprint@ "index"es ('Anns') 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; don't change the tcEnv we
        -- were givem.
        tcEnv <$ smuggling dflags minImpFilePath
  where

    -- The original imports
    imports :: [LImportDecl GhcRn]
    imports = tcg_rn_imports tcEnv

    -- 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

          -- If a module is imported open, then pattern synonyms can be
          -- imported from it with the @PatternSynonym@ language extension.  If,
          -- however, a pattern synonym is imported explicitly, the extension is
          -- required. So we switch on the @PatternSynonym@ extension for
          -- parsing the minimal imports.
          let dflags' = xopt_set dflags PatternSynonyms

          -- 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,
              -- with open imports processed as specified
              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'
              let newModulePath = modulePath -<.> ext
              liftIO $ writeFile newModulePath newContent

              liftIO $
                compilationProgressMsg
                  dflags
                  ( "smuggler2: output written to " ++ newModulePath )

              -- 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 importAvails = tcg_imports tcEnv --  actually not needed for the Nothing case
              let this_mod = tcg_mod tcEnv
              exports <- exports_from_avail Nothing rdr_env importAvails 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,
    -- @import Data.List ()@) and handles the preservation of open imports
    -- (eg, @import Prelude@)
    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 . leaveOpen)
                          (filter letThrough 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
        -- Ignore explicit instance only imports, unless the 'MinimiseImports'
        -- option is specified
        letThrough :: LImportDecl pass -> Bool
        letThrough (L _ i) = notImplicit i && (keepInstanceOnlyImports || notInstancesOnly i)
        --
        leaveOpen :: LImportDecl pass -> LImportDecl pass
        leaveOpen (L l decl) = L l $ case ideclHiding decl of
          Just (False, L _ _) -- ie, not hiding
            | thisModule `elem` kModules || thisModule `elem` mModules -> decl {ideclHiding = Nothing}
          _ -> decl
          where
            thisModule = unLoc (ideclName decl)
            mModules = makeOpenImports options
            lModules = leaveOpenImports options
            oModules = unLoc . ideclName <$> filter isOpen (unLoc <$> imports) -- original open imports
              where
                isOpen = isNothing . ideclHiding
            kModules = lModules `intersect` oModules -- original open imports to leave open

    -- 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