content: {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE ParallelListComp #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE DerivingVia #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# OPTIONS_GHC -Wno-orphans #-}
  module Lib
    ( -- * Running/Connecting to a debuggee
      Debuggee
    , debuggeeRun
    , debuggeeConnect
    , snapshotConnect
    , debuggeeClose
    , withDebuggeeRun
    , withDebuggeeConnect
    , socketDirectory
    , snapshotDirectory
  
      -- * Pause/Resume
    , GD.pause
    , GD.resume
    , GD.pausePoll
    , GD.withPause
  
      -- * Querying the paused debuggee
    , rootClosures
    , savedClosures
  
      -- * Closures
    , Closure
    , DebugClosure(..)
    , closureShowAddress
    , closureExclusiveSize
    , closureRetainerSize
    , closureSourceLocation
    , SourceInformation(..)
    , closureReferences
    , closurePretty
    , fillConstrDesc
    , InfoTablePtr
    , ListItem(..)
    , closureInfoPtr
    , infoSourceLocation
  
      -- * Common initialisation
    , initialTraversal
    , HG.HeapGraph(..)
      -- * Dominator Tree
    , dominatorRootClosures
    , closureDominatees
    , runAnalysis
    , Analysis(..)
    , Size(..)
    , RetainerSize(..)
      -- * Reverse Edge Map
    , HG.mkReverseGraph
    , reverseClosureReferences
    , lookupHeapGraph
  
      -- * Profiling
    , profile
  
      -- * Retainers
    , retainersOfConstructor
    , retainersOfConstructorExact
  
      -- * Snapshot
    , snapshot
  
    -- * Types
    , Ptr(..)
    , toPtr
    , dereferencePtr
    , ConstrDesc(..)
    , ConstrDescCont
    , GenPapPayload(..)
    , StackCont
    , PayloadCont
    , ClosurePtr
    , HG.StackHI
    , HG.PapHI
    , HG.HeapGraphIndex
      --
    ) where
  
  import           Data.List.NonEmpty (NonEmpty(..))
  import qualified Data.Graph as G
  import           Data.Maybe (fromMaybe, mapMaybe)
  import qualified GHC.Debug.Types as GD
  import           GHC.Debug.Types hiding (Closure, DebugClosure)
  import           GHC.Debug.Convention (socketDirectory, snapshotDirectory)
  import           GHC.Debug.Client.Monad (request, run, Debuggee)
  import qualified GHC.Debug.Client.Monad as GD
  import qualified GHC.Debug.Client.Query as GD
  import qualified GHC.Debug.Profile as GD
  import qualified GHC.Debug.Retainers as GD
  import qualified GHC.Debug.Snapshot as GD
  import qualified GHC.Debug.Types.Graph as HG
  import qualified GHC.Debug.Dominators as HG
  import qualified Data.HashMap.Strict as HM
  import Data.Tree
  import Control.Monad
  import System.FilePath
  import System.Directory
  
  data Analysis = Analysis
    { analysisDominatorRoots :: ![ClosurePtr]
    , analysisDominatees :: !(ClosurePtr -> [ClosurePtr])
    -- ^ Unsorted dominatees of a closure
    , analysisSizes :: !(ClosurePtr -> (Size, RetainerSize))
    -- ^ Size and retainer size (via dominator tree) of closures
    }
  
  initialTraversal :: Debuggee -> IO (HG.HeapGraph Size)
  initialTraversal e = run e $ do
      -- Calculate the dominator tree with retainer sizes
      -- TODO perhaps this conversion to a graph can be done in GHC.Debug.Types.Graph
      _ <- GD.precacheBlocks
      rs <- request RequestRoots
      let derefFuncM cPtr = do
            c <- GD.dereferenceClosure cPtr
            quadtraverse GD.dereferencePapPayload GD.dereferenceConDesc GD.dereferenceStack pure c
      hg <- case rs of
        [] -> error "Empty roots"
        (x:xs) -> HG.multiBuildHeapGraph derefFuncM Nothing (x :| xs)
      return hg
  
  -- This function is very very very slow, it needs to be optimised.
  runAnalysis :: Debuggee -> HG.HeapGraph Size -> IO Analysis
  runAnalysis e hg = run e $ do
      let drs :: [G.Tree (ClosurePtr, (Size, RetainerSize))]
          drs = fmap (\ent -> (HG.hgeClosurePtr ent, HG.hgeData ent)) <$> HG.retainerSize hg
  
          !hmGraph = HM.unions (map snd $ foldTree buildGraphNode <$> (HG.retainerSize hg))
  
          buildGraphNode :: HG.HeapGraphEntry v
                         -> [(ClosurePtr, HM.HashMap ClosurePtr (v, [ClosurePtr]))]
                         -> (ClosurePtr, HM.HashMap ClosurePtr (v, [ClosurePtr]))
          buildGraphNode hge subtrees =
              (cptr, HM.insert cptr v (HM.unions submaps))
            where
              cptr = HG.hgeClosurePtr hge
              v = (HG.hgeData hge, children)
              (children, submaps) = unzip subtrees
  
          cPtrToData
            = fromMaybe ((-12221, RetainerSize (-12221)), [])
            -- ^ TODO I would expect the mapping to be complete unless out analysis misses some closures.
            . flip HM.lookup hmGraph
  
      return $ Analysis
                [drPtr | G.Node (drPtr, _) _ <- drs]
                ((\(_,x) -> x) . cPtrToData)
                ((\(x,_) -> x) . cPtrToData)
  
  -- | Bracketed version of @debuggeeRun@. Runs a debuggee, connects to it, runs
  -- the action, kills the process, then closes the debuggee.
  withDebuggeeRun :: FilePath  -- ^ path to executable to run as the debuggee
                  -> FilePath  -- ^ filename of socket (e.g. @"/tmp/ghc-debug"@)
                  -> (Debuggee -> IO a)
                  -> IO a
  withDebuggeeRun exeName socketName action = GD.withDebuggeeRun exeName socketName action
  
  -- | Bracketed version of @debuggeeConnect@. Connects to a debuggee, runs the
  -- action, then closes the debuggee.
  withDebuggeeConnect :: FilePath  -- ^ filename of socket (e.g. @"/tmp/ghc-debug"@)
                     -> (Debuggee -> IO a)
                     -> IO a
  withDebuggeeConnect socketName action = GD.withDebuggeeConnect socketName action
  
  -- | Run a debuggee and connect to it. Use @debuggeeClose@ when you're done.
  debuggeeRun :: FilePath  -- ^ path to executable to run as the debuggee
              -> FilePath  -- ^ filename of socket (e.g. @"/tmp/ghc-debug"@)
              -> IO Debuggee
  debuggeeRun exeName socketName = GD.debuggeeRun exeName socketName
  
  -- | Run a debuggee and connect to it. Use @debuggeeClose@ when you're done.
  debuggeeConnect :: FilePath  -- ^ filename of socket (e.g. @"/tmp/ghc-debug"@)
                  -> IO Debuggee
  debuggeeConnect socketName = GD.debuggeeConnect socketName
  
  snapshotConnect :: FilePath -> IO Debuggee
  snapshotConnect snapshotName = GD.snapshotInit snapshotName
  
  -- | Close the connection to the debuggee.
  debuggeeClose :: Debuggee -> IO ()
  debuggeeClose = GD.debuggeeClose
  
  -- | Request the debuggee's root pointers.
  rootClosures :: Debuggee -> IO [Closure]
  rootClosures e = run e $ do
    closurePtrs <- request RequestRoots
    closures <- GD.dereferenceClosures closurePtrs
    return [ Closure closurePtr' closure
              | closurePtr' <- closurePtrs
              | closure <- closures
              ]
  
  -- | A client can save objects by calling a special RTS method
  -- This function returns the closures it saved.
  savedClosures :: Debuggee -> IO [Closure]
  savedClosures e = run e $ do
    closurePtrs <- request RequestSavedObjects
    closures <- GD.dereferenceClosures closurePtrs
    return $ zipWith Closure
              closurePtrs
              closures
  
  profile :: Debuggee -> FilePath -> IO ()
  profile dbg fp = do
    c <- run dbg $ do
      roots <- GD.gcRoots
      GD.censusClosureType roots
    GD.writeCensusByClosureType fp c
  
  snapshot :: Debuggee -> FilePath -> IO ()
  snapshot dbg fp = do
    dir <- snapshotDirectory
    createDirectoryIfMissing True dir
    GD.makeSnapshot dbg (dir </> fp)
  
  retainersOfConstructor :: Maybe [ClosurePtr] -> Debuggee -> String -> IO [[Closure]]
  retainersOfConstructor mroots dbg con_name = do
    run dbg $ do
      roots <- maybe GD.gcRoots return mroots
      stack <- GD.findRetainersOfConstructor (Just 100) roots con_name
      traverse (\cs -> zipWith Closure cs <$> (GD.dereferenceClosures cs)) stack
  
  retainersOfConstructorExact :: Debuggee -> String -> IO [[Closure]]
  retainersOfConstructorExact dbg con_name = do
    run dbg $ do
      roots <- GD.gcRoots
      stack <- GD.findRetainersOfConstructorExact (Just 100) roots con_name
      traverse (\cs -> zipWith Closure cs <$> (GD.dereferenceClosures cs)) stack
  
  -- -- | Request the description for an info table.
  -- -- The `InfoTablePtr` is just used for the equality
  -- requestConstrDesc :: Debuggee -> PayloadWithKey InfoTablePtr ClosurePtr -> IO ConstrDesc
  -- requestConstrDesc (Debuggee e _) = run e $ request RequestConstrDesc
  
  -- -- | Lookup source information of an info table
  -- requestSourceInfo :: Debuggee -> InfoTablePtr -> IO [String]
  -- requestSourceInfo (Debuggee e _) = run e $ request RequestSourceInfo
  
  -- -- | Request a set of closures.
  -- requestClosures :: Debuggee -> [ClosurePtr] -> IO [RawClosure]
  -- requestClosures (Debuggee e _) = run e $ request RequestClosures
  
  type Closure = DebugClosure PayloadCont ConstrDescCont StackCont ClosurePtr
  
  data ListItem a b c d = ListData | ListOnlyInfo InfoTablePtr | ListFullClosure (DebugClosure a b c d)
  
  data DebugClosure p cd s c
    = Closure
      { _closurePtr :: ClosurePtr
      , _closureSized :: DebugClosureWithSize p cd s c
      }
    | Stack
      { _stackPtr :: StackCont
      , _stackStack :: GD.GenStackFrames c
      }
  
  toPtr :: DebugClosure p cd s c -> Ptr
  toPtr (Closure cp _) = CP cp
  toPtr (Stack sc _)   = SP sc
  
  data Ptr = CP ClosurePtr | SP StackCont deriving (Eq, Ord)
  
  deriving instance Eq StackCont
  deriving instance Ord StackCont
  
  dereferencePtr :: Debuggee -> Ptr -> IO (DebugClosure PayloadCont ConstrDescCont StackCont ClosurePtr)
  dereferencePtr dbg (CP cp) = run dbg (Closure <$> pure cp <*> GD.dereferenceClosure cp)
  dereferencePtr dbg (SP sc) = run dbg (Stack <$> pure sc <*> GD.dereferenceStack sc)
  
  instance Quadtraversable DebugClosure where
    quadtraverse p f g h (Closure cp c) = Closure cp <$> quadtraverse p f g h c
    quadtraverse _ _ _ h (Stack sp s) = Stack sp <$> traverse h s
  
  closureShowAddress :: DebugClosure p cd s c -> String
  closureShowAddress (Closure c _) = show c
  closureShowAddress (Stack   s _) = show s
  
  -- | Get the exclusive size (not including any referenced closures) of a closure.
  closureExclusiveSize :: DebugClosure p cd s c -> Size
  closureExclusiveSize (Stack{}) = Size (-1)
  closureExclusiveSize (Closure _ c) = (GD.dcSize c)
  
  -- | Get the retained size (including all dominated closures) of a closure.
  closureRetainerSize :: Analysis -> DebugClosure p cd s c -> RetainerSize
  closureRetainerSize analysis c = snd (closureExcAndRetainerSizes analysis c)
  
  closureExcAndRetainerSizes :: Analysis -> DebugClosure p cd s c -> (Size, RetainerSize)
  closureExcAndRetainerSizes _ Stack{} = (Size (-1), RetainerSize (-1))
    -- ^ TODO How should we handle stack size? only used space on the stack?
    -- Include underflow frames? Return Maybe?
  closureExcAndRetainerSizes analysis (Closure cPtr _) =
    let getSizes = analysisSizes analysis
    in getSizes cPtr
  
  closureSourceLocation :: Debuggee -> DebugClosure p cd s c -> IO (Maybe SourceInformation)
  closureSourceLocation _ (Stack _ _) = return Nothing
  closureSourceLocation e (Closure _ c) = run e $ do
    request (RequestSourceInfo (tableId (info (noSize c))))
  
  closureInfoPtr :: DebugClosure p cd s c -> Maybe InfoTablePtr
  closureInfoPtr (Stack {}) = Nothing
  closureInfoPtr (Closure _ c) = Just (tableId (info (noSize c)))
  
  infoSourceLocation :: Debuggee -> InfoTablePtr -> IO (Maybe SourceInformation)
  infoSourceLocation e ip = run e $ request (RequestSourceInfo ip)
  
  -- | Get the directly referenced closures (with a label) of a closure.
  closureReferences :: Debuggee -> DebugClosure PayloadCont ConstrDesc StackCont ClosurePtr -> IO [(String, ListItem PayloadCont ConstrDescCont StackCont ClosurePtr)]
  closureReferences e (Stack _ stack) = run e $ do
    let action (GD.SPtr ptr) = ("Pointer", ListFullClosure . Closure ptr <$> GD.dereferenceClosure ptr)
        action (GD.SNonPtr dat) = ("Data:" ++ show dat, return ListData)
  
        frame_items frame = ("Info: " ++ show (tableId (frame_info frame)), return (ListOnlyInfo (tableId (frame_info frame)))) :
                            map action (GD.values frame)
  
        add_frame_ix ix (lbl, x) = ("Frame " ++ show ix ++ " " ++ lbl, x)
    let lblAndPtrs = [ map (add_frame_ix frameIx) (frame_items frame)
                        | (frameIx, frame) <- zip [(0::Int)..] (GD.getFrames stack)
                     ]
  --  traverse GD.dereferenceClosures (snd <$> lblAndPtrs)
    traverse (traverse id) (concat lblAndPtrs)
    {-
    return $ zipWith (\(lbl,ptr) c -> (lbl, Closure ptr c))
              lblAndPtrs
              closures
              -}
  closureReferences e (Closure _ closure) = run e $ do
    let refPtrs = closureReferencesAndLabels (unDCS closure)
    forM refPtrs $ \(label, ptr) -> case ptr of
      Left cPtr -> do
        refClosure' <- GD.dereferenceClosure cPtr
        return (label, ListFullClosure $ Closure cPtr refClosure')
      Right sPtr -> do
        refStack' <- GD.dereferenceStack sPtr
        return (label, ListFullClosure $ Stack sPtr refStack')
  
  reverseClosureReferences :: HG.HeapGraph Size
                           -> HG.ReverseGraph
                           -> Debuggee
                           -> DebugClosure HG.PapHI ConstrDesc HG.StackHI (Maybe HG.HeapGraphIndex)
                           -> IO [(String, DebugClosure
                                              HG.PapHI
                                              ConstrDesc HG.StackHI
                                              (Maybe HG.HeapGraphIndex))]
  reverseClosureReferences hg rm _ c =
    case c of
      Stack {} -> error "Nope - Stack"
      Closure cp _ -> case (HG.reverseEdges cp rm) of
                        Nothing -> return []
                        Just es ->
                          let revs = mapMaybe (flip HG.lookupHeapGraph hg) es
                          in return [(show n, Closure (HG.hgeClosurePtr hge)
                                                 (DCS (HG.hgeData hge) (HG.hgeClosure hge) ))
                                      | (n, hge) <- zip [0 :: Int ..] revs]
  
  lookupHeapGraph :: HG.HeapGraph Size -> ClosurePtr -> Maybe (DebugClosure HG.PapHI ConstrDesc HG.StackHI (Maybe HG.HeapGraphIndex))
  lookupHeapGraph hg cp =
    case HG.lookupHeapGraph cp hg of
      Just (HG.HeapGraphEntry ptr d s) -> Just (Closure ptr (DCS s d))
      Nothing -> Nothing
  
  fillConstrDesc :: Debuggee
                 -> DebugClosure pap ConstrDescCont s c
                 -> IO (DebugClosure pap ConstrDesc s c)
  fillConstrDesc e closure = do
    run e $ GD.quadtraverse pure GD.dereferenceConDesc pure pure closure
  
  -- | Pretty print a closure
  closurePretty :: Show c => DebugClosure p ConstrDesc s c ->  String
  closurePretty (Stack _ _) = "STACK"
  closurePretty (Closure _ closure) =
    HG.ppClosure
      "??"
      (\_ refPtr -> show refPtr)
      0
      (unDCS closure)
  
  -- $dominatorTree
  --
  -- Closure `a` dominates closure `b` if all paths from GC roots to `b` pass
  -- through `a`. This means that if `a` is GCed then all dominated closures can
  -- be GCed. The relationship is transitive. Transitive edges are omitted in the
  -- "dominator tree".
  --
  -- see http://kohlerm.blogspot.com/2009/02/memory-leaks-are-easy-to-find.html
  
  -- | The roots of the dominator tree.
  dominatorRootClosures :: Debuggee -> Analysis -> IO [Closure]
  dominatorRootClosures e analysis = run e $ do
    let domRoots = analysisDominatorRoots analysis
    closures <- GD.dereferenceClosures domRoots
    return [ Closure closurePtr' closure
              | closurePtr' <- domRoots
              | closure <- closures
              ]
  
  -- | Get the dominatess of a closure i.e. the children in the dominator tree.
  closureDominatees :: Debuggee -> Analysis -> DebugClosure p cd s ClosurePtr -> IO [Closure]
  closureDominatees _ _ (Stack{}) = error "TODO dominator tree does not yet support STACKs"
  closureDominatees e analysis (Closure cPtr _) = run e $ do
    let cPtrToDominatees = analysisDominatees analysis
        cPtrs = cPtrToDominatees cPtr
    closures <- GD.dereferenceClosures cPtrs
    return [ Closure closurePtr' closure
              | closurePtr' <- cPtrs
              | closure <- closures
              ]
  
  --
  -- Internal Stuff
  --
  
  closureReferencesAndLabels :: GD.DebugClosure pap string stack pointer -> [(String, Either pointer stack)]
  closureReferencesAndLabels closure = case closure of
    TSOClosure {..} ->
      [ ("Stack", Left tsoStack)
      , ("Link", Left _link)
      , ("Global Link", Left global_link)
      , ("TRec", Left trec)
      , ("Blocked Exceptions", Left blocked_exceptions)
      , ("Blocking Queue", Left bq)
      ]
    StackClosure{..} -> [("Frames", Right frames )]
    WeakClosure {..} -> [ ("Key", Left key)
                        , ("Value", Left value)
                        , ("C Finalizers", Left cfinalizers)
                        , ("Finalizer", Left finalizer)
                        ] ++
                        [ ("Link", Left link)
                        | Just link <- [mlink] -- TODO do we want to show NULL pointers some how?
                        ]
    TVarClosure {..} -> [("val", Left current_value)]
    MutPrimClosure {..} -> withArgLables ptrArgs
    ConstrClosure {..} -> withFieldLables ptrArgs
    ThunkClosure {..} -> withArgLables ptrArgs
    SelectorClosure {..} -> [("Selectee", Left selectee)]
    IndClosure {..} -> [("Indirectee", Left indirectee)]
    BlackholeClosure {..} -> [("Indirectee", Left indirectee)]
    APClosure {..} -> ("Function", Left fun) : [] -- TODO withBitmapLables ap_payload
    PAPClosure {..} -> ("Function", Left fun) : [] -- TODO: withBitmapLables pap_payload
    APStackClosure {..} -> ("Function", Left fun) : ("Frames", Right payload) : []
    BCOClosure {..} -> [ ("Instructions", Left instrs)
                        , ("Literals", Left literals)
                        , ("Byte Code Objects", Left bcoptrs)
                        ]
    ArrWordsClosure {} -> []
    MutArrClosure {..} -> withIxLables mccPayload
    SmallMutArrClosure {..} -> withIxLables mccPayload
    MutVarClosure {..} -> [("Value", Left var)]
    MVarClosure {..} -> [ ("Queue Head", Left queueHead)
                        , ("Queue Tail", Left queueTail)
                        , ("Value", Left value)
                        ]
    FunClosure {..} -> withArgLables ptrArgs
    BlockingQueueClosure {..} -> [ ("Link", Left link)
                                  , ("Black Hole", Left blackHole)
                                  , ("Owner", Left owner)
                                  , ("Queue", Left queue)
                                  ]
    OtherClosure {..} -> ("",) . Left <$> hvalues
    TRecChunkClosure{}  -> [] --TODO
    UnsupportedClosure {} -> []
    where
    withIxLables elements   = [("[" <> show i <> "]" , Left x) | (i, x) <- zip [(0::Int)..] elements]
    withArgLables ptrArgs   = [("Argument " <> show i, Left x) | (i, x) <- zip [(0::Int)..] ptrArgs]
    withFieldLables ptrArgs = [("Field " <> show i   , Left x) | (i, x) <- zip [(0::Int)..] ptrArgs]
  --  withBitmapLables pap = [("Argument " <> show i   , Left x) | (i, SPtr x) <- zip [(0::Int)..] (getValues pap)]
  
  --
  