content: module NixTree.PathStats
    ( PathStats (..),
      calculatePathStats,
      whyDepends,
      shortestPathTo,
      module NixTree.StorePath,
    )
  where
  
  import Data.List (minimumBy)
  import qualified Data.List.NonEmpty as NE
  import qualified Data.Map.Lazy as M
  import qualified Data.Set as S
  import NixTree.StorePath
  
  data IntermediatePathStats = IntermediatePathStats
    { ipsAllRefs :: M.Map StoreName (StorePath StoreName ())
    }
  
  data PathStats = PathStats
    { psTotalSize :: !Int,
      psAddedSize :: !Int,
      psImmediateParents :: [StoreName],
      psDisambiguationChars :: !Int
    }
    deriving (Show, Generic, NFData)
  
  mkIntermediateEnv ::
    (StoreName -> Bool) ->
    StoreEnv () ->
    StoreEnv IntermediatePathStats
  mkIntermediateEnv env =
    seBottomUp $ \curr ->
      IntermediatePathStats
        { ipsAllRefs =
            M.unions
              ( M.fromList
                  [ (spName, void sp)
                    | sp@StorePath {spName} <- spRefs curr,
                      env spName
                  ]
                  : map (ipsAllRefs . spPayload) (spRefs curr)
              )
        }
  
  mkFinalEnv :: StoreEnv IntermediatePathStats -> StoreEnv PathStats
  mkFinalEnv env =
    let totalSize = calculateEnvSize env
        immediateParents = calculateImmediateParents (sePaths env)
        disambiguationChars = seDisambiguationChars env
     in flip seBottomUp env $ \StorePath {spName, spSize, spPayload} ->
          let filteredSize =
                seFetchRefs env (/= spName) (seRoots env)
                  & calculateRefsSize
              addedSize = totalSize - filteredSize
           in PathStats
                { psTotalSize =
                    spSize
                      + calculateRefsSize (ipsAllRefs spPayload),
                  psAddedSize = addedSize,
                  psImmediateParents =
                    maybe [] S.toList $ M.lookup spName immediateParents,
                  psDisambiguationChars =
                    M.lookup spName disambiguationChars
                      & maybe 0 id
                }
    where
      calculateEnvSize :: StoreEnv IntermediatePathStats -> Int
      calculateEnvSize e =
        seGetRoots e
          & toList
          & map
            ( \sp@StorePath {spName, spPayload} ->
                M.insert
                  spName
                  (void sp)
                  (ipsAllRefs spPayload)
            )
          & M.unions
          & calculateRefsSize
      calculateRefsSize :: (Functor f, Foldable f) => f (StorePath a b) -> Int
      calculateRefsSize = sum . fmap spSize
      calculateImmediateParents ::
        (Foldable f) =>
        f (StorePath StoreName b) ->
        M.Map StoreName (S.Set StoreName)
      calculateImmediateParents =
        foldl'
          ( \m StorePath {spName, spRefs} ->
              M.unionWith
                (<>)
                m
                (M.fromList (map (,S.singleton spName) spRefs))
          )
          M.empty
  
      seShortNames :: StoreEnv a -> M.Map Text [StoreName]
      seShortNames env =
        let paths = seAll env & toList
         in foldl'
              ( \m StorePath {spName} ->
                  let (_, shortName) = storeNameToSplitShortText spName
                   in M.alter
                        ( \case
                            Nothing -> Just [spName]
                            Just xs -> Just (spName : xs)
                        )
                        shortName
                        m
              )
              M.empty
              paths
  
      seDisambiguationChars :: StoreEnv a -> M.Map StoreName Int
      seDisambiguationChars env =
        M.toList (seShortNames env)
          & map snd
          & concatMap
            ( \xs ->
                let chrs = disambiguate xs
                 in map (\x -> (x, chrs)) xs
            )
          & M.fromList
  
      disambiguate :: [StoreName] -> Int
      disambiguate xs = go 0
        where
          go n =
            if isGood n
              then n
              else go (n + 2)
  
          isGood n =
            xs
              & map (storeNameToShortTextWithDisambiguation n)
              & allUnique
  
          allUnique xx =
            let unique = S.fromList xx
             in length unique == length xx
  
  calculatePathStats :: StoreEnv () -> StoreEnv PathStats
  calculatePathStats = mkFinalEnv . mkIntermediateEnv (const True)
  
  -- TODO: This can be precomputed.
  shortestPathTo :: StoreEnv a -> StoreName -> NonEmpty (StorePath StoreName a)
  shortestPathTo env name =
    seBottomUp
      ( \curr ->
          let currOut = curr {spRefs = spName <$> spRefs curr}
           in if spName curr == name
                then Just (1 :: Int, currOut :| [])
                else
                  mapMaybe spPayload (spRefs curr)
                    & \case
                      [] -> Nothing
                      xs -> case minimumBy (comparing fst) xs of
                        (c, p) -> Just (c + 1, currOut NE.<| p)
      )
      env
      & seGetRoots
      & fmap spPayload
      & NE.toList
      & catMaybes
      & minimumBy (comparing fst)
      & snd
      & NE.reverse
  
  -- Why depends implementation
  
  -- We iterate the dependency graph bottom up. Every node contains a set of paths which represent
  -- the why-depends output from that node down. The set of paths is represented as a "Treeish" object,
  -- which is a trie-like structure.
  whyDepends :: forall a. StoreEnv a -> StoreName -> [NonEmpty (StorePath StoreName a)]
  whyDepends env name =
    seBottomUp @_ @(Maybe (Treeish (StorePath StoreName a)))
      ( \curr ->
          if spName curr == name
            then Just $ mkTreeish (curr {spRefs = map spName (spRefs curr)}) []
            else
              mapMaybe spPayload (spRefs curr)
                & NE.nonEmpty
                <&> NE.toList
                <&> mkTreeish (curr {spRefs = map spName (spRefs curr)})
                <&> capTreeish 1_000_000
      )
      env
      & seGetRoots
      & fmap spPayload
      & NE.toList
      & catMaybes
      & take 10000
      & concatMap treeishToList
  
  -- A trie-like structure which also caches the size.
  data Treeish a = Treeish Int a [Treeish a]
  
  mkTreeish :: a -> [Treeish a] -> Treeish a
  mkTreeish a ts = Treeish (1 + sum (map (\(Treeish i _ _) -> i) ts)) a ts
  
  treeishSize :: Treeish a -> Int
  treeishSize (Treeish i _ _) = i
  
  capTreeish :: Int -> Treeish a -> Treeish a
  capTreeish cap (Treeish i a ts)
    | i <= cap = Treeish i a ts
    | otherwise = Treeish cap a (go cap ts)
    where
      go _ [] = []
      go remaining (x : xs) =
        let x' = capTreeish remaining x
            remaining' = remaining - treeishSize x'
         in if remaining > 0
              then x' : go remaining' xs
              else [x']
  
  treeishToList :: Treeish a -> [NonEmpty a]
  treeishToList (Treeish _ a []) = [a :| []]
  treeishToList (Treeish _ a xs) = map (a NE.<|) (concatMap treeishToList xs)
  