    {-# LANGUAGE CPP, DeriveDataTypeable, DeriveFunctor #-}
    module Extract (Module(..), extract) where

    import           Prelude hiding (mod, concat)
    import           Control.Monad
    import           Control.Applicative
    import           Control.Exception
    import           Data.List (partition, isSuffixOf)
    import           Data.Maybe
    import           Data.Foldable (concat)

    import           Control.DeepSeq (deepseq, NFData(rnf))
    import           Data.Generics

    #if __GLASGOW_HASKELL__ < 707
    import           GHC hiding (flags, Module, Located)
    import           MonadUtils (liftIO, MonadIO)
    #else
    import           GHC hiding (Module, Located)
    import           DynFlags
    import           MonadUtils (liftIO)
    #endif
    import           Exception (ExceptionMonad)
    import           System.Directory
    import           System.FilePath
    import           NameSet (NameSet)
    import           Coercion (Coercion)
    import           FastString (unpackFS)
    import           Digraph (flattenSCCs)

    import           System.Posix.Internals (c_getpid)

    import           GhcUtil (withGhc)
    import           Location hiding (unLoc)

    import           Util (convertDosLineEndings)
    import           Sandbox (getSandboxArguments)

    -- | A wrapper around `SomeException`, to allow for a custom `Show` instance.
    newtype ExtractError = ExtractError SomeException
      deriving Typeable

    instance Show ExtractError where
      show (ExtractError e) =
        unlines [
            "Ouch! Hit an error thunk in GHC's AST while extracting documentation."
          , ""
          , "    " ++ msg
          , ""
          , "This is most likely a bug in doctest."
          , ""
          , "Please report it here: https://github.com/sol/doctest-haskell/issues/new"
          ]
        where
          msg = case fromException e of
            Just (Panic s) -> "GHC panic: " ++ s
            _              -> show e

    instance Exception ExtractError

    -- | Documentation for a module grouped together with the modules name.
    data Module a = Module {
      moduleName    :: String
    , moduleSetup   :: Maybe a
    , moduleContent :: [a]
    } deriving (Eq, Functor)

    instance NFData a => NFData (Module a) where
      rnf (Module name setup content) = name `deepseq` setup `deepseq` content `deepseq` ()

    -- | Parse a list of modules.
    parse :: [String] -> IO [TypecheckedModule]
    parse args = withGhc args $ \modules_ -> withTempOutputDir $ do

      -- ignore additional object files
      let modules = filter (not . isSuffixOf ".o") modules_

      mapM (`guessTarget` Nothing) modules >>= setTargets
      mods <- depanal [] False

      mods' <- if needsTemplateHaskell mods then enableCompilation mods else return mods

      let sortedMods = flattenSCCs (topSortModuleGraph False mods' Nothing)
      reverse <$> mapM (parseModule >=> typecheckModule >=> loadModule) sortedMods
      where
        -- copied from Haddock/Interface.hs
        enableCompilation :: ModuleGraph -> Ghc ModuleGraph
        enableCompilation modGraph = do
    #if __GLASGOW_HASKELL__ < 707
          let enableComp d = d { hscTarget = defaultObjectTarget }
    #else
          let enableComp d = let platform = targetPlatform d
                             in d { hscTarget = defaultObjectTarget platform }
    #endif
          modifySessionDynFlags enableComp
          -- We need to update the DynFlags of the ModSummaries as well.
          let upd m = m { ms_hspp_opts = enableComp (ms_hspp_opts m) }
          let modGraph' = map upd modGraph
          return modGraph'

        -- copied from Haddock/GhcUtils.hs
        modifySessionDynFlags :: (DynFlags -> DynFlags) -> Ghc ()
        modifySessionDynFlags f = do
          dflags <- getSessionDynFlags
    #if __GLASGOW_HASKELL__ < 707
          _ <- setSessionDynFlags (f dflags)
    #else
          -- GHCi 7.7 now uses dynamic linking.
          let dflags' = case lookup "GHC Dynamic" (compilerInfo dflags) of
                Just "YES" -> gopt_set dflags Opt_BuildDynamicToo
                _          -> dflags
          _ <- setSessionDynFlags (f dflags')
    #endif
          return ()

        withTempOutputDir :: Ghc a -> Ghc a
        withTempOutputDir action = do
          tmp <- liftIO getTemporaryDirectory
          x   <- liftIO c_getpid
          let dir = tmp </> ".doctest-" ++ show x
          modifySessionDynFlags (setOutputDir dir)
          gbracket_
            (liftIO $ createDirectory dir)
            (liftIO $ removeDirectoryRecursive dir)
            action

        -- | A variant of 'gbracket' where the return value from the first computation
        -- is not required.
        gbracket_ :: ExceptionMonad m => m a -> m b -> m c -> m c
        gbracket_ before_ after thing = gbracket before_ (const after) (const thing)

        setOutputDir f d = d {
            objectDir  = Just f
          , hiDir      = Just f
          , stubDir    = Just f
          , includePaths = f : includePaths d
          }

    -- | Extract all docstrings from given list of files/modules.
    --
    -- This includes the docstrings of all local modules that are imported from
    -- those modules (possibly indirect).
    extract :: [String] -> IO [Module (Located String)]
    extract args = do
      sandboxArgs <- getSandboxArguments
      let args'  = args ++ sandboxArgs
      mods <- parse args'
      let docs = map (fmap (fmap convertDosLineEndings) . extractFromModule . tm_parsed_module) mods

      (docs `deepseq` return docs) `catches` [
          -- Re-throw AsyncException, otherwise execution will not terminate on
          -- SIGINT (ctrl-c).  All AsyncExceptions are re-thrown (not just
          -- UserInterrupt) because all of them indicate severe conditions and
          -- should not occur during normal operation.
          Handler (\e -> throw (e :: AsyncException))
        , Handler (throwIO . ExtractError)
        ]

    -- | Extract all docstrings from given module and attach the modules name.
    extractFromModule :: ParsedModule -> Module (Located String)
    extractFromModule m = Module name (listToMaybe $ map snd setup) (map snd docs)
      where
        isSetup = (== Just "setup") . fst
        (setup, docs) = partition isSetup (docStringsFromModule m)
        name = (moduleNameString . GHC.moduleName . ms_mod . pm_mod_summary) m

    -- | Extract all docstrings from given module.
    docStringsFromModule :: ParsedModule -> [(Maybe String, Located String)]
    docStringsFromModule mod = map (fmap (toLocated . fmap unpackDocString)) docs
      where
        source   = (unLoc . pm_parsed_source) mod

        -- we use dlist-style concatenation here
        docs     = header ++ exports ++ decls

        -- We process header, exports and declarations separately instead of
        -- traversing the whole source in a generic way, to ensure that we get
        -- everything in source order.
        header  = [(Nothing, x) | Just x <- [hsmodHaddockModHeader source]]
        exports = [(Nothing, L loc doc) | L loc (IEDoc doc) <- concat (hsmodExports source)]
        decls   = extractDocStrings (hsmodDecls source)

    type Selector a = a -> ([(Maybe String, LHsDocString)], Bool)

    -- | Ignore a subtree.
    ignore :: Selector a
    ignore = const ([], True)

    -- | Collect given value and descend into subtree.
    select :: a -> ([a], Bool)
    select x = ([x], False)

    -- | Extract all docstrings from given value.
    extractDocStrings :: Data a => a -> [(Maybe String, LHsDocString)]
    extractDocStrings = everythingBut (++) (([], False) `mkQ` fromLHsDecl
      `extQ` fromLDocDecl
      `extQ` fromLHsDocString
      `extQ` (ignore :: Selector NameSet)
      `extQ` (ignore :: Selector PostTcKind)

      -- HsExpr never contains any documentation, but it may contain error thunks.
      --
      -- Problematic are (non comprehensive):
      --
      --  * parallel list comprehensions
      --  * infix operators
      --
      `extQ` (ignore :: Selector (HsExpr RdrName))

      -- undefined before type checking
      `extQ` (ignore :: Selector Coercion)

    #if __GLASGOW_HASKELL__ >= 706
      -- hswb_kvs and hswb_tvs may be error thunks
      `extQ` (ignore :: Selector (HsWithBndrs [LHsType RdrName]))
      `extQ` (ignore :: Selector (HsWithBndrs [LHsType Name]))
      `extQ` (ignore :: Selector (HsWithBndrs (LHsType RdrName)))
      `extQ` (ignore :: Selector (HsWithBndrs (LHsType Name)))
    #endif
      )
      where
        fromLHsDecl :: Selector (LHsDecl RdrName)
        fromLHsDecl (L loc decl) = case decl of

          -- Top-level documentation has to be treated separately, because it has
          -- no location information attached.  The location information is
          -- attached to HsDecl instead.
          DocD x -> select (fromDocDecl loc x)

          _ -> (extractDocStrings decl, True)

        fromLDocDecl :: Selector LDocDecl
        fromLDocDecl (L loc x) = select (fromDocDecl loc x)

        fromLHsDocString :: Selector LHsDocString
        fromLHsDocString x = select (Nothing, x)

        fromDocDecl :: SrcSpan -> DocDecl -> (Maybe String, LHsDocString)
        fromDocDecl loc x = case x of
          DocCommentNamed name doc -> (Just name, L loc doc)
          _                        -> (Nothing, L loc $ docDeclDoc x)

    -- | Convert a docstring to a plain string.
    unpackDocString :: HsDocString -> String
    unpackDocString (HsDocString s) = unpackFS s