{-# 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
                            -- * 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(..))

-- | 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 PapHI ConstrDesc StackHI (Maybe HeapGraphIndex),
        hgeData :: a}
    deriving (Show, Functor, Foldable, Traversable)
type HeapGraphIndex = ClosurePtr

type StackHI = GenStackFrames (Maybe HeapGraphIndex)
type PapHI =  GenPapPayload (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 PapPayload ConstrDesc StackFrames 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' <- quadtraverse (traverse new_add) pure (traverse 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 (printData (hgeData hge)) 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 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 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 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 p ConstrDesc s c -> Bool
isNil ConstrClosure{ constrDesc = ConstrDesc {pkg = "ghc-prim", modl = "GHC.Types", name = "[]"}, dataArgs = _, ptrArgs = []} = True
isNil _ = False

isCons :: DebugClosure 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 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 :: String -> (Int -> c -> String) -> Int -> DebugClosure p ConstrDesc s c -> String
ppClosure herald 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 $
        "_thunk(" : herald : ")" : 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 {..} ->
        "_fun" ++ 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