content: {-# LANGUAGE DeriveAnyClass    #-}
  {-# LANGUAGE DeriveGeneric     #-}
  {-# LANGUAGE NamedFieldPuns    #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards   #-}
  {-# LANGUAGE TupleSections     #-}
  {-# LANGUAGE TypeFamilies      #-}
  
  -- | Haskell Imports and Exports tool
  --
  -- This tool parses imports and exports from Haskell source files and provides
  -- analysis on these imports. For example, you can check whether consistent
  -- import aliases are used across your codebase.
  
  module Main (main) where
  
  import qualified Data.Aeson                              as JSON
  import qualified Data.ByteString.Lazy.Char8              as LBS8
  import qualified Data.Csv                                as Csv
  import qualified Data.Map                                as Map
  import qualified Data.Set                                as Set
  import qualified Data.Text                               as T
  import qualified Data.Text.IO                            as T
  import qualified Dot
  import qualified GHC
  import qualified GHC.Paths
  import qualified Language.Haskell.GHC.ExactPrint.Parsers as ExactPrint
  import qualified Options.Applicative                     as O
  import qualified System.FilePath                         as FP
  
  import Data.Aeson.Encode.Pretty   (encodePretty)
  import Data.Function              ((&))
  import Data.List                  (intercalate)
  import Data.Maybe                 (catMaybes, mapMaybe)
  import Data.Text                  (Text)
  import GHC.Generics               (Generic)
  import GHC.Hs.Extension           (GhcPs)
  import GHC.Types.Error            (getMessages)
  import GHC.Types.Name.Occurrence  (occNameString)
  import GHC.Types.Name.Reader      (rdrNameOcc)
  import GHC.Unit.Module.Name       (moduleNameString)
  import GHC.Utils.Error            (pprMsgEnvelopeBagWithLoc)
  import System.Directory.Recursive (getFilesRecursive)
  import System.Exit                (exitFailure)
  
  -- TYPES
  
  data Options =
    Options
      { command :: Command
      , sources :: [FilePath]
      }
  
  data Command
    = Dump OutputFormat
    | GraphSymbols
    | GraphModules
    | CheckAliases
    | CheckWildcards [Text]
  
  data OutputFormat = OutputCsv | OutputJson
  
  data ImportedSymbol =
    ImportedSymbol
      { impFromModule :: Text
      , impModule     :: Text
      , impQualified  :: ImportQualified
      , impAlias      :: Maybe Text
      , impType       :: ImportType
      , impSymbol     :: Maybe Text
      , impInternal   :: ModuleInternal
      , impSource     :: FilePath
      , impFile       :: FilePath
      }
      deriving (Generic, Csv.ToNamedRecord, Csv.DefaultOrdered, JSON.ToJSON)
  
  data ImportQualified
    = Qualified
    | NotQualified
    deriving (Eq, Generic, JSON.ToJSON)
  
  instance Csv.ToField ImportQualified where
    toField Qualified    = "qualified"
    toField NotQualified = "not qualified"
  
  data ModuleInternal
    = Internal
    | External
    deriving (Eq, Generic, JSON.ToJSON)
  
  instance Csv.ToField ModuleInternal where
    toField Internal = "internal"
    toField External = "external"
  
  data ImportType
    = Wildcard
    | Hiding
    | Explicit
    deriving (Eq, Generic, JSON.ToJSON)
  
  instance Csv.ToField ImportType where
    toField Wildcard = "wildcard"
    toField Hiding   = "hiding"
    toField Explicit = "explicit"
  
  -- | Mapping of modules to their aliases and to the files they are found in
  type ModuleAliases = [(Text, [(Text, [FilePath])])]
  
  -- | Mapping of modules to files
  type WildcardImports = [(FilePath, [Text])]
  
  
  -- MAIN
  
  main :: IO ()
  main =
    run =<< O.customExecParser prefs infoOpts
    where
      prefs = O.prefs $ O.subparserInline <> O.showHelpOnEmpty
      infoOpts =
        O.info (O.helper <*> opts) $
          O.fullDesc
          <> O.header "hsie - Swiss army knife for HaSkell Imports and Exports"
          <> O.progDesc "Parse Haskell code to analyze imports and exports"
      opts =
        Options <$> commandOption <*> O.some srcOption
      srcOption =
        O.argument O.str $
          O.metavar "SRCDIR"
          <> O.help "Haskell source directory"
          <> O.action "directory"
      commandOption =
        O.subparser $
          command "dump-imports" "Dump imported symbols as CSV or JSON"
            (Dump <$> jsonOutputFlag)
          <> command "graph-modules" "Print dot graph of module imports"
               (pure GraphModules)
          <> command "graph-symbols" "Print dot graph of symbol imports"
               (pure GraphSymbols)
          <> command "check-aliases"
               "Check that aliases of imported modules are consistent"
               (pure CheckAliases)
          <> command "check-wildcards"
               "Check that no modules are imported as unqualified wildcards"
               (CheckWildcards <$> O.many okModuleOption)
      command name desc options =
        O.command name . O.info (O.helper <*> options) $ O.progDesc desc
      jsonOutputFlag =
        O.flag OutputCsv OutputJson $
          O.long "json" <> O.short 'j' <> O.help "Output JSON"
      okModuleOption =
        O.strOption $
          O.long "ok"
          <> O.short 'o'
          <> O.metavar "OKMODULE"
          <> O.help "Module that is ok to import as unqualified wildcard"
  
  run :: Options -> IO ()
  run Options{command, sources} =
    runCommand command . markInternal . concat =<< mapM sourceSymbols sources
    where
      runCommand :: Command -> [ImportedSymbol] -> IO ()
      runCommand (Dump format) = LBS8.putStr . dump format
      runCommand GraphSymbols = T.putStr . symbolsGraph
      runCommand GraphModules = T.putStr . Dot.encode . modulesGraph
      runCommand CheckAliases = runInconsistentAliases . inconsistentAliases
      runCommand (CheckWildcards okModules) = runWildcards . wildcards okModules
  
      runInconsistentAliases :: ModuleAliases -> IO ()
      runInconsistentAliases [] = T.putStrLn "No inconsistent module aliases found."
      runInconsistentAliases xs = T.putStr (formatInconsistentAliases xs) >> exitFailure
  
      runWildcards :: WildcardImports -> IO ()
      runWildcards [] = T.putStrLn "No unwanted wildcard imports found."
      runWildcards xs = T.putStr (formatWildcards xs) >> exitFailure
  
  -- | Mark imports from modules that are among the analyzed ones as internal.
  markInternal :: [ImportedSymbol] -> [ImportedSymbol]
  markInternal symbols =
    fmap mark symbols
    where
      mark s = s { impInternal = if isInternal s then Internal else External }
      isInternal = flip Set.member internalModules . impModule
      internalModules = Set.fromList $ fmap impFromModule symbols
  
  
  -- SYMBOLS
  
  -- | Parse all imported symbols from a source of Haskell source files
  sourceSymbols :: FilePath -> IO [ImportedSymbol]
  sourceSymbols source = do
    files <- filterExts [".hs", ".imports"] <$> getFilesRecursive source
    concat <$> mapM moduleSymbols files
    where
      filterExts exts = filter $ flip elem exts . FP.takeExtension
      moduleSymbols filepath = do
        GHC.HsModule{..} <- parseModule filepath
        return $ concatMap (importSymbols source filepath . GHC.unLoc) hsmodImports
  
  -- | Parse a Haskell module
  parseModule :: FilePath -> IO GHC.HsModule
  parseModule filepath = do
    result <- ExactPrint.parseModule GHC.Paths.libdir filepath
    case result of
      Right hsmod ->
        return $ GHC.unLoc hsmod
      Left errs ->
        fail $ "Errors with " <> show filepath <> ":\n    "
          <> show (pprMsgEnvelopeBagWithLoc $ getMessages errs)
  
  -- | Symbols imported in an import declaration.
  --
  -- If the import is a wildcard, i.e. no symbols are selected for import, then
  -- only one item is returned.
  importSymbols :: FilePath -> FilePath -> GHC.ImportDecl GhcPs -> [ImportedSymbol]
  importSymbols source filepath GHC.ImportDecl{..} =
    case ideclHiding of
      Just (hiding, syms) ->
        symbol (if hiding then Hiding else Explicit) . Just . GHC.unLoc <$> GHC.unLoc syms
      Nothing ->
        [ symbol Wildcard Nothing ]
    where
      symbol hiding sym =
        ImportedSymbol
          { impFile = relativePath filepath
          , impSource = source
          , impFromModule = T.pack $ moduleFromPath filepath
          , impModule = T.pack . moduleNameString . GHC.unLoc $ ideclName
          , impQualified = if ideclQualified /= GHC.NotQualified then Qualified else NotQualified
          , impAlias = T.pack . moduleNameString . GHC.unLoc <$> ideclAs
          , impInternal = External
          , impType = hiding
          , impSymbol = T.pack . occNameString . rdrNameOcc . GHC.ieName <$> sym
          }
      moduleFromPath =
        intercalate "." . FP.splitDirectories . FP.dropExtension . relativePath
      relativePath = FP.makeRelative source
  
  
  -- DUMP
  
  -- | Dump list of symbols as CSV or JSON
  dump :: OutputFormat -> [ImportedSymbol] -> LBS8.ByteString
  dump OutputCsv  = Csv.encodeDefaultOrderedByName
  dump OutputJson = encodePretty
  
  
  -- ALIASES
  
  -- | Find modules that are imported under different aliases
  inconsistentAliases :: [ImportedSymbol] -> ModuleAliases
  inconsistentAliases symbols =
    foldr (insertSetMapMap . moduleAlias) Map.empty symbols
      & Map.map (aliases . Map.toList)
      & Map.filter ((<) 1 . length)
      & Map.toList
    where
      moduleAlias ImportedSymbol{..} =
        (impModule, impAlias, FP.joinPath [impSource, impFile])
      insertSetMapMap (k1, k2, v) =
        Map.insertWith (Map.unionWith Set.union) k1
          (Map.singleton k2 $ Set.singleton v)
      aliases :: [(Maybe Text, Set.Set FilePath)] -> [(Text, [FilePath])]
      aliases = mapMaybe (\(k, v) -> fmap (, Set.toList v) k)
  
  formatInconsistentAliases :: ModuleAliases -> Text
  formatInconsistentAliases modules =
    "The following imports have inconsistent aliases:\n\n"
      <> T.concat (fmap formatModule modules)
    where
      formatModule (modName, aliases) =
        "Module '"
          <> modName
          <> "' has the aliases:\n"
          <> T.concat (fmap formatAlias aliases)
          <> "\n"
      formatAlias (alias, sourceFiles) =
        "  '"
          <> alias
          <> "' in file"
          <> (if length sourceFiles > 2 then "s" else "")
          <> ":\n"
          <> T.concat (fmap formatFile sourceFiles)
      formatFile sourceFile =
        "    " <> T.pack sourceFile <> "\n"
  
  
  -- WILDCARDS
  
  -- | Find modules that are imported as wildcards, excluding whitelisted modules.
  --
  -- Wildcard imports are ones that are not qualified and do not specify which
  -- symbols should be imported.
  wildcards :: [Text] -> [ImportedSymbol] -> WildcardImports
  wildcards okModules =
    groupByFile . filter isWildcard . filter (not . isOkModule)
    where
      isWildcard ImportedSymbol{..} =
        impQualified == NotQualified && impType /= Explicit
      isOkModule = flip Set.member (Set.fromList okModules) . impModule
      groupByFile = Map.toList . fmap Set.toList . foldr insertMap Map.empty
      insertMap ImportedSymbol{..} =
        Map.insertWith Set.union impFile (Set.singleton impModule)
  
  formatWildcards :: WildcardImports -> Text
  formatWildcards files =
    "Modules in the following files were imported as wildcards:\n\n"
      <> T.concat (fmap formatFile files)
    where
      formatFile (filepath, modules) =
        "In " <> T.pack filepath <> ":\n" <> T.concat (fmap formatModule modules) <> "\n"
      formatModule moduleName = "  " <> moduleName <> "\n"
  
  
  -- GRAPHS
  
  modulesGraph :: [ImportedSymbol] -> Dot.DotGraph
  modulesGraph symbols =
    Dot.DotGraph Dot.Strict Dot.Directed (Just "Modules") $ fmap edge edges
    where
      edge (from, to) =
        Dot.StatementEdge $ Dot.EdgeStatement
          (Dot.ListTwo (edgeNode from) (edgeNode to) mempty) mempty
      edgeNode t = Dot.EdgeNode $ Dot.NodeId (Dot.Id t) Nothing
      edges = unique . fmap edgeTuple . filter ((==) Internal . impInternal) $ symbols
      edgeTuple ImportedSymbol{..} = (impFromModule, impModule)
      unique = Set.toList . Set.fromList
  
  -- Building Text directly as the Dot package currently doesn't support subgraphs.
  symbolsGraph :: [ImportedSymbol] -> Text
  symbolsGraph symbols =
    "digraph Symbols {\n"
      <> "  rankdir=LR\n"
      <> "  ranksep=5\n"
      <> T.concat (fmap edge edges)
      <> T.concat (fmap cluster symbolsByModule)
      <> "}\n"
    where
      edge (from, to, symbol) =
        "  "
          <> quoted from
          <> " -> "
          <> quoted (to <> maybe "" ("." <>) symbol)
          <> "\n"
      cluster (moduleName, clusterSymbols) =
        "  subgraph "
          <> quoted ("cluster_" <> moduleName)
          <> " {\n"
          <> "    " <> quoted moduleName <> "\n"
          <> T.concat (fmap (clusterNode moduleName) clusterSymbols)
          <> "  }\n"
      clusterNode moduleName symbol =
        "    " <> quoted (moduleName <> "." <> symbol) <> "\n"
      quoted t = "\"" <> t <> "\""
      edges = unique . fmap edgeTuple . filter ((==) Internal . impInternal) $ symbols
      edgeTuple ImportedSymbol{..} = (impFromModule, impModule, impSymbol)
      unique = Set.toList . Set.fromList
      symbolsByModule =
        Map.toList . Map.map (catMaybes . Set.toList) . foldr insertMap Map.empty $ edges
      insertMap (_, to, symbol) = Map.insertWith Set.union to $ Set.singleton symbol
  