content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  module Calligraphy.Phases.Parse
    ( parseHieFiles,
      ppParseError,
      ppParsePhaseDebugInfo,
      ParseError (..),
      ParsePhaseDebugInfo (..),
    )
  where
  
  import qualified Calligraphy.Compat.GHC as GHC
  import Calligraphy.Compat.Lib (isDerivingNode, isInlineNode, isInstanceNode, isMinimalNode, isTypeSignatureNode, mergeSpans, sourceInfo)
  import qualified Calligraphy.Compat.Lib as GHC
  import Calligraphy.Prelude hiding (Decl, DeclType)
  import Calligraphy.Util.LexTree (LexTree, TreeError (..), foldLexTree)
  import qualified Calligraphy.Util.LexTree as LT
  import Calligraphy.Util.Printer
  import Calligraphy.Util.Types (CallGraph (..), Decl (..), DeclType (..), GHCKey (..), Key (..), Loc (..), Module (..), forT_, forestT, over, rekeyCalls)
  import Control.Monad.Except
  import Control.Monad.State
  import Data.Array (Array)
  import qualified Data.Array as Array
  import Data.EnumMap (EnumMap)
  import qualified Data.EnumMap as EnumMap
  import Data.EnumSet (EnumSet)
  import qualified Data.EnumSet as EnumSet
  import qualified Data.Foldable as Foldable
  import Data.List (unzip4)
  import Data.Map (Map)
  import qualified Data.Map as M
  import qualified Data.Map as Map
  import Data.Semigroup
  import Data.Set (Set)
  import qualified Data.Set as Set
  import Data.Tree (Forest)
  import qualified Data.Tree as Tree
  
  -- | A declaration extracted from the source code.
  --
  -- A single symbol can apparently declare a name multiple times in the same place, with multiple distinct keys D:
  -- This happens, for example, with default methods; the name refers to both the method name and the default implementation's name.
  -- We have to account for that to _some_ degree, which is why keys is a set.
  -- The actual resolution of these happens wit 'dedup' in mkForest
  data RawDecl = RawDecl
    { _rdName :: !String,
      rdKeys :: !(EnumSet GHCKey),
      _rdTyp :: !DeclType,
      rdStart :: !Loc,
      rdEnd :: !Loc
    }
  
  data ParseError = TreeError
    { _peModuleName :: String,
      _peModulePath :: FilePath,
      _peError :: TreeError Loc RawDecl
    }
  
  ppParseError :: Prints ParseError
  ppParseError (TreeError str path err) = do
    strLn $ "Parse error in module " <> str <> " (" <> path <> ")"
    indent $ ppTreeError err
    where
      ppTreeError :: Prints (TreeError Loc RawDecl)
      ppTreeError (InvalidBounds l decl r) = do
        strLn $ "Invalid bounds " <> show (l, r) <> "while inserting"
        indent $ ppRawDecl decl
      ppTreeError (OverlappingBounds a b l r) = do
        strLn $ "Clashing bounds: (" <> show (l, r) <> ")"
        strLn "Node 1:"
        indent $ ppRawDecl a
        strLn "Node 2:"
        indent $ ppRawDecl b
      ppTreeError MidSplit = strLn "MidSplit"
      ppTreeError (LexicalError l decl r t) = do
        strLn "Lexical error while inserting"
        strLn "Node:"
        indent $ ppRawDecl decl
        strLn "Bounds:"
        indent $ showLn (l, r)
        strLn "Tree:"
        indent $ do
          ppLexTree t
  
  ppRawDecl :: Prints RawDecl
  ppRawDecl (RawDecl name keys typ st end) = do
    strLn name
    indent $ do
      strLn $ "Type: " <> show typ
      strLn $ "Span: " <> show (st, end)
      strLn $ "Keys: " <> unwords (show <$> EnumSet.toList keys)
  
  ppLexTree :: Prints (LexTree Loc RawDecl)
  ppLexTree = foldLexTree (pure ()) $ \ls l decl m r rs -> do
    ls
    showLn (l, r)
    ppRawDecl decl
    indent m
    rs
  
  ghcNameKey :: GHC.Name -> GHCKey
  ghcNameKey = GHCKey . fromIntegral . GHC.getKey . GHC.nameUnique
  
  newtype ParsePhaseDebugInfo = ParsePhaseDebugInfo {modulesLexTrees :: [(String, LexTree Loc RawDecl)]}
  
  ppParsePhaseDebugInfo :: Prints ParsePhaseDebugInfo
  ppParsePhaseDebugInfo (ParsePhaseDebugInfo mods) = Foldable.forM_ mods $ \(modName, ltree) -> do
    strLn modName
    indent $ ppLexTree ltree
  
  data ParsedFile = ParsedFile
    { _pfModuleName :: String,
      _pfFilePath :: FilePath,
      _pfDecls :: Forest Decl,
      _pfCalls :: Set (GHCKey, GHCKey),
      _pfTypings :: EnumMap GHCKey (EnumSet GHCKey),
      _pfDebugTree :: LexTree Loc RawDecl
    }
  
  -- | Assigns and maintains a mapping of GHCKeys to Key
  type HieParse a = StateT (Key, EnumMap GHCKey Key) (Either ParseError) a
  
  parseHieFiles ::
    [GHC.HieFile] ->
    Either ParseError (ParsePhaseDebugInfo, CallGraph)
  parseHieFiles files = (\(parsed, (_, keymap)) -> mkCallGraph parsed keymap) <$> runStateT (mapM parseHieFile files) (Key 0, mempty)
    where
      mkCallGraph :: [ParsedFile] -> EnumMap GHCKey Key -> (ParsePhaseDebugInfo, CallGraph)
      mkCallGraph parsed keymap =
        let (mods, debugs, calls, typings) = unzip4 (fmap (\(ParsedFile name path decls call typing ltree) -> (Module name path decls, (name, ltree), call, typing)) parsed)
            typeEdges = rekeyCalls keymap . Set.fromList $ do
              (term, types) <- EnumMap.toList (mconcat typings)
              typ <- EnumSet.toList types
              pure (term, typ)
         in (ParsePhaseDebugInfo debugs, CallGraph mods (rekeyCalls keymap (mconcat calls)) typeEdges)
  
  parseHieFile :: GHC.HieFile -> HieParse ParsedFile
  parseHieFile file@(GHC.HieFile filepath mdl _ _ avails _) = do
    lextree <- either (throwError . TreeError modname filepath) pure $ structure decls
    let calls = resolveCalls (fmap (EnumSet.findMin . rdKeys) lextree)
    forest <- forestT (mkDecl exportKeys) (mkForest lextree)
    pure $ ParsedFile modname filepath forest calls types lextree
    where
      modname = GHC.moduleNameString (GHC.moduleName mdl)
      exportKeys = EnumSet.fromList $ fmap ghcNameKey $ avails >>= GHC.availNames
      Collect decls useSites types = collect file
  
      resolveCalls :: LexTree Loc GHCKey -> Set (GHCKey, GHCKey)
      resolveCalls lextree = flip foldMap useSites $ \(loc, callee) ->
        case LT.lookup loc lextree of
          Nothing -> mempty
          Just rep -> Set.singleton (rep, callee)
  
      mkForest :: LexTree Loc RawDecl -> Forest RawDecl
      mkForest = over forestT fromKV . dedup . over forestT toKV . LT.toForest
        where
          toKV (_, RawDecl name keys typ s e, _) = (name, (keys, Max typ, First s, First e))
          fromKV (name, (keys, Max typ, First s, First e)) = RawDecl name keys typ s e
  
      -- TODO this is the only part that touches the state, maybe it's worth lifting it out
      mkDecl :: EnumSet GHCKey -> RawDecl -> HieParse Decl
      mkDecl exportSet (RawDecl str ghcKeys typ start _) = do
        key <- fresh
        Foldable.forM_ (EnumSet.toList ghcKeys) (assoc key)
        let exported = any (flip EnumSet.member exportSet) (EnumSet.toList ghcKeys)
        pure $ Decl str key ghcKeys exported typ start
  
      fresh :: HieParse Key
      fresh = state $ \(Key n, m) -> (Key n, (Key (n + 1), m))
  
      assoc :: Key -> GHCKey -> HieParse ()
      assoc key ghckey = modify $ fmap (EnumMap.insert ghckey key)
  
  dedup :: (Ord k, Semigroup v) => Forest (k, v) -> Forest (k, v)
  dedup = fromDedup . toDedup
    where
      fromDedup = fmap (\(k, (v, d)) -> Tree.Node (k, v) (fromDedup d)) . Map.toList . unDedup
      toDedup = Dedup . Map.fromListWith (<>) . fmap (\(Tree.Node (k, v) f) -> (k, (v, toDedup f)))
  
  newtype Dedup k v = Dedup {unDedup :: Map k (v, Dedup k v)}
  
  instance (Ord k, Semigroup v) => Semigroup (Dedup k v) where
    Dedup a <> Dedup b = Dedup (Map.unionWith (<>) a b)
  
  structure :: [RawDecl] -> Either (TreeError Loc RawDecl) (LexTree Loc RawDecl)
  structure = foldM (\ !t decl -> LT.insertWith f (rdStart decl) decl (rdEnd decl) t) LT.emptyLexTree
    where
      f (RawDecl na ka ta sa ea) prev@(RawDecl nb kb tb _ _)
        | ta == tb && na == nb = Just (RawDecl na (ka <> kb) ta sa ea)
        | otherwise = Just prev
  
  -- | This is the best way I can find of checking whether the name was written by a programmer or not.
  -- GHC internally classifies names extensively, but none of those mechanisms seem to allow to distinguish GHC-generated names.
  isGenerated :: GHC.Name -> Bool
  isGenerated = elem '$' . GHC.getOccString
  
  data Collect = Collect
    { _decls :: [RawDecl],
      _uses :: [(Loc, GHCKey)],
      _types :: EnumMap GHCKey (EnumSet GHCKey)
    }
  
  -- | Collect declarations, uses, and types in a HIE file
  collect :: GHC.HieFile -> Collect
  collect (GHC.HieFile _ _ typeArr (GHC.HieASTs asts) _ _) = execState (forT_ traverse asts go) (Collect mempty mempty mempty)
    where
      tellDecl :: GHC.Name -> DeclType -> GHC.RealSrcSpan -> State Collect ()
      tellDecl nm typ spn = modify $ \(Collect decls uses types) -> Collect (decl : decls) uses types
        where
          decl =
            RawDecl
              (GHC.getOccString nm)
              (EnumSet.singleton . ghcNameKey $ nm)
              typ
              (Loc (GHC.srcSpanStartLine spn) (GHC.srcSpanStartCol spn))
              (Loc (GHC.srcSpanEndLine spn) (GHC.srcSpanEndCol spn))
  
      tellUse :: GHC.RealSrcLoc -> GHCKey -> State Collect ()
      tellUse loc key = modify $ \(Collect decls uses types) -> Collect decls ((Loc (GHC.srcLocLine loc) (GHC.srcLocCol loc), key) : uses) types
  
      tellType :: GHC.Name -> GHC.TypeIndex -> State Collect ()
      tellType name ix = modify $ \(Collect decls uses types) -> Collect decls uses (EnumMap.insertWith (<>) (ghcNameKey name) (typeMap EnumMap.! ix) types)
  
      typeMap = resolveTypes typeArr
  
      ignoreNode nodeInfo = any ($ nodeInfo) [isInstanceNode, isTypeSignatureNode, isInlineNode, isMinimalNode, isDerivingNode]
  
      go :: GHC.HieAST GHC.TypeIndex -> State Collect ()
      go node@(GHC.Node _ _ children) =
        forT_ sourceInfo node $ \nodeInfo ->
          unless (ignoreNode nodeInfo) $ do
            Foldable.forM_ (M.toList $ GHC.nodeIdentifiers nodeInfo) $ \case
              (Right name, GHC.IdentifierDetails ty info) | not (isGenerated name) -> do
                mapM_ (tellType name) ty
                case classifyIdentifier info of
                  IdnIgnore -> pure ()
                  IdnUse -> tellUse (GHC.realSrcSpanStart $ GHC.nodeSpan node) (ghcNameKey name)
                  IdnDecl typ sp
                    | GHC.isPointSpan sp -> pure ()
                    | otherwise -> tellDecl name typ sp
              _ -> pure ()
            mapM_ go children
  
  -- TODO This can be faster by storing intermediate restuls, but that has proven tricky to get right.
  resolveTypes :: Array GHC.TypeIndex GHC.HieTypeFlat -> EnumMap GHC.TypeIndex (EnumSet GHCKey)
  resolveTypes typeArray = EnumMap.fromList [(ix, evalState (go ix) mempty) | ix <- Array.indices typeArray]
    where
      keys :: GHC.HieType a -> EnumSet GHCKey
      keys (GHC.HTyConApp (GHC.IfaceTyCon name _) _) = EnumSet.singleton (ghcNameKey name)
      keys (GHC.HForAllTy ((name, _), _) _) = EnumSet.singleton (ghcNameKey name)
      -- These are variables, which we ignore, but it can't hurt
      keys (GHC.HTyVarTy name) = EnumSet.singleton (ghcNameKey name)
      keys _ = mempty
      go :: GHC.TypeIndex -> State (EnumSet GHC.TypeIndex) (EnumSet GHCKey)
      go current =
        gets (EnumSet.member current) >>= \case
          True -> pure mempty
          False -> do
            modify (EnumSet.insert current)
            let ty = typeArray Array.! current
            mappend (keys ty) . mconcat <$> mapM go (Foldable.toList ty)
  
  data IdentifierType
    = IdnDecl !DeclType !GHC.Span
    | IdnUse
    | IdnIgnore
  
  instance Semigroup IdentifierType where
    IdnIgnore <> a = a
    IdnUse <> IdnIgnore = IdnUse
    IdnUse <> a = a
    IdnDecl typ sp <> IdnDecl typ' sp' = IdnDecl (max typ typ') (mergeSpans sp sp')
    IdnDecl typ sp <> _ = IdnDecl typ sp
  
  instance Monoid IdentifierType where mempty = IdnIgnore
  
  classifyIdentifier :: Set GHC.ContextInfo -> IdentifierType
  classifyIdentifier = foldMap classify
    where
      classify :: GHC.ContextInfo -> IdentifierType
      classify (GHC.Decl GHC.DataDec (Just sp)) = IdnDecl DataDecl sp
      classify (GHC.Decl GHC.PatSynDec (Just sp)) = IdnDecl DataDecl sp
      classify (GHC.Decl GHC.FamDec (Just sp)) = IdnDecl DataDecl sp
      classify (GHC.Decl GHC.SynDec (Just sp)) = IdnDecl DataDecl sp
      classify (GHC.Decl GHC.ConDec (Just sp)) = IdnDecl ConDecl sp
      classify (GHC.Decl GHC.ClassDec (Just sp)) = IdnDecl ClassDecl sp
      classify (GHC.ClassTyDecl (Just sp)) = IdnDecl ValueDecl sp
      classify (GHC.ValBind GHC.RegularBind _ (Just sp)) = IdnDecl ValueDecl sp
      classify (GHC.RecField GHC.RecFieldDecl (Just sp)) = IdnDecl RecDecl sp
      -- Use
      classify (GHC.RecField GHC.RecFieldAssign _) = IdnUse
      classify (GHC.RecField GHC.RecFieldOcc _) = IdnUse
      classify GHC.Use = IdnUse
      -- Ignore
      classify _ = IdnIgnore
  