{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# LANGUAGE ImportQualifiedPost #-}

module Main where

import GHC.Driver.Env
import GHC.Unit.Module
import GHC.Driver.Session
import GHC.Driver.Main
import GHC
import Control.Monad
import Control.Monad.IO.Class
import System.Environment
import GHC.Unit.Module.Deps
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set

-- Example invocation:
--  inplace/bin/count-deps `inplace/bin/ghc-stage2 --print-libdir` "GHC.Parser"
main :: IO ()
main = do
  args <- getArgs
  case args of
    [libdir, modName, "--dot"] -> printDeps libdir modName True
    [libdir, modName] -> printDeps libdir modName False
    _ -> fail "usage: count-deps libdir module [--dot]"

dotSpec :: String -> Map.Map String [String] -> String
dotSpec name g =
 "digraph \"" ++ name ++ "\" {\n" ++
    Map.foldlWithKey' f "" g ++ "}\n"
  where
    f acc k ns = acc ++ concat ["  " ++ show k ++ " -> " ++ show n ++ ";\n" | n <- ns]

printDeps :: String -> String -> Bool -> IO ()
printDeps libdir modName dot = do
  modGraph <-
    Map.map (map moduleNameString) .
      Map.mapKeys moduleNameString <$> calcDeps modName libdir
  if not dot then
    do
      let modules = Map.keys modGraph
      putStrLn $ "Found " ++ modName ++ " module dependencies"
      forM_ modules putStrLn
  else
    -- * Copy the digraph output to a file ('deps.dot' say)
    -- * To render it, use a command along the lines of
    --   'tred deps.dot > deps-tred.dot && dot -Tpdf -o deps.pdf deps-tred.dot'
    putStr $ dotSpec modName modGraph

calcDeps :: String -> FilePath -> IO (Map.Map ModuleName [ModuleName])
calcDeps modName libdir =
  defaultErrorHandler defaultFatalMessager defaultFlushOut $ do
    runGhc (Just libdir) $ do
        df <- getSessionDynFlags
        logger <- getLogger
        (df, _, _) <- parseDynamicFlags logger df [noLoc "-package=ghc"]
        setSessionDynFlags df
        case lookup "Project Unit Id" (compilerInfo df) of
          Nothing -> fail "failed to find ghc's unit-id in the compiler info"
          Just ghcUnitId -> do
            env <- getSession
            loop ghcUnitId env Map.empty [mkModuleName modName]
  where
    -- Source imports are only guaranteed to show up in the 'mi_deps'
    -- of modules that import them directly and don’t propagate
    -- transitively so we loop.
    loop :: String -> HscEnv -> Map.Map ModuleName [ModuleName] -> [ModuleName] -> Ghc (Map.Map ModuleName [ModuleName])
    loop ghcUnitId env modules (m : ms) =
      if m `Map.member` modules
        then loop ghcUnitId env modules ms
        else do
          mi <- liftIO $ hscGetModuleInterface env (mkModule ghcUnitId m)
          let deps = modDeps mi
          modules <- return $ Map.insert m [] modules
          loop ghcUnitId env (Map.insert m deps modules) $ ms ++ filter (not . (`Map.member` modules)) deps
    loop _ _ modules [] = return modules

    mkModule :: String -> ModuleName -> Module
    mkModule ghcUnitId = Module (stringToUnit ghcUnitId)

    modDeps :: ModIface -> [ModuleName]
    modDeps mi = map (gwib_mod . snd) $ Set.toList $ dep_direct_mods (mi_deps mi)