{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

-- | Dependency filtering is removing all nodes that are not part of a certain dependency tree
module Calligraphy.Phases.DependencyFilter
  ( DependencyFilterConfig,
    DependencyFilterError (..),
    ppFilterError,
    dependencyFilter,
    pDependencyFilterConfig,
  )
where

import Calligraphy.Prelude hiding (Decl, DeclType, Node, filter)
import Calligraphy.Util.Optparse (boolFlags)
import Calligraphy.Util.Printer
import Calligraphy.Util.Types
import Control.Monad.State.Strict
import Data.Bifunctor (bimap)
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.NonEmpty (NonEmpty, nonEmpty)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Tree (Tree)
import qualified Data.Tree as Tree
import Data.Tuple (swap)
import Options.Applicative

data DependencyFilterConfig = DependencyFilterConfig
  { _depRoot :: Maybe (NonEmpty String),
    _revDepRoot :: Maybe (NonEmpty String),
    _depDepth :: Maybe Int,
    _followParent :: Bool,
    _followChildren :: Bool,
    _followCalls :: Bool,
    _followTypes :: Bool
  }

pDependencyFilterConfig :: Parser DependencyFilterConfig
pDependencyFilterConfig =
  DependencyFilterConfig
    <$> (fmap nonEmpty . many)
      ( strOption
          ( long "forward-root"
              <> short 'f'
              <> metavar "NAME"
              <> help "Name of a dependency filter root. Specifying a dependency filter root hides everything that's not a (transitive) dependency of a root. The name can be qualified. This argument can be repeated."
          )
      )
    <*> (fmap nonEmpty . many)
      ( strOption
          ( long "reverse-root"
              <> short 'r'
              <> metavar "NAME"
              <> help "Name of a reverse dependency filter root. Specifying a dependency filter root hides everything that's not a reverse (transitive) dependency of a root. The name can be qualified. This argument can be repeated."
          )
      )
    <*> optional (option auto (long "max-depth" <> help "Maximum search depth for transitive dependencies."))
    <*> boolFlags True "follow-parent" "In calculating (transitive) dependencies, follow edges to from a child to its parent." mempty
    <*> boolFlags True "follow-child" "In calculating (transitive) dependencies, follow edges from a parent to its children." mempty
    <*> boolFlags True "follow-value" "In calculating (transitive) dependencies, follow normal edges." mempty
    <*> boolFlags False "follow-type" "In calculating (transitive) dependencies, follow type edges." mempty

newtype DependencyFilterError = UnknownRootName String

ppFilterError :: Prints DependencyFilterError
ppFilterError (UnknownRootName root) = strLn $ "Unknown root name: " <> root

pruneModules :: (Decl -> Bool) -> CallGraph -> CallGraph
pruneModules p (CallGraph modules calls types) = removeDeadCalls $ CallGraph modules' calls types
  where
    modules' = over (traverse . modForest) (>>= go) modules
    go :: Tree Decl -> [Tree Decl]
    go (Tree.Node decl children) = do
      let children' = children >>= go
       in if p decl then pure (Tree.Node decl children') else children'

-- | Remove all calls and typings (i.e. edges) where one end is not present in the graph.
-- This is intended to be used after an operation that may have removed nodes from the graph.
removeDeadCalls :: CallGraph -> CallGraph
removeDeadCalls (CallGraph mods calls types) = CallGraph mods calls' types'
  where
    outputKeys = execState (forT_ (traverse . modDecls) mods (modify . EnumSet.insert . declKey)) mempty
    calls' = Set.filter (\(a, b) -> EnumSet.member a outputKeys && EnumSet.member b outputKeys) calls
    types' = Set.filter (\(a, b) -> EnumSet.member a outputKeys && EnumSet.member b outputKeys) types

dependencyFilter :: DependencyFilterConfig -> CallGraph -> Either DependencyFilterError CallGraph
dependencyFilter (DependencyFilterConfig mfw mbw maxDepth useParent useChild useCalls useTypes) mods@(CallGraph modules calls types) = do
  fwFilter <- forM mfw $ flip mkDepFilter edges
  bwFilter <- forM mbw $ flip mkDepFilter (Set.map swap edges)
  pure $
    let p = case (fwFilter, bwFilter) of
          (Nothing, Nothing) -> const True
          (Just fa, Nothing) -> fa
          (Nothing, Just fb) -> fb
          (Just fa, Just fb) -> \decl -> fa decl || fb decl
     in pruneModules p mods
  where
    names :: Map String (EnumSet Key)
    names = Map.unionsWith mappend (fmap resolveNames modules)
    mkDepFilter :: NonEmpty String -> Set (Key, Key) -> Either DependencyFilterError (Decl -> Bool)
    mkDepFilter rootNames edges = do
      rootKeys <- forM rootNames $ \name -> maybe (Left $ UnknownRootName name) (pure . EnumSet.toList) (Map.lookup name names)
      let ins = transitives maxDepth (mconcat $ Foldable.toList rootKeys) edges
      pure $ \decl -> EnumSet.member (declKey decl) ins

    edges =
      mconcat
        [ if useParent then parentEdges else mempty,
          if useChild then childEdges else mempty,
          if useCalls then calls else mempty,
          if useTypes then types else mempty
        ]

    parentEdges, childEdges :: Set (Key, Key)
    (parentEdges, childEdges) = execState (forT_ (traverse . modForest . traverse) modules go) mempty
      where
        go :: Tree Decl -> State (Set (Key, Key), Set (Key, Key)) ()
        go (Tree.Node parent children) =
          forM_ children $ \childNode@(Tree.Node child _) -> do
            let kParent = declKey parent
                kChild = declKey child
            modify $ bimap (Set.insert (kParent, kChild)) (Set.insert (kChild, kParent))
            go childNode

-- | Create a map of all names, and the keys they correspond to.
-- For every name in the source, this introduces two entries; one naked, and one qualified with the module name.
resolveNames :: Module -> Map String (EnumSet Key)
resolveNames (Module modName _ forest) =
  flip execState mempty $
    flip forestT forest $
      \(Decl name key _ _ _ _) ->
        modify $
          Map.insertWith (<>) (modName <> "." <> name) (EnumSet.singleton key)
            . Map.insertWith (<>) name (EnumSet.singleton key)

transitives :: forall a. Enum a => Maybe Int -> [a] -> Set (a, a) -> EnumSet a
transitives maxDepth roots deps = go 0 mempty (EnumSet.fromList roots)
  where
    go :: Int -> EnumSet a -> EnumSet a -> EnumSet a
    go depth old new
      | EnumSet.null new = old
      | maybe False (< depth) maxDepth = old
      | otherwise =
          let old' = old <> new
              new' = EnumSet.foldr (\a -> maybe id mappend $ EnumMap.lookup a adjacencies) mempty new
           in go (depth + 1) old' (new' EnumSet.\\ old')
    adjacencies :: EnumMap a (EnumSet a)
    adjacencies = foldr (\(from, to) -> EnumMap.insertWith (<>) from (EnumSet.singleton to)) mempty deps