{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}

-- | Get information on modules, identifiers, etc.

module GhciInfo (collectInfo,getModInfo,showppr) where

import           ConLike
import           Control.Exception
import           Control.Monad
import qualified CoreUtils
import           Data.Data
import qualified Data.Generics
import           Data.List
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import           Data.Maybe
import           Data.Time
import           DataCon
import           Desugar
import           DynFlags
import           GHC
import           GhcMonad
import           GhciTypes
import           HscTypes
import           Intero.Compat
import           Outputable
import           Prelude hiding (mod)
import           System.Directory
import           TcHsSyn
import           Var

#if __GLASGOW_HASKELL__ <= 802
import           NameSet
#endif

#if MIN_VERSION_ghc(7,8,3)
#else
import           Bag
#endif

-- | Collect type info data for the loaded modules.
collectInfo :: (GhcMonad m)
            => Map ModuleName ModInfo -> [ModuleName] -> m (Map ModuleName ModInfo)
collectInfo ms loaded =
  do df <- getSessionDynFlags
     -- Generate for all modules in interpreted mode.
     invalidated <-
       liftIO (if hscTarget df == HscInterpreted
                  then return loaded
                  else filterM cacheInvalid loaded)
     if null invalidated
        then return ms
        else do liftIO (putStrLn ("Collecting type info for " ++
                                  show (length invalidated) ++
                                  " module(s) ... "))
                foldM (\m name ->
                         gcatch (do info <- getModInfo name
                                    return (M.insert name info m))
                                (\(e :: SomeException) ->
                                   do liftIO (putStrLn ("Error while getting type info from " ++
                                                        showppr df name ++
                                                        ": " ++ show e))
                                      return m))
                      ms
                      invalidated
  where cacheInvalid name =
          case M.lookup name ms of
            Nothing -> return True
            Just mi ->
              do let fp =
#if MIN_VERSION_ghc(8,0,4)
                       ml_hi_file (ms_location (modinfoSummary mi))
#else
                       ml_obj_file (ms_location (modinfoSummary mi))
#endif
                     last' = modinfoLastUpdate mi
                 exists <- doesFileExist fp
                 if exists
                    then do mod <- getModificationTime fp
                            return (mod > last')
                    else return True

-- | Get info about the module: summary, types, etc.
getModInfo :: (GhcMonad m) => ModuleName -> m ModInfo
getModInfo name =
  do m <- getModSummary name
     p <- parseModule m
     let location = getModuleLocation (parsedSource p)
     typechecked <- typecheckModuleSilent p
     let Just (_, imports, _, _) = renamedSource typechecked
     allTypes <- processAllTypeCheckedModule typechecked
     let i = tm_checked_module_info typechecked
     now <- liftIO getCurrentTime
     return (ModInfo m allTypes i now imports location)

-- | Type-check the module without logging messages.
typecheckModuleSilent :: GhcMonad m => ParsedModule -> m TypecheckedModule
#if MIN_VERSION_ghc(8,0,1)
typecheckModuleSilent parsed = do
  typecheckModule
    parsed
    { GHC.pm_mod_summary =
        (GHC.pm_mod_summary parsed)
        { HscTypes.ms_hspp_opts =
            (HscTypes.ms_hspp_opts (GHC.pm_mod_summary parsed))
            {log_action = nullLogAction}
        }
    }
  where
    nullLogAction _df _reason _sev _span _style _msgdoc = return ()
#else
typecheckModuleSilent parsed = do
  typecheckModule
    parsed
    { GHC.pm_mod_summary =
        (GHC.pm_mod_summary parsed)
        { HscTypes.ms_hspp_opts =
            (HscTypes.ms_hspp_opts (GHC.pm_mod_summary parsed))
            {log_action = nullLogAction}
        }
    }
  where
    nullLogAction _df _reason _sev _span _style = return ()
#endif

getModuleLocation :: ParsedSource -> SrcSpan
getModuleLocation pSource = case hsmodName (unLoc pSource) of
  Just located -> getLoc located
  Nothing -> noSrcSpan

-- | Get ALL source spans in the module.
processAllTypeCheckedModule :: GhcMonad m
                            => TypecheckedModule -> m [SpanInfo]
processAllTypeCheckedModule tcm =
  do let tcs = tm_typechecked_source tcm
         bs = listifyAllSpans tcs :: [LHsBind StageReaderId]
         es = listifyAllSpans tcs :: [LHsExpr StageReaderId]
         ps = listifyAllSpans tcs :: [LPat StageReaderId]
     bts <- mapM (getTypeLHsBind tcm) bs
     ets <- mapM (getTypeLHsExpr tcm) es
     pts <- mapM (getTypeLPat tcm) ps
     return (mapMaybe toSpanInfo (sortBy cmp (concat bts ++ catMaybes (concat [ets,pts]))))
  where cmp (_,a,_) (_,b,_)
          | a `isSubspanOf` b = LT
          | b `isSubspanOf` a = GT
          | otherwise = EQ

getTypeLHsBind :: (GhcMonad m)
               => TypecheckedModule
               -> LHsBind StageReaderId
               -> m [(Maybe Id,SrcSpan,Type)]
#if MIN_VERSION_ghc(7,8,3)
getTypeLHsBind _ (L _spn FunBind{fun_id = pid,fun_matches = MG{}}) =
  return (return (Just (unLoc pid),getLoc pid,varType (unLoc pid)))
#else
getTypeLHsBind _ (L _spn FunBind{fun_id = pid,fun_matches = MG _ _ _typ}) =
  return (return (Just (unLoc pid),getLoc pid,varType (unLoc pid)))
#endif
#if MIN_VERSION_ghc(7,8,3)
#else
getTypeLHsBind m (L _spn AbsBinds{abs_binds = binds}) =
  fmap concat
       (mapM (getTypeLHsBind m)
             (map snd (bagToList binds)))
#endif
getTypeLHsBind _ _ = return []
-- getTypeLHsBind _ x =
--   do df <- getSessionDynFlags
--      error ("getTypeLHsBind: unhandled case: " ++
--             showppr df x)

getTypeLHsExpr :: (GhcMonad m)
               => TypecheckedModule
               -> LHsExpr StageReaderId
               -> m (Maybe (Maybe Id,SrcSpan,Type))
getTypeLHsExpr _ e =
  do hs_env <- getSession
     (_,mbe) <- liftIO (deSugarExpr hs_env e)
     case mbe of
       Nothing -> return Nothing
       Just expr ->
         return (Just (case unwrapVar (unLoc e) of
#if __GLASGOW_HASKELL__ >= 806
                         HsVar _ (L _ i) -> Just i
#elif __GLASGOW_HASKELL__ >= 800
                         HsVar (L _ i) -> Just i
#else
                         HsVar i -> Just i
#endif
                         _ -> Nothing
                      ,getLoc e
                      ,CoreUtils.exprType expr))
  where
#if __GLASGOW_HASKELL__ >= 806
    unwrapVar (HsWrap _ _ var) = var
#else
    unwrapVar (HsWrap _ var) = var
#endif
    unwrapVar e' = e'

-- | Get id and type for patterns.
getTypeLPat :: (GhcMonad m)
            => TypecheckedModule -> LPat StageReaderId -> m (Maybe (Maybe Id,SrcSpan,Type))
getTypeLPat _ (L spn pat) =
  return (Just (getMaybeId pat,spn,getPatType pat))
  where
    getPatType (ConPatOut (L _ (RealDataCon dc)) _ _ _ _ _ _) =
      dataConRepType dc
    getPatType pat' = hsPatType pat'
#if __GLASGOW_HASKELL__ >= 806
    getMaybeId (VarPat _ (L _ vid)) = Just vid
#elif __GLASGOW_HASKELL__ >= 800
    getMaybeId (VarPat (L _ vid)) = Just vid
#else
    getMaybeId (VarPat vid) = Just vid
#endif
    getMaybeId _ = Nothing

-- | Get ALL source spans in the source.
listifyAllSpans :: Typeable a
                => TypecheckedSource -> [Located a]
listifyAllSpans tcs =
  listifyStaged TypeChecker p tcs
  where p (L spn _) = isGoodSrcSpan spn

listifyStaged :: Typeable r
              => Stage -> (r -> Bool) -> Data.Generics.GenericQ [r]
#if __GLASGOW_HASKELL__ <= 802
listifyStaged s p =
  everythingStaged
    s
    (++)
    []
    ([] `Data.Generics.mkQ`
     (\x -> [x | p x]))
#else
listifyStaged _ p = Data.Generics.listify p
#endif

------------------------------------------------------------------------------
-- The following was taken from 'ghc-syb-utils'
--
-- ghc-syb-utils:
--     https://github.com/nominolo/ghc-syb

-- | Ghc Ast types tend to have undefined holes, to be filled
--   by later compiler phases. We tag Asts with their source,
--   so that we can avoid such holes based on who generated the Asts.
data Stage
  = Parser
  | Renamer
  | TypeChecker
  deriving (Eq,Ord,Show)

-- | Like 'everything', but avoid known potholes, based on the 'Stage' that
--   generated the Ast.
#if __GLASGOW_HASKELL__ <= 802
everythingStaged :: Stage -> (r -> r -> r) -> r -> Data.Generics.GenericQ r -> Data.Generics.GenericQ r
everythingStaged stage k z f x
  | (const False `Data.Generics.extQ` postTcType `Data.Generics.extQ` fixity `Data.Generics.extQ` nameSet) x = z
  | otherwise = foldl k (f x) (gmapQ (everythingStaged stage k z f) x)
  where nameSet    = const (stage `elem` [Parser,TypeChecker]) :: NameSet -> Bool
#if __GLASGOW_HASKELL__ >= 709
        postTcType = const (stage<TypeChecker)                 :: PostTc Id Type -> Bool
#else
        postTcType = const (stage<TypeChecker)                 :: PostTcType -> Bool
#endif
        fixity     = const (stage<Renamer)                     :: GHC.Fixity -> Bool
#endif


-- | Pretty print the types into a 'SpanInfo'.
toSpanInfo :: (Maybe Id,SrcSpan,Type) -> Maybe SpanInfo
toSpanInfo (n,mspan,typ) =
  case mspan of
    RealSrcSpan spn ->
      Just (SpanInfo (srcSpanStartLine spn)
                     (srcSpanStartCol spn - 1)
                     (srcSpanEndLine spn)
                     (srcSpanEndCol spn - 1)
                     (Just typ)
                     n)
    _ -> Nothing

-- | Pretty print something to string.
showppr :: Outputable a
        => DynFlags -> a -> String
showppr dflags =
  showSDocForUser dflags neverQualify .
  ppr
