{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecursiveDo #-}
{-# LANGUAGE NamedFieldPuns #-}
module GHC.Debug.Types.Graph( -- * Types
                              HeapGraph(..)
                            , HeapGraphEntry(..)
                            , HeapGraphIndex
                            , PapHI
                            , StackHI
                            , SrtHI
                            -- * Building a heap graph
                            , DerefFunction
                            , buildHeapGraph
                            , multiBuildHeapGraph
                            , generalBuildHeapGraph

                            -- * Printing a heap graph
                            , ppHeapGraph
                            , ppClosure

                            -- * Utility
                            , lookupHeapGraph
                            , traverseHeapGraph
                            , updateHeapGraph
                            , heapGraphSize
                            , annotateHeapGraph

                            -- * Reverse Graph
                            , ReverseGraph
                            , mkReverseGraph
                            , reverseEdges
                            )
                            where

import Data.Char
import Data.List (intercalate, foldl', sort, group, sortBy, groupBy)
import Data.Maybe       ( catMaybes )
import Data.Function
import qualified Data.HashMap.Strict as M
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Control.Monad
import Control.Monad.Fix
import Control.Monad.Trans.State
import Control.Monad.Trans.Class
import GHC.Debug.Types.Ptr
import GHC.Debug.Types.Closures
import qualified Data.List.NonEmpty as NE
import Data.List.NonEmpty (NonEmpty(..))
import Data.Bitraversable

-- | For heap graphs, i.e. data structures that also represent sharing and
-- cyclic structures, these are the entries. If the referenced value is
-- @Nothing@, then we do not have that value in the map, most likely due to
-- exceeding the recursion bound passed to 'buildHeapGraph'.
--
-- Besides a pointer to the stored value and the closure representation we
-- have a slot for arbitrary data, for the user's convenience.
data HeapGraphEntry a = HeapGraphEntry {
        hgeClosurePtr :: ClosurePtr,
        hgeClosure :: DebugClosure SrtHI PapHI ConstrDesc StackHI (Maybe HeapGraphIndex),
        hgeData :: a}
    deriving (Show, Functor, Foldable, Traversable)
type HeapGraphIndex = ClosurePtr

type StackHI = GenStackFrames (GenSrtPayload (Maybe HeapGraphIndex)) (Maybe HeapGraphIndex)
type PapHI =  GenPapPayload (Maybe HeapGraphIndex)
type SrtHI =  GenSrtPayload (Maybe HeapGraphIndex)

-- | The whole graph. The suggested interface is to only use 'lookupHeapGraph',
-- as the internal representation may change. Nevertheless, we export it here:
-- Sometimes the user knows better what he needs than we do.
data HeapGraph a = HeapGraph
                      { roots :: !(NE.NonEmpty ClosurePtr)
                      , graph :: !(IM.IntMap (HeapGraphEntry a)) }
    deriving (Show, Foldable, Traversable, Functor)

traverseHeapGraph :: Applicative m =>
                    (HeapGraphEntry a -> m (HeapGraphEntry b))
                  -> HeapGraph a
                  -> m (HeapGraph b)
traverseHeapGraph f (HeapGraph r im) = HeapGraph r <$> traverse f im


lookupHeapGraph :: HeapGraphIndex -> HeapGraph a -> Maybe (HeapGraphEntry a)
lookupHeapGraph (ClosurePtr i) (HeapGraph _r m) = IM.lookup (fromIntegral i) m

insertHeapGraph :: HeapGraphIndex -> HeapGraphEntry a -> HeapGraph a -> HeapGraph a
insertHeapGraph (ClosurePtr i) a (HeapGraph r m) = HeapGraph r (IM.insert (fromIntegral i) a m)

updateHeapGraph :: (HeapGraphEntry a -> Maybe (HeapGraphEntry a))
                -> HeapGraphIndex
                -> HeapGraph a
                -> HeapGraph a
updateHeapGraph f (ClosurePtr i) (HeapGraph r m) = HeapGraph r (IM.update f (fromIntegral i) m)

heapGraphSize :: HeapGraph a -> Int
heapGraphSize (HeapGraph _ g) = IM.size g

-- | Creates a 'HeapGraph' for the value in the box, but not recursing further
-- than the given limit.
buildHeapGraph
   :: (MonadFix m)
   => DerefFunction m a
   -> Maybe Int
   -> ClosurePtr -- ^ The value to start with
   -> m (HeapGraph a)
buildHeapGraph deref limit initialBox =
  multiBuildHeapGraph deref limit (NE.singleton initialBox)

-- TODO: It is a bit undesirable that the ConstrDesc field is already
-- dereferenced, but also, not such a big deal. It could lead to additional
-- requests to the debuggee which are not necessary and causes a mismatch
-- with the step-by-step decoding functions in `Client.hs`
type DerefFunction m a = ClosurePtr -> m (DebugClosureWithExtra a SrtPayload PapPayload ConstrDesc (GenStackFrames SrtPayload ClosurePtr) ClosurePtr)

-- | Creates a 'HeapGraph' for the values in multiple boxes, but not recursing
--   further than the given limit.
multiBuildHeapGraph
    :: (MonadFix m)
    => DerefFunction m a
    -> Maybe Int
    -> NonEmpty ClosurePtr -- ^ Starting values with associated data entry
    -> m (HeapGraph a)
multiBuildHeapGraph deref limit rs =
  generalBuildHeapGraph deref limit (HeapGraph rs IM.empty) rs
{-# INLINE multiBuildHeapGraph #-}

-- | Adds the given annotation to the entry at the given index, using the
-- 'mappend' operation of its 'Monoid' instance.
annotateHeapGraph ::  (a -> a) -> HeapGraphIndex -> HeapGraph a -> HeapGraph a
annotateHeapGraph f i hg = updateHeapGraph go i hg
  where
    go hge = Just $ hge { hgeData = f (hgeData hge) }

{-# INLINE generalBuildHeapGraph #-}
generalBuildHeapGraph
    :: forall m a .  (MonadFix m)
    => DerefFunction m a
    -> Maybe Int
    -> HeapGraph a
    -> NonEmpty ClosurePtr
    -> m (HeapGraph a)
generalBuildHeapGraph deref limit hg addBoxes = do
    -- First collect all boxes from the existing heap graph
    (_is, hg') <- runStateT (mapM (add limit) addBoxes) hg
    return hg'
  where
    add :: Maybe Int -> ClosurePtr -> StateT (HeapGraph a) m (Maybe ClosurePtr)
    add (Just 0) _ = return Nothing
    add n cp = do
        -- If the box is in the map, return the index
        hm <- get
        case lookupHeapGraph cp hm of
            Just {} -> return (Just cp)
            -- FIXME GHC BUG: change `mdo` to `do` below:
            --       "GHC internal error: ‘c’ is not in scope during type checking, but it passed the renamer"
            Nothing -> mdo
                -- Look up the closure
                c <- lift $ deref cp
                let new_add = add (subtract 1 <$> n)
                -- NOTE: We tie-the-knot here with RecursiveDo so that we don't
                -- get into an infinite loop with cycles in the heap.
                rec modify' (insertHeapGraph cp (HeapGraphEntry cp c' e))
                    -- Add the resulting closure below to the map (above):
                    DCS e c' <- quintraverse (traverse new_add) (traverse new_add) pure (bitraverse (traverse new_add) new_add) new_add c
                return (Just cp)

-- | Pretty-prints a HeapGraph. The resulting string contains newlines. Example
-- for @let s = \"Ki\" in (s, s, cycle \"Ho\")@:
--
-- >let x1 = "Ki"
-- >    x6 = C# 'H' : C# 'o' : x6
-- >in (x1,x1,x6)
ppHeapGraph :: (a -> String) -> HeapGraph a -> String
ppHeapGraph printData (HeapGraph (heapGraphRoot :| rs) m) = letWrapper ++ "(" ++ printData (hgeData (iToE heapGraphRoot)) ++ ") " ++ roots
  where
    -- All variables occuring more than once
    bindings = boundMultipleTimes (HeapGraph (heapGraphRoot :| rs) m) [heapGraphRoot]

    roots = unlines [
              "r" ++ show n ++ ":(" ++ printData (hgeData (iToE r)) ++ ") " ++ ppRef 0 (Just r) ++ "\n"
              | (n, r) <- zip [0 :: Int ..] (heapGraphRoot : rs) ]

    letWrapper =
        if null bindings
        then ""
        else "let " ++ intercalate "\n    " (map ppBinding bindings) ++ "\nin "

    bindingLetter i = case hgeClosure (iToE i) of
        ThunkClosure {} -> 't'
        SelectorClosure {} -> 't'
        APClosure {} -> 't'
        PAPClosure {} -> 'f'
        BCOClosure {} -> 't'
        FunClosure {} -> 'f'
        _ -> 'x'

    ppBindingMap = M.fromList $
        concatMap (zipWith (\j (i,c) -> (i, c : show j)) [(1::Int)..]) $
        groupBy ((==) `on` snd) $
        sortBy (compare `on` snd)
        [ (i, bindingLetter i) | i <- bindings ]

    ppVar i = ppBindingMap M.! i
    ppBinding i = ppVar i ++ "(" ++ printData (hgeData (iToE i)) ++  ") = " ++ ppEntry 0 (iToE i)

    ppEntry prec hge
        | Just s <- isString (hgeClosure hge) = show s
        | Just l <- isList (hgeClosure hge)   = "[" ++ intercalate "," (map (ppRef 0) l) ++ "]"
        | otherwise = ppClosure ppRef prec (hgeClosure hge)
      where
        _app [a] = a  ++ "()"
        _app xs = addBraces (10 <= prec) (unwords xs)

    ppRef _ Nothing = "..."
    ppRef prec (Just i) | i `elem` bindings = ppVar i
                        | otherwise = ppEntry prec (iToE i)
    iToE (ClosurePtr i) = m IM.! (fromIntegral i)

    iToUnboundE cp@(ClosurePtr i)
        | cp `elem` bindings = Nothing
        | otherwise         = IM.lookup (fromIntegral i) m

    isList :: DebugClosure srt p ConstrDesc s (Maybe HeapGraphIndex) -> Maybe [Maybe HeapGraphIndex]
    isList c
        | isNil c =
            return []
        | otherwise = do
            (h,t) <- isCons c
            ti <- t
            e <- iToUnboundE ti
            t' <- isList (hgeClosure e)
            return $ (:) h t'

    isString :: DebugClosure srt p ConstrDesc s (Maybe HeapGraphIndex) -> Maybe String
    isString e = do
        list <- isList e
        -- We do not want to print empty lists as "" as we do not know that they
        -- are really strings.
        if null list
        then Nothing
        else mapM (isChar . hgeClosure <=< iToUnboundE <=< id) list


-- | In the given HeapMap, list all indices that are used more than once. The
-- second parameter adds external references, commonly @[heapGraphRoot]@.
boundMultipleTimes :: HeapGraph a -> [HeapGraphIndex] -> [HeapGraphIndex]
boundMultipleTimes (HeapGraph _rs m) roots = map head $ filter (not.null) $ group $ sort $
     roots ++ concatMap (catMaybes . allClosures . hgeClosure) (IM.elems m)

-- Utilities

addBraces :: Bool -> String -> String
addBraces True t = "(" ++ t ++ ")"
addBraces False t = t

braceize :: [String] -> String
braceize [] = ""
braceize xs = "{" ++ intercalate "," xs ++ "}"

isChar :: DebugClosure srt p ConstrDesc s c -> Maybe Char
isChar ConstrClosure{ constrDesc = ConstrDesc {pkg = "ghc-prim", modl = "GHC.Types", name = "C#"}, dataArgs = [ch], ptrArgs = []} = Just (chr (fromIntegral ch))
isChar _ = Nothing

isNil :: DebugClosure srt p ConstrDesc s c -> Bool
isNil ConstrClosure{ constrDesc = ConstrDesc {pkg = "ghc-prim", modl = "GHC.Types", name = "[]"}, dataArgs = _, ptrArgs = []} = True
isNil _ = False

isCons :: DebugClosure srt p ConstrDesc s c -> Maybe (c, c)
isCons ConstrClosure{ constrDesc = ConstrDesc {pkg = "ghc-prim", modl = "GHC.Types", name = ":"}, dataArgs = [], ptrArgs = [h,t]} = Just (h,t)
isCons _ = Nothing

isTup :: DebugClosure srt p ConstrDesc s c -> Maybe [c]
isTup ConstrClosure{ dataArgs = [], ..} =
    if length (name constrDesc) >= 3 &&
       head (name constrDesc) == '(' && last (name constrDesc) == ')' &&
       all (==',') (tail (init (name constrDesc)))
    then Just ptrArgs else Nothing
isTup _ = Nothing



-- | A pretty-printer that tries to generate valid Haskell for evalutated data.
-- It assumes that for the included boxes, you already replaced them by Strings
-- using 'Data.Foldable.map' or, if you need to do IO, 'Data.Foldable.mapM'.
--
-- The parameter gives the precedendence, to avoid avoidable parenthesises.
ppClosure :: (Int -> c -> String) -> Int -> DebugClosure (GenSrtPayload c) p ConstrDesc s c -> String
ppClosure showBox prec c = case c of
    _ | Just ch <- isChar c -> app
        ["C#", show ch]
    _ | Just (h,t) <- isCons c -> addBraces (5 <= prec) $
        showBox 5 h ++ " : " ++ showBox 4 t
    _ | Just vs <- isTup c ->
        "(" ++ intercalate "," (map (showBox 0) vs) ++ ")"
    ConstrClosure {..} -> app $
        name constrDesc : map (showBox 10) ptrArgs ++ map show dataArgs
    ThunkClosure {..} -> app $
        let srt_string = case getSrt srt of
                            Nothing -> []
                            Just s  -> ["{", showBox 10 s, "}"]
        in ["_thunk" ++ unwords srt_string] ++ map (showBox 10) ptrArgs ++ map show dataArgs

    SelectorClosure {..} -> app
        ["_sel", showBox 10 selectee]
    IndClosure {..} -> app
        ["_ind", showBox 10 indirectee]
    BlackholeClosure {..} -> app
        ["_bh",  showBox 10 indirectee]
    APClosure {..} -> app $ map (showBox 10) $
        [fun]
        -- TODO: Payload
    PAPClosure {..} -> app $ map (showBox 10) $
        [fun] -- TODO payload
    APStackClosure {..} -> app $ map (showBox 10) $
        [fun] -- TODO: stack
    TRecChunkClosure {} -> "_trecChunk" --TODO
    BCOClosure {..} -> app
        ["_bco", showBox 10 bcoptrs]
    ArrWordsClosure {..} -> app
        ["ARR_WORDS", "("++show bytes ++ " bytes)", ((show $ arrWordsBS arrWords)) ]
    MutArrClosure {..} -> app
        --["toMutArray", "("++show (length mccPayload) ++ " ptrs)",  intercalate "," (shorten (map (showBox 10) mccPayload))]
        ["[", intercalate ", " (shorten (map (showBox 10) mccPayload)),"]"]
    SmallMutArrClosure {..} -> app
        ["[", intercalate ", " (shorten (map (showBox 10) mccPayload)),"]"]
    MutVarClosure {..} -> app
        ["_mutVar", showBox 10 var]
    MVarClosure {..} -> app
        ["MVar", showBox 10 value]
    FunClosure {..} ->
        let srt_string = case getSrt srt of
                            Nothing -> []
                            Just s  -> ["{", showBox 10 s, "}"]
        in "_fun" ++ (unwords srt_string) ++  braceize (map (showBox 0) ptrArgs ++ map show dataArgs)
    BlockingQueueClosure {} ->
        "_blockingQueue"
    OtherClosure {} ->
        "_other"
    TSOClosure {} -> "TSO"
    StackClosure {..} -> app ["Stack(", show stack_size, ")"] -- TODO
    WeakClosure {} -> "_wk" -- TODO
    TVarClosure {} -> "_tvar" -- TODO
    MutPrimClosure {} -> "_mutPrim" -- TODO
    UnsupportedClosure {info} -> (show info)


  where
    app [a] = a  ++ "()"
    app xs = addBraces (10 <= prec) (unwords xs)

    shorten xs = if length xs > 20 then take 20 xs ++ ["(and more)"] else xs


-- Reverse Edges
--
closurePtrToInt :: ClosurePtr -> Int
closurePtrToInt (ClosurePtr p) = fromIntegral p

intToClosurePtr :: Int -> ClosurePtr
intToClosurePtr i = mkClosurePtr (fromIntegral i)

newtype ReverseGraph = ReverseGraph (IM.IntMap IS.IntSet)

reverseEdges :: ClosurePtr -> ReverseGraph -> Maybe [ClosurePtr]
reverseEdges cp (ReverseGraph rg) =
  map intToClosurePtr . IS.toList <$> IM.lookup (closurePtrToInt cp) rg

mkReverseGraph :: HeapGraph a -> ReverseGraph
mkReverseGraph (HeapGraph _ hg) = ReverseGraph graph
  where
    graph = IM.foldlWithKey' collectNodes IM.empty hg
    collectNodes newMap k h =
      let bs = allClosures (hgeClosure h)
      in foldl' (\m ma ->
                    case ma of
                      Nothing -> m
                      Just a -> IM.insertWith IS.union (closurePtrToInt a) (IS.singleton k) m) newMap bs