{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
module Info
    ( getIdentifierInfo
    , getType
    ) where

import Data.Generics (GenericQ, mkQ, extQ, gmapQ)
import Data.List (find, sortBy, intersperse)
import Data.Maybe (catMaybes, fromMaybe)
import Data.Typeable (Typeable)
import MonadUtils (liftIO)
import qualified CoreUtils
import qualified Desugar
#if __GLASGOW_HASKELL__ >= 706
import qualified DynFlags
#endif
#if __GLASGOW_HASKELL__ >= 708
import qualified HsExpr
#else
import qualified TcRnTypes
#endif
import qualified GHC
import qualified HscTypes
import qualified NameSet
import qualified Outputable
import qualified PprTyThing
import qualified Pretty
import qualified TcHsSyn

import GhcTypes (getModSummaries, TypecheckI)

getIdentifierInfo :: FilePath -> String -> GHC.Ghc (Either String String)
getIdentifierInfo file identifier =
    withModSummary file $ \m -> do
#if __GLASGOW_HASKELL__ >= 706
        GHC.setContext [GHC.IIModule (GHC.moduleName (GHC.ms_mod m))]
#elif __GLASGOW_HASKELL__ >= 704
        GHC.setContext [GHC.IIModule (GHC.ms_mod m)]
#else
        GHC.setContext [GHC.ms_mod m] []
#endif
        GHC.handleSourceError (return . Left . show) $
            fmap Right (infoThing identifier)

getType :: FilePath -> (Int, Int) -> GHC.Ghc (Either String [((Int, Int, Int, Int), String)])
getType file (line, col) =
    withModSummary file $ \m -> do
        p <- GHC.parseModule m
        typechecked <- GHC.typecheckModule p
        types <- processTypeCheckedModule typechecked (line, col)
        return (Right types)

withModSummary :: String -> (HscTypes.ModSummary -> GHC.Ghc (Either String a)) -> GHC.Ghc (Either String a)
withModSummary file action = do
    modSummary <- getModuleSummary file
    case modSummary of
        Nothing -> return (Left "Module not found in module graph")
        Just m -> action m

getModuleSummary :: FilePath -> GHC.Ghc (Maybe GHC.ModSummary)
getModuleSummary file = do
    modSummaries <- getModSummaries
    case find (moduleSummaryMatchesFilePath file) modSummaries of
        Nothing -> return Nothing
        Just moduleSummary -> return (Just moduleSummary)

moduleSummaryMatchesFilePath :: FilePath -> GHC.ModSummary -> Bool
moduleSummaryMatchesFilePath file moduleSummary =
    let location = GHC.ms_location moduleSummary
        location_file = GHC.ml_hs_file location
    in case location_file of
        Just f -> f == file
        Nothing -> False

------------------------------------------------------------------------------
-- Most of the following code was taken from the source code of 'ghc-mod' (with
-- some stylistic changes)
--
-- ghc-mod:
--     http://www.mew.org/~kazu/proj/ghc-mod/
--     https://github.com/kazu-yamamoto/ghc-mod/

processTypeCheckedModule :: GHC.TypecheckedModule -> (Int, Int) -> GHC.Ghc [((Int, Int, Int, Int), String)]
processTypeCheckedModule tcm (line, col) = do
    let tcs = GHC.tm_typechecked_source tcm
        bs = listifySpans tcs (line, col) :: [GHC.LHsBind TypecheckI]
        es = listifySpans tcs (line, col) :: [GHC.LHsExpr TypecheckI]
        ps = listifySpans tcs (line, col) :: [GHC.LPat TypecheckI]
    bts <- mapM (getTypeLHsBind tcm) bs
    ets <- mapM (getTypeLHsExpr tcm) es
    pts <- mapM (getTypeLPat tcm) ps
#if __GLASGOW_HASKELL__ >= 706
    dflags <- DynFlags.getDynFlags
    return $ map (toTup dflags) $
#else
    return $ map toTup $
#endif
        sortBy cmp $ catMaybes $ concat [ets, bts, pts]
    where
    cmp (a, _) (b, _)
        | a `GHC.isSubspanOf` b = LT
        | b `GHC.isSubspanOf` a = GT
        | otherwise = EQ

#if __GLASGOW_HASKELL__ >= 706
toTup :: GHC.DynFlags -> (GHC.SrcSpan, GHC.Type) -> ((Int, Int, Int, Int), String)
toTup dflags (spn, typ) = (fourInts spn, pretty dflags typ)
#else
toTup :: (GHC.SrcSpan, GHC.Type) -> ((Int, Int, Int, Int), String)
toTup (spn, typ) = (fourInts spn, pretty typ)
#endif

fourInts :: GHC.SrcSpan -> (Int, Int, Int, Int)
fourInts = fromMaybe (0, 0, 0, 0) . getSrcSpan

getSrcSpan :: GHC.SrcSpan -> Maybe (Int, Int, Int, Int)
getSrcSpan (GHC.RealSrcSpan spn) =
    Just (GHC.srcSpanStartLine spn
         , GHC.srcSpanStartCol spn
         , GHC.srcSpanEndLine spn
         , GHC.srcSpanEndCol spn)
getSrcSpan _ = Nothing

getTypeLHsBind :: GHC.TypecheckedModule -> GHC.LHsBind TypecheckI -> GHC.Ghc (Maybe (GHC.SrcSpan, GHC.Type))
#if __GLASGOW_HASKELL__ >= 806
getTypeLHsBind _ (GHC.L spn GHC.FunBind{GHC.fun_matches = grp}) = return $ Just (spn, HsExpr.mg_res_ty $ HsExpr.mg_ext grp)
#else
#if __GLASGOW_HASKELL__ >= 708
getTypeLHsBind _ (GHC.L spn GHC.FunBind{GHC.fun_matches = grp}) = return $ Just (spn, HsExpr.mg_res_ty grp)
#else
getTypeLHsBind _ (GHC.L spn GHC.FunBind{GHC.fun_matches = GHC.MatchGroup _ typ}) = return $ Just (spn, typ)
#endif
#endif
getTypeLHsBind _ _ = return Nothing

getTypeLHsExpr :: GHC.TypecheckedModule -> GHC.LHsExpr TypecheckI -> GHC.Ghc (Maybe (GHC.SrcSpan, GHC.Type))
#if __GLASGOW_HASKELL__ >= 708
getTypeLHsExpr _ e = do
#else
getTypeLHsExpr tcm e = do
#endif
    hs_env <- GHC.getSession
#if __GLASGOW_HASKELL__ >= 708
    (_, mbe) <- liftIO $ Desugar.deSugarExpr hs_env e
#else
    let modu   = GHC.ms_mod $ GHC.pm_mod_summary $ GHC.tm_parsed_module tcm
        rn_env = TcRnTypes.tcg_rdr_env $ fst $ GHC.tm_internals_ tcm
        ty_env = TcRnTypes.tcg_type_env $ fst $ GHC.tm_internals_ tcm
    (_, mbe) <- liftIO $ Desugar.deSugarExpr hs_env modu rn_env ty_env e
#endif
    return ()
    case mbe of
        Nothing -> return Nothing
        Just expr -> return $ Just (GHC.getLoc e, CoreUtils.exprType expr)

getTypeLPat :: GHC.TypecheckedModule -> GHC.LPat TypecheckI -> GHC.Ghc (Maybe (GHC.SrcSpan, GHC.Type))
getTypeLPat _ (GHC.L spn pat) = return $ Just (spn, TcHsSyn.hsPatType pat)

listifySpans :: Typeable a => GHC.TypecheckedSource -> (Int, Int) -> [GHC.Located a]
listifySpans tcs lc = listifyStaged TypeChecker p tcs
    where
    p (GHC.L spn _) = GHC.isGoodSrcSpan spn && spn `GHC.spans` lc

listifyStaged :: Typeable r => Stage -> (r -> Bool) -> GenericQ [r]
listifyStaged s p = everythingStaged s (++) [] ([] `mkQ` (\x -> [x | p x]))

#if __GLASGOW_HASKELL__ >= 706
pretty :: GHC.DynFlags -> GHC.Type -> String
pretty dflags =
#else
pretty :: GHC.Type -> String
pretty =
#endif
#if __GLASGOW_HASKELL__ >= 800
    Pretty.renderStyle
      Pretty.style{ Pretty.lineLength = 0, Pretty.mode = Pretty.OneLineMode }
#elif __GLASGOW_HASKELL__ >= 708
    Pretty.showDoc Pretty.OneLineMode 0
#else
    Pretty.showDocWith Pretty.OneLineMode
#endif
#if __GLASGOW_HASKELL__ >= 706
    . Outputable.withPprStyleDoc dflags
#else
    . Outputable.withPprStyleDoc
#endif
#if __GLASGOW_HASKELL__ >= 802
        (Outputable.mkUserStyle dflags Outputable.neverQualify Outputable.AllTheWay)
#else
        (Outputable.mkUserStyle Outputable.neverQualify Outputable.AllTheWay)
#endif
#if __GLASGOW_HASKELL__ >= 708
    . PprTyThing.pprTypeForUser
#else
    . PprTyThing.pprTypeForUser False
#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.
everythingStaged :: Stage -> (r -> r -> r) -> r -> GenericQ r -> GenericQ r
everythingStaged stage k z f x
#if __GLASGOW_HASKELL__ >= 806
-- This is a hack, ghc 8.6 changed representation from PostTc
-- to a whole bunch of individial types and I don't really want
-- to handle all of them, at least for the moment since I'm not using
-- this functionality
  | (const False `extQ` fixity `extQ` nameSet) x = z
#else
  | (const False `extQ` postTcType `extQ` fixity `extQ` nameSet) x = z
#endif
  | otherwise = foldl k (f x) (gmapQ (everythingStaged stage k z f) x)
  where nameSet    = const (stage `elem` [Parser,TypeChecker]) :: NameSet.NameSet -> Bool
#if __GLASGOW_HASKELL__ >= 806
        -- there's no more "simple" PostTc type in ghc 8.6
#elif __GLASGOW_HASKELL__ >= 709
        postTcType = const (stage<TypeChecker)                 :: GHC.PostTc TypecheckI GHC.Type -> Bool
#else
        postTcType = const (stage<TypeChecker)                 :: GHC.PostTcType -> Bool
#endif
        fixity     = const (stage<Renamer)                     :: GHC.Fixity -> Bool

------------------------------------------------------------------------------
-- The following code was taken from GHC's ghc/InteractiveUI.hs (with some
-- stylistic changes)

infoThing :: String -> GHC.Ghc String
infoThing str = do
    names <- GHC.parseName str
#if __GLASGOW_HASKELL__ >= 803
    mb_stuffs <- mapM (GHC.getInfo False) names
    let filtered = filterOutChildren (\(t,_f,_i,_,_) -> t) (catMaybes mb_stuffs)
#elif __GLASGOW_HASKELL__ >= 708
    mb_stuffs <- mapM (GHC.getInfo False) names
    let filtered = filterOutChildren (\(t,_f,_i,_) -> t) (catMaybes mb_stuffs)
#else
    mb_stuffs <- mapM GHC.getInfo names
    let filtered = filterOutChildren (\(t,_f,_i) -> t) (catMaybes mb_stuffs)
#endif
    unqual <- GHC.getPrintUnqual
#if __GLASGOW_HASKELL__ >= 706
    dflags <- DynFlags.getDynFlags
    return $ Outputable.showSDocForUser dflags unqual $
#else
    return $ Outputable.showSDocForUser unqual $
#endif
#if __GLASGOW_HASKELL__ >= 708
        Outputable.vcat (intersperse (Outputable.text "") $ map pprInfo filtered)
#else
        Outputable.vcat (intersperse (Outputable.text "") $ map (pprInfo False) filtered)
#endif

  -- Filter out names whose parent is also there Good
  -- example is '[]', which is both a type and data
  -- constructor in the same type
filterOutChildren :: (a -> HscTypes.TyThing) -> [a] -> [a]
filterOutChildren get_thing xs
  = filter (not . has_parent) xs
  where
    all_names = NameSet.mkNameSet (map (GHC.getName . get_thing) xs)
#if __GLASGOW_HASKELL__ >= 704
    has_parent x = case HscTypes.tyThingParent_maybe (get_thing x) of
#else
    has_parent x = case PprTyThing.pprTyThingParent_maybe (get_thing x) of
#endif
                     Just p  -> GHC.getName p `NameSet.elemNameSet` all_names
                     Nothing -> False

#if __GLASGOW_HASKELL__ >= 803
pprInfo :: (HscTypes.TyThing, GHC.Fixity, [GHC.ClsInst], [GHC.FamInst], Outputable.SDoc) -> Outputable.SDoc
pprInfo (thing, fixity, insts, _, _) =
    PprTyThing.pprTyThingInContextLoc thing
#elif __GLASGOW_HASKELL__ >= 708
pprInfo :: (HscTypes.TyThing, GHC.Fixity, [GHC.ClsInst], [GHC.FamInst]) -> Outputable.SDoc
pprInfo (thing, fixity, insts, _) =
    PprTyThing.pprTyThingInContextLoc thing
#elif __GLASGOW_HASKELL__ >= 706
pprInfo :: PprTyThing.PrintExplicitForalls -> (HscTypes.TyThing, GHC.Fixity, [GHC.ClsInst]) -> Outputable.SDoc
pprInfo pefas (thing, fixity, insts) =
    PprTyThing.pprTyThingInContextLoc pefas thing
#else
pprInfo :: PprTyThing.PrintExplicitForalls -> (HscTypes.TyThing, GHC.Fixity, [GHC.Instance]) -> Outputable.SDoc
pprInfo pefas (thing, fixity, insts) =
    PprTyThing.pprTyThingInContextLoc pefas thing
#endif
        Outputable.$$ show_fixity fixity
        Outputable.$$ Outputable.vcat (map GHC.pprInstance insts)
    where
    show_fixity fix
        | fix == GHC.defaultFixity = Outputable.empty
        | otherwise                = Outputable.ppr fix Outputable.<+> Outputable.ppr (GHC.getName thing)