{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveFunctor #-}
module GHC.Iface.Ext.Utils where

import GHC.Prelude

import GHC.Core.Map
import GHC.Driver.Session    ( DynFlags )
import GHC.Data.FastString   ( FastString, mkFastString )
import GHC.Iface.Type
import GHC.Core.Multiplicity
import GHC.Types.Name hiding (varName)
import GHC.Types.Name.Set
import GHC.Utils.Outputable hiding ( (<>) )
import qualified GHC.Utils.Outputable as O
import GHC.Types.SrcLoc
import GHC.CoreToIface
import GHC.Core.TyCon
import GHC.Core.TyCo.Rep
import GHC.Core.Type
import GHC.Types.Var
import GHC.Types.Var.Env

import GHC.Iface.Ext.Types

import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.IntMap.Strict as IM
import qualified Data.Array as A
import Data.Data                  ( typeOf, typeRepTyCon, Data(toConstr) )
import Data.Maybe                 ( maybeToList, mapMaybe)
import Data.Monoid
import Data.List                  (find)
import Data.Traversable           ( for )
import Data.Coerce
import Control.Monad.Trans.State.Strict hiding (get)
import Control.Monad.Trans.Reader
import qualified Data.Tree as Tree

type RefMap a = M.Map Identifier [(Span, IdentifierDetails a)]

generateReferencesMap
  :: Foldable f
  => f (HieAST a)
  -> RefMap a
generateReferencesMap = foldr (\ast m -> M.unionWith (++) (go ast) m) M.empty
  where
    go ast = M.unionsWith (++) (this : map go (nodeChildren ast))
      where
        this = fmap (pure . (nodeSpan ast,)) $ sourcedNodeIdents $ sourcedNodeInfo ast

renderHieType :: DynFlags -> HieTypeFix -> String
renderHieType dflags ht = renderWithStyle (initSDocContext dflags defaultUserStyle) (ppr $ hieTypeToIface ht)

resolveVisibility :: Type -> [Type] -> [(Bool,Type)]
resolveVisibility kind ty_args
  = go (mkEmptyTCvSubst in_scope) kind ty_args
  where
    in_scope = mkInScopeSet (tyCoVarsOfTypes ty_args)

    go _   _                   []     = []
    go env ty                  ts
      | Just ty' <- coreView ty
      = go env ty' ts
    go env (ForAllTy (Bndr tv vis) res) (t:ts)
      | isVisibleArgFlag vis = (True , t) : ts'
      | otherwise            = (False, t) : ts'
      where
        ts' = go (extendTvSubst env tv t) res ts

    go env (FunTy { ft_res = res }) (t:ts) -- No type-class args in tycon apps
      = (True,t) : (go env res ts)

    go env (TyVarTy tv) ts
      | Just ki <- lookupTyVar env tv = go env ki ts
    go env kind (t:ts) = (True, t) : (go env kind ts) -- Ill-kinded

foldType :: (HieType a -> a) -> HieTypeFix -> a
foldType f (Roll t) = f $ fmap (foldType f) t

selectPoint :: HieFile -> (Int,Int) -> Maybe (HieAST Int)
selectPoint hf (sl,sc) = getFirst $
  flip foldMap (M.toList (getAsts $ hie_asts hf)) $ \(fs,ast) -> First $
      case selectSmallestContaining (sp fs) ast of
        Nothing -> Nothing
        Just ast' -> Just ast'
 where
   sloc fs = mkRealSrcLoc fs sl sc
   sp fs = mkRealSrcSpan (sloc fs) (sloc fs)

findEvidenceUse :: NodeIdentifiers a -> [Name]
findEvidenceUse ni = [n | (Right n, dets) <- xs, any isEvidenceUse (identInfo dets)]
 where
   xs = M.toList ni

data EvidenceInfo a
  = EvidenceInfo
  { evidenceVar :: Name
  , evidenceSpan :: RealSrcSpan
  , evidenceType :: a
  , evidenceDetails :: Maybe (EvVarSource, Scope, Maybe Span)
  } deriving (Eq,Ord,Functor)

instance (Outputable a) => Outputable (EvidenceInfo a) where
  ppr (EvidenceInfo name span typ dets) =
    hang (ppr name <+> text "at" <+> ppr span O.<> text ", of type:" <+> ppr typ) 4 $
      pdets $$ (pprDefinedAt name)
    where
      pdets = case dets of
        Nothing -> text "is a usage of an external evidence variable"
        Just (src,scp,spn) -> text "is an" <+> ppr (EvidenceVarBind src scp spn)

getEvidenceTreesAtPoint :: HieFile -> RefMap a -> (Int,Int) -> Tree.Forest (EvidenceInfo a)
getEvidenceTreesAtPoint hf refmap point =
  [t | Just ast <- pure $ selectPoint hf point
     , n        <- findEvidenceUse (sourcedNodeIdents $ sourcedNodeInfo ast)
     , Just t   <- pure $ getEvidenceTree refmap n
     ]

getEvidenceTree :: RefMap a -> Name -> Maybe (Tree.Tree (EvidenceInfo a))
getEvidenceTree refmap var = go emptyNameSet var
  where
    go seen var
      | var `elemNameSet` seen = Nothing
      | otherwise = do
          xs <- M.lookup (Right var) refmap
          case find (any isEvidenceBind . identInfo . snd) xs of
            Just (sp,dets) -> do
              typ <- identType dets
              (evdet,children) <- getFirst $ foldMap First $ do
                 det <- S.toList $ identInfo dets
                 case det of
                   EvidenceVarBind src@(EvLetBind (getEvBindDeps -> xs)) scp spn ->
                     pure $ Just ((src,scp,spn),mapMaybe (go $ extendNameSet seen var) xs)
                   EvidenceVarBind src scp spn -> pure $ Just ((src,scp,spn),[])
                   _ -> pure Nothing
              pure $ Tree.Node (EvidenceInfo var sp typ (Just evdet)) children
            -- It is externally bound
            Nothing -> getFirst $ foldMap First $ do
              (sp,dets) <- xs
              if (any isEvidenceUse $ identInfo dets)
                then do
                  case identType dets of
                    Nothing -> pure Nothing
                    Just typ -> pure $ Just $ Tree.Node (EvidenceInfo var sp typ Nothing) []
                else pure Nothing

hieTypeToIface :: HieTypeFix -> IfaceType
hieTypeToIface = foldType go
  where
    go (HTyVarTy n) = IfaceTyVar $ occNameFS $ getOccName n
    go (HAppTy a b) = IfaceAppTy a (hieToIfaceArgs b)
    go (HLitTy l) = IfaceLitTy l
    go (HForAllTy ((n,k),af) t) = let b = (occNameFS $ getOccName n, k)
                                  in IfaceForAllTy (Bndr (IfaceTvBndr b) af) t
    go (HFunTy w a b)   = IfaceFunTy VisArg   w       a    b
    go (HQualTy pred b) = IfaceFunTy InvisArg many_ty pred b
    go (HCastTy a) = a
    go HCoercionTy = IfaceTyVar "<coercion type>"
    go (HTyConApp a xs) = IfaceTyConApp a (hieToIfaceArgs xs)

    -- This isn't fully faithful - we can't produce the 'Inferred' case
    hieToIfaceArgs :: HieArgs IfaceType -> IfaceAppArgs
    hieToIfaceArgs (HieArgs xs) = go' xs
      where
        go' [] = IA_Nil
        go' ((True ,x):xs) = IA_Arg x Required $ go' xs
        go' ((False,x):xs) = IA_Arg x Specified $ go' xs

data HieTypeState
  = HTS
    { tyMap      :: !(TypeMap TypeIndex)
    , htyTable   :: !(IM.IntMap HieTypeFlat)
    , freshIndex :: !TypeIndex
    }

initialHTS :: HieTypeState
initialHTS = HTS emptyTypeMap IM.empty 0

freshTypeIndex :: State HieTypeState TypeIndex
freshTypeIndex = do
  index <- gets freshIndex
  modify' $ \hts -> hts { freshIndex = index+1 }
  return index

compressTypes
  :: HieASTs Type
  -> (HieASTs TypeIndex, A.Array TypeIndex HieTypeFlat)
compressTypes asts = (a, arr)
  where
    (a, (HTS _ m i)) = flip runState initialHTS $
      for asts $ \typ -> do
        i <- getTypeIndex typ
        return i
    arr = A.array (0,i-1) (IM.toList m)

recoverFullType :: TypeIndex -> A.Array TypeIndex HieTypeFlat -> HieTypeFix
recoverFullType i m = go i
  where
    go i = Roll $ fmap go (m A.! i)

getTypeIndex :: Type -> State HieTypeState TypeIndex
getTypeIndex t
  | otherwise = do
      tm <- gets tyMap
      case lookupTypeMap tm t of
        Just i -> return i
        Nothing -> do
          ht <- go t
          extendHTS t ht
  where
    extendHTS t ht = do
      i <- freshTypeIndex
      modify' $ \(HTS tm tt fi) ->
        HTS (extendTypeMap tm t i) (IM.insert i ht tt) fi
      return i

    go (TyVarTy v) = return $ HTyVarTy $ varName v
    go ty@(AppTy _ _) = do
      let (head,args) = splitAppTys ty
          visArgs = HieArgs $ resolveVisibility (typeKind head) args
      ai <- getTypeIndex head
      argsi <- mapM getTypeIndex visArgs
      return $ HAppTy ai argsi
    go (TyConApp f xs) = do
      let visArgs = HieArgs $ resolveVisibility (tyConKind f) xs
      is <- mapM getTypeIndex visArgs
      return $ HTyConApp (toIfaceTyCon f) is
    go (ForAllTy (Bndr v a) t) = do
      k <- getTypeIndex (varType v)
      i <- getTypeIndex t
      return $ HForAllTy ((varName v,k),a) i
    go (FunTy { ft_af = af, ft_mult = w, ft_arg = a, ft_res = b }) = do
      ai <- getTypeIndex a
      bi <- getTypeIndex b
      wi <- getTypeIndex w
      return $ case af of
                 InvisArg -> case w of Many -> HQualTy ai bi; _ -> error "Unexpected non-unrestricted predicate"
                 VisArg   -> HFunTy wi ai bi
    go (LitTy a) = return $ HLitTy $ toIfaceTyLit a
    go (CastTy t _) = do
      i <- getTypeIndex t
      return $ HCastTy i
    go (CoercionTy _) = return HCoercionTy

resolveTyVarScopes :: M.Map FastString (HieAST a) -> M.Map FastString (HieAST a)
resolveTyVarScopes asts = M.map go asts
  where
    go ast = resolveTyVarScopeLocal ast asts

resolveTyVarScopeLocal :: HieAST a -> M.Map FastString (HieAST a) -> HieAST a
resolveTyVarScopeLocal ast asts = go ast
  where
    resolveNameScope dets = dets{identInfo =
      S.map resolveScope (identInfo dets)}
    resolveScope (TyVarBind sc (UnresolvedScope names Nothing)) =
      TyVarBind sc $ ResolvedScopes
        [ LocalScope binding
        | name <- names
        , Just binding <- [getNameBinding name asts]
        ]
    resolveScope (TyVarBind sc (UnresolvedScope names (Just sp))) =
      TyVarBind sc $ ResolvedScopes
        [ LocalScope binding
        | name <- names
        , Just binding <- [getNameBindingInClass name sp asts]
        ]
    resolveScope scope = scope
    go (Node info span children) = Node info' span $ map go children
      where
        info' = SourcedNodeInfo (updateNodeInfo <$> getSourcedNodeInfo info)
        updateNodeInfo i = i { nodeIdentifiers = idents }
          where
            idents = M.map resolveNameScope $ nodeIdentifiers i

getNameBinding :: Name -> M.Map FastString (HieAST a) -> Maybe Span
getNameBinding n asts = do
  (_,msp) <- getNameScopeAndBinding n asts
  msp

getNameScope :: Name -> M.Map FastString (HieAST a) -> Maybe [Scope]
getNameScope n asts = do
  (scopes,_) <- getNameScopeAndBinding n asts
  return scopes

getNameBindingInClass
  :: Name
  -> Span
  -> M.Map FastString (HieAST a)
  -> Maybe Span
getNameBindingInClass n sp asts = do
  ast <- M.lookup (srcSpanFile sp) asts
  getFirst $ foldMap First $ do
    child <- flattenAst ast
    dets <- maybeToList
      $ M.lookup (Right n) $ sourcedNodeIdents $ sourcedNodeInfo child
    let binding = foldMap (First . getBindSiteFromContext) (identInfo dets)
    return (getFirst binding)

getNameScopeAndBinding
  :: Name
  -> M.Map FastString (HieAST a)
  -> Maybe ([Scope], Maybe Span)
getNameScopeAndBinding n asts = case nameSrcSpan n of
  RealSrcSpan sp _ -> do -- @Maybe
    ast <- M.lookup (srcSpanFile sp) asts
    defNode <- selectLargestContainedBy sp ast
    getFirst $ foldMap First $ do -- @[]
      node <- flattenAst defNode
      dets <- maybeToList
        $ M.lookup (Right n) $ sourcedNodeIdents $ sourcedNodeInfo node
      scopes <- maybeToList $ foldMap getScopeFromContext (identInfo dets)
      let binding = foldMap (First . getBindSiteFromContext) (identInfo dets)
      return $ Just (scopes, getFirst binding)
  _ -> Nothing

getScopeFromContext :: ContextInfo -> Maybe [Scope]
getScopeFromContext (ValBind _ sc _) = Just [sc]
getScopeFromContext (PatternBind a b _) = Just [a, b]
getScopeFromContext (ClassTyDecl _) = Just [ModuleScope]
getScopeFromContext (Decl _ _) = Just [ModuleScope]
getScopeFromContext (TyVarBind a (ResolvedScopes xs)) = Just $ a:xs
getScopeFromContext (TyVarBind a _) = Just [a]
getScopeFromContext (EvidenceVarBind _ a _) = Just [a]
getScopeFromContext _ = Nothing

getBindSiteFromContext :: ContextInfo -> Maybe Span
getBindSiteFromContext (ValBind _ _ sp) = sp
getBindSiteFromContext (PatternBind _ _ sp) = sp
getBindSiteFromContext _ = Nothing

flattenAst :: HieAST a -> [HieAST a]
flattenAst n =
  n : concatMap flattenAst (nodeChildren n)

smallestContainingSatisfying
  :: Span
  -> (HieAST a -> Bool)
  -> HieAST a
  -> Maybe (HieAST a)
smallestContainingSatisfying sp cond node
  | nodeSpan node `containsSpan` sp = getFirst $ mconcat
      [ foldMap (First . smallestContainingSatisfying sp cond) $
          nodeChildren node
      , First $ if cond node then Just node else Nothing
      ]
  | sp `containsSpan` nodeSpan node = Nothing
  | otherwise = Nothing

selectLargestContainedBy :: Span -> HieAST a -> Maybe (HieAST a)
selectLargestContainedBy sp node
  | sp `containsSpan` nodeSpan node = Just node
  | nodeSpan node `containsSpan` sp =
      getFirst $ foldMap (First . selectLargestContainedBy sp) $
        nodeChildren node
  | otherwise = Nothing

selectSmallestContaining :: Span -> HieAST a -> Maybe (HieAST a)
selectSmallestContaining sp node
  | nodeSpan node `containsSpan` sp = getFirst $ mconcat
      [ foldMap (First . selectSmallestContaining sp) $ nodeChildren node
      , First (Just node)
      ]
  | sp `containsSpan` nodeSpan node = Nothing
  | otherwise = Nothing

definedInAsts :: M.Map FastString (HieAST a) -> Name -> Bool
definedInAsts asts n = case nameSrcSpan n of
  RealSrcSpan sp _ -> srcSpanFile sp `elem` M.keys asts
  _ -> False

getEvidenceBindDeps :: ContextInfo -> [Name]
getEvidenceBindDeps (EvidenceVarBind (EvLetBind xs) _ _) =
  getEvBindDeps xs
getEvidenceBindDeps _ = []

isEvidenceBind :: ContextInfo -> Bool
isEvidenceBind EvidenceVarBind{} = True
isEvidenceBind _ = False

isEvidenceContext :: ContextInfo -> Bool
isEvidenceContext EvidenceVarUse = True
isEvidenceContext EvidenceVarBind{} = True
isEvidenceContext _ = False

isEvidenceUse :: ContextInfo -> Bool
isEvidenceUse EvidenceVarUse = True
isEvidenceUse _ = False

isOccurrence :: ContextInfo -> Bool
isOccurrence Use = True
isOccurrence EvidenceVarUse = True
isOccurrence _ = False

scopeContainsSpan :: Scope -> Span -> Bool
scopeContainsSpan NoScope _ = False
scopeContainsSpan ModuleScope _ = True
scopeContainsSpan (LocalScope a) b = a `containsSpan` b

-- | One must contain the other. Leaf nodes cannot contain anything
combineAst :: HieAST Type -> HieAST Type -> HieAST Type
combineAst a@(Node aInf aSpn xs) b@(Node bInf bSpn ys)
  | aSpn == bSpn = Node (aInf `combineSourcedNodeInfo` bInf) aSpn (mergeAsts xs ys)
  | aSpn `containsSpan` bSpn = combineAst b a
combineAst a (Node xs span children) = Node xs span (insertAst a children)

-- | Insert an AST in a sorted list of disjoint Asts
insertAst :: HieAST Type -> [HieAST Type] -> [HieAST Type]
insertAst x = mergeAsts [x]

nodeInfo :: HieAST Type -> NodeInfo Type
nodeInfo = foldl' combineNodeInfo emptyNodeInfo . getSourcedNodeInfo . sourcedNodeInfo

emptyNodeInfo :: NodeInfo a
emptyNodeInfo = NodeInfo S.empty [] M.empty

sourcedNodeIdents :: SourcedNodeInfo a -> NodeIdentifiers a
sourcedNodeIdents = M.unionsWith (<>) . fmap nodeIdentifiers . getSourcedNodeInfo

combineSourcedNodeInfo :: SourcedNodeInfo Type -> SourcedNodeInfo Type -> SourcedNodeInfo Type
combineSourcedNodeInfo = coerce $ M.unionWith combineNodeInfo

-- | Merge two nodes together.
--
-- Precondition and postcondition: elements in 'nodeType' are ordered.
combineNodeInfo :: NodeInfo Type -> NodeInfo Type -> NodeInfo Type
(NodeInfo as ai ad) `combineNodeInfo` (NodeInfo bs bi bd) =
  NodeInfo (S.union as bs) (mergeSorted ai bi) (M.unionWith (<>) ad bd)
  where
    mergeSorted :: [Type] -> [Type] -> [Type]
    mergeSorted la@(a:as) lb@(b:bs) = case nonDetCmpType a b of
                                        LT -> a : mergeSorted as lb
                                        EQ -> a : mergeSorted as bs
                                        GT -> b : mergeSorted la bs
    mergeSorted as [] = as
    mergeSorted [] bs = bs


{- | Merge two sorted, disjoint lists of ASTs, combining when necessary.

In the absence of position-altering pragmas (ex: @# line "file.hs" 3@),
different nodes in an AST tree should either have disjoint spans (in
which case you can say for sure which one comes first) or one span
should be completely contained in the other (in which case the contained
span corresponds to some child node).

However, since Haskell does have position-altering pragmas it /is/
possible for spans to be overlapping. Here is an example of a source file
in which @foozball@ and @quuuuuux@ have overlapping spans:

@
module Baz where

# line 3 "Baz.hs"
foozball :: Int
foozball = 0

# line 3 "Baz.hs"
bar, quuuuuux :: Int
bar = 1
quuuuuux = 2
@

In these cases, we just do our best to produce sensible `HieAST`'s. The blame
should be laid at the feet of whoever wrote the line pragmas in the first place
(usually the C preprocessor...).
-}
mergeAsts :: [HieAST Type] -> [HieAST Type] -> [HieAST Type]
mergeAsts xs [] = xs
mergeAsts [] ys = ys
mergeAsts xs@(a:as) ys@(b:bs)
  | span_a `containsSpan`   span_b = mergeAsts (combineAst a b : as) bs
  | span_b `containsSpan`   span_a = mergeAsts as (combineAst a b : bs)
  | span_a `rightOf`        span_b = b : mergeAsts xs bs
  | span_a `leftOf`         span_b = a : mergeAsts as ys

  -- These cases are to work around ASTs that are not fully disjoint
  | span_a `startsRightOf`  span_b = b : mergeAsts as ys
  | otherwise                      = a : mergeAsts as ys
  where
    span_a = nodeSpan a
    span_b = nodeSpan b

rightOf :: Span -> Span -> Bool
rightOf s1 s2
  = (srcSpanStartLine s1, srcSpanStartCol s1)
       >= (srcSpanEndLine s2, srcSpanEndCol s2)
    && (srcSpanFile s1 == srcSpanFile s2)

leftOf :: Span -> Span -> Bool
leftOf s1 s2
  = (srcSpanEndLine s1, srcSpanEndCol s1)
       <= (srcSpanStartLine s2, srcSpanStartCol s2)
    && (srcSpanFile s1 == srcSpanFile s2)

startsRightOf :: Span -> Span -> Bool
startsRightOf s1 s2
  = (srcSpanStartLine s1, srcSpanStartCol s1)
       >= (srcSpanStartLine s2, srcSpanStartCol s2)

-- | combines and sorts ASTs using a merge sort
mergeSortAsts :: [HieAST Type] -> [HieAST Type]
mergeSortAsts = go . map pure
  where
    go [] = []
    go [xs] = xs
    go xss = go (mergePairs xss)
    mergePairs [] = []
    mergePairs [xs] = [xs]
    mergePairs (xs:ys:xss) = mergeAsts xs ys : mergePairs xss

simpleNodeInfo :: FastString -> FastString -> NodeInfo a
simpleNodeInfo cons typ = NodeInfo (S.singleton (cons, typ)) [] M.empty

locOnly :: Monad m => SrcSpan -> ReaderT NodeOrigin m [HieAST a]
locOnly (RealSrcSpan span _) = do
  org <- ask
  let e = mkSourcedNodeInfo org $ emptyNodeInfo
  pure [Node e span []]
locOnly _ = pure []

mkScope :: SrcSpan -> Scope
mkScope (RealSrcSpan sp _) = LocalScope sp
mkScope _ = NoScope

mkLScope :: Located a -> Scope
mkLScope = mkScope . getLoc

combineScopes :: Scope -> Scope -> Scope
combineScopes ModuleScope _ = ModuleScope
combineScopes _ ModuleScope = ModuleScope
combineScopes NoScope x = x
combineScopes x NoScope = x
combineScopes (LocalScope a) (LocalScope b) =
  mkScope $ combineSrcSpans (RealSrcSpan a Nothing) (RealSrcSpan b Nothing)

mkSourcedNodeInfo :: NodeOrigin -> NodeInfo a -> SourcedNodeInfo a
mkSourcedNodeInfo org ni = SourcedNodeInfo $ M.singleton org ni

{-# INLINEABLE makeNode #-}
makeNode
  :: (Monad m, Data a)
  => a                       -- ^ helps fill in 'nodeAnnotations' (with 'Data')
  -> SrcSpan                 -- ^ return an empty list if this is unhelpful
  -> ReaderT NodeOrigin m [HieAST b]
makeNode x spn = do
  org <- ask
  pure $ case spn of
    RealSrcSpan span _ -> [Node (mkSourcedNodeInfo org $ simpleNodeInfo cons typ) span []]
    _ -> []
  where
    cons = mkFastString . show . toConstr $ x
    typ = mkFastString . show . typeRepTyCon . typeOf $ x

{-# INLINEABLE makeTypeNode #-}
makeTypeNode
  :: (Monad m, Data a)
  => a                       -- ^ helps fill in 'nodeAnnotations' (with 'Data')
  -> SrcSpan                 -- ^ return an empty list if this is unhelpful
  -> Type                    -- ^ type to associate with the node
  -> ReaderT NodeOrigin m [HieAST Type]
makeTypeNode x spn etyp = do
  org <- ask
  pure $ case spn of
    RealSrcSpan span _ ->
      [Node (mkSourcedNodeInfo org $ NodeInfo (S.singleton (cons,typ)) [etyp] M.empty) span []]
    _ -> []
  where
    cons = mkFastString . show . toConstr $ x
    typ = mkFastString . show . typeRepTyCon . typeOf $ x