Diagnostics.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : Fractal.Layer.Diagnostics
-- Description : Diagnostics and visualization for layer initialization
-- Stability   : experimental
-- Portability : Portable
--
-- This module provides tools for visualizing and debugging layer initialization.
-- It can track how layers are composed, which services are shared, and render
-- the dependency tree in various formats.
--
-- == Example Usage
--
-- @
-- appLayer :: Layer IO () (Config, Database, WebServer)
-- appLayer = configLayer >>> (dbLayer &&& webLayer)
--
-- main :: IO ()
-- main = withLayerDiagnostics appLayer () $ \(env, diags) -> do
--   -- Print ASCII tree
--   putStrLn $ renderLayerTree diags
--
--   -- Or export as JSON
--   BSL.writeFile "layer-tree.json" (encode diags)
--
--   -- Use the environment
--   let (cfg, db, ws) = env
--   serveRequests ws
-- @
module Fractal.Layer.Diagnostics
  ( -- * Diagnostics Types
    LayerDiagnostics (..),
    LayerNode (..),
    LayerNodeType (..),
    ResourceStatus (..),

    -- * Diagnostics Interceptor
    createDiagnosticsInterceptor,
    DiagnosticsCollector,
    newDiagnosticsCollector,
    finalizeDiagnostics,

    -- * Running with Diagnostics
    withLayerDiagnostics,
    buildLayerDiagnostics,

    -- * Rendering
    renderLayerTree,
    renderLayerTreeDetailed,

    -- * Snapshots
    snapshotDiagnostics,

    -- * JSON Export
    diagnosticsToJSON,
  )
where

import Control.Monad.IO.Class
import Data.Aeson (ToJSON (..), FromJSON (..), object, (.=), (.:), Value, withObject, withText)
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Control.Concurrent.MVar
import Data.Text (Text)
import qualified Data.Text as T
import Data.Time.Clock (NominalDiffTime, getCurrentTime, diffUTCTime, UTCTime)
import Data.Typeable
import Fractal.Layer.Interceptor
import UnliftIO (MonadUnliftIO)

-- Import Layer types and functions (defined later in dependency order)
import Fractal.Layer (Layer, runLayerWithInterceptor, withLayerAndInterceptor)

-------------------------------------------------------------------------------
-- Types
-------------------------------------------------------------------------------

-- | Complete diagnostics for a layer initialization
data LayerDiagnostics = LayerDiagnostics
  { rootNode :: LayerNode
  -- ^ The root of the initialization tree
  , totalDuration :: Double
  -- ^ Total time taken to initialize (seconds)
  , totalResources :: Int
  -- ^ Total number of resources allocated
  , sharedResources :: Int
  -- ^ Number of resources that were shared/cached
  }
  deriving (Show)

-- | A node in the layer initialization tree
data LayerNode = LayerNode
  { nodeId :: Text
  -- ^ Unique identifier for this node
  , nodeName :: Text
  -- ^ Human-readable name (e.g., "DatabaseLayer", "ConfigService")
  , nodeType :: LayerNodeType
  -- ^ Type of layer node
  , resourceType :: Maybe TypeRep
  -- ^ The type of resource produced (if available)
  , status :: ResourceStatus
  -- ^ Status of this resource
  , duration :: Maybe Double
  -- ^ Time taken to initialize this layer (seconds)
  , children :: [LayerNode]
  -- ^ Child layers
  , metadata :: HashMap Text Text
  -- ^ Additional metadata
  }
  deriving (Show)

-- | The type of layer node
data LayerNodeType
  = ResourceNode
  -- ^ A resource that was allocated
  | EffectNode
  -- ^ An effect that was run
  | ServiceNode
  -- ^ A cached service
  | ComposedNode
  -- ^ A composed layer (>>> or &&&)
  | ParallelNode
  -- ^ Parallel composition (&&&)
  | SequentialNode
  -- ^ Sequential composition (>>>)
  deriving (Show, Eq)

-- | Status of a resource
data ResourceStatus
  = Initializing
  -- ^ Currently being initialized
  | Initialized
  -- ^ Successfully initialized
  | Failed String
  -- ^ Failed to initialize with error
  | SharedReference Text
  -- ^ Reference to a shared resource (with node ID)
  deriving (Show, Eq)

-------------------------------------------------------------------------------
-- JSON Instances
-------------------------------------------------------------------------------

instance ToJSON LayerDiagnostics where
  toJSON LayerDiagnostics{..} =
    object
      [ "root" .= rootNode
      , "totalDuration" .= totalDuration
      , "totalResources" .= totalResources
      , "sharedResources" .= sharedResources
      ]

instance ToJSON LayerNode where
  toJSON LayerNode{..} =
    object
      [ "id" .= nodeId
      , "name" .= nodeName
      , "type" .= nodeType
      , "resourceType" .= (show <$> resourceType)
      , "status" .= status
      , "duration" .= duration
      , "children" .= children
      , "metadata" .= metadata
      ]

instance ToJSON LayerNodeType where
  toJSON = \case
    ResourceNode -> "resource"
    EffectNode -> "effect"
    ServiceNode -> "service"
    ComposedNode -> "composed"
    ParallelNode -> "parallel"
    SequentialNode -> "sequential"

instance ToJSON ResourceStatus where
  toJSON = \case
    Initializing -> object ["status" .= ("initializing" :: Text)]
    Initialized -> object ["status" .= ("initialized" :: Text)]
    Failed err -> object ["status" .= ("failed" :: Text), "error" .= err]
    SharedReference nodeId -> object ["status" .= ("shared" :: Text), "reference" .= nodeId]

-- FromJSON instances for deserialization
instance FromJSON LayerDiagnostics where
  parseJSON = withObject "LayerDiagnostics" $ \v -> LayerDiagnostics
    <$> v .: "root"
    <*> v .: "totalDuration"
    <*> v .: "totalResources"
    <*> v .: "sharedResources"

instance FromJSON LayerNode where
  parseJSON = withObject "LayerNode" $ \v -> LayerNode
    <$> v .: "id"
    <*> v .: "name"
    <*> v .: "type"
    <*> pure Nothing  -- resourceType is serialized as String, skip for now
    <*> v .: "status"
    <*> v .: "duration"
    <*> v .: "children"
    <*> v .: "metadata"

instance FromJSON LayerNodeType where
  parseJSON = withText "LayerNodeType" $ \case
    "resource" -> pure ResourceNode
    "effect" -> pure EffectNode
    "service" -> pure ServiceNode
    "composed" -> pure ComposedNode
    "parallel" -> pure ParallelNode
    "sequential" -> pure SequentialNode
    _ -> fail "Unknown layer node type"

instance FromJSON ResourceStatus where
  parseJSON = withObject "ResourceStatus" $ \v -> do
    status <- v .: "status"
    case status :: Text of
      "initializing" -> pure Initializing
      "initialized" -> pure Initialized
      "failed" -> Failed <$> v .: "error"
      "shared" -> SharedReference <$> v .: "reference"
      _ -> fail "Unknown status"

-------------------------------------------------------------------------------
-- Diagnostics Collector (Interceptor-based)
-------------------------------------------------------------------------------

-- | Mutable state for collecting diagnostics during layer initialization
data DiagnosticsCollectorState = DiagnosticsCollectorState
  { collectorNodeStack :: ![LayerNode]
  -- ^ Stack of nodes being built (top = current)
  , collectorNextId :: !Int
  -- ^ Next available node ID
  , collectorStartTime :: !UTCTime
  -- ^ When collection started
  , collectorServiceMap :: !(HashMap TypeRep Text)
  -- ^ Map of service types to their node IDs (for tracking sharing)
  , collectorTotalResources :: !Int
  -- ^ Total number of resources allocated
  , collectorSharedResources :: !Int
  -- ^ Number of shared/reused resources
  }

-- | Opaque handle to a diagnostics collector
newtype DiagnosticsCollector = DiagnosticsCollector (MVar DiagnosticsCollectorState)

-- | Create a new diagnostics collector
newDiagnosticsCollector :: MonadIO m => m DiagnosticsCollector
newDiagnosticsCollector = liftIO $ do
  start <- getCurrentTime
  ref <- newMVar DiagnosticsCollectorState
    { collectorNodeStack = [rootNode]
    , collectorNextId = 1
    , collectorStartTime = start
    , collectorServiceMap = HashMap.empty
    , collectorTotalResources = 0
    , collectorSharedResources = 0
    }
  pure $ DiagnosticsCollector ref
  where
    rootNode = LayerNode
      { nodeId = "root"
      , nodeName = "Root"
      , nodeType = ComposedNode
      , resourceType = Nothing
      , status = Initializing
      , duration = Nothing
      , children = []
      , metadata = HashMap.empty
      }

-- | Finalize diagnostics and get the complete tree
finalizeDiagnostics :: MonadIO m => DiagnosticsCollector -> m LayerDiagnostics
finalizeDiagnostics (DiagnosticsCollector ref) = liftIO $ do
  state <- readMVar ref
  endTime <- getCurrentTime
  let totalDur = realToFrac (diffUTCTime endTime (collectorStartTime state))
      rootNode = case collectorNodeStack state of
        [] -> error "Diagnostics collector has empty stack"
        (node:_) -> node { status = Initialized, duration = Just totalDur }
  pure $ LayerDiagnostics
    { rootNode = rootNode
    , totalDuration = totalDur
    , totalResources = collectorTotalResources state
    , sharedResources = collectorSharedResources state
    }

-- | Create a LayerInterceptor that collects diagnostics
createDiagnosticsInterceptor :: MonadIO m => DiagnosticsCollector -> LayerInterceptor m
createDiagnosticsInterceptor collector = LayerInterceptor
  { onResourceAcquire = \ctx -> liftIO $
      diagStartNode collector ctx ResourceNode
  , onResourceAcquireComplete = \_ dur -> liftIO $
      diagEndNode collector dur Initialized
  , onResourceRelease = \_ -> pure ()
  , onEffectRun = \ctx -> liftIO $
      diagStartNode collector ctx EffectNode
  , onEffectComplete = \_ dur -> liftIO $
      diagEndNode collector dur Initialized
  , onServiceCreate = \ctx -> liftIO $ do
      diagStartNode collector ctx ServiceNode
      diagRegisterService collector ctx
  , onServiceReuse = \name tr -> liftIO $
      diagAddSharedRef collector name tr
  , onCompositionStart = \typ -> liftIO $ do
      let nt = case typ of
            Sequential -> SequentialNode
            Parallel -> ParallelNode
          ctx = OperationContext (T.pack (show typ)) Nothing []
      diagStartNode collector ctx nt
  , onCompositionEnd = \_ dur -> liftIO $
      diagEndNode collector dur Initialized
  }

diagStartNode :: DiagnosticsCollector -> OperationContext -> LayerNodeType -> IO ()
diagStartNode (DiagnosticsCollector ref) ctx nt =
  modifyMVar_ ref $ \state -> do
    let newNode = LayerNode
          { nodeId = T.pack ("node-" <> show (collectorNextId state))
          , nodeName = operationName ctx
          , nodeType = nt
          , resourceType = operationType ctx
          , status = Initializing
          , duration = Nothing
          , children = []
          , metadata = HashMap.fromList (operationMetadata ctx)
          }
    pure state
      { collectorNodeStack = newNode : collectorNodeStack state
      , collectorNextId = collectorNextId state + 1
      }

diagEndNode :: DiagnosticsCollector -> NominalDiffTime -> ResourceStatus -> IO ()
diagEndNode (DiagnosticsCollector ref) dur st =
  modifyMVar_ ref $ \state ->
    pure $ case collectorNodeStack state of
      [] -> state
      [root] ->
        state { collectorNodeStack =
            [root { status = st, duration = Just (realToFrac dur) }]
        }
      (current:parent:rest) ->
        let completed = current { status = st, duration = Just (realToFrac dur) }
        in state { collectorNodeStack =
            parent { children = children parent ++ [completed] } : rest
        }

diagRegisterService :: DiagnosticsCollector -> OperationContext -> IO ()
diagRegisterService (DiagnosticsCollector ref) ctx =
  modifyMVar_ ref $ \s -> pure $ case operationType ctx of
    Just tr -> s
      { collectorServiceMap = HashMap.insert tr (operationName ctx) (collectorServiceMap s)
      , collectorTotalResources = collectorTotalResources s + 1
      }
    Nothing -> s
      { collectorTotalResources = collectorTotalResources s + 1
      }

diagAddSharedRef :: DiagnosticsCollector -> Text -> TypeRep -> IO ()
diagAddSharedRef (DiagnosticsCollector ref) name tr =
  modifyMVar_ ref $ \state ->
    case HashMap.lookup tr (collectorServiceMap state) of
      Just originalNodeId -> do
        let sn = LayerNode
              { nodeId = T.pack ("shared-" <> show (collectorNextId state))
              , nodeName = name
              , nodeType = ServiceNode
              , resourceType = Just tr
              , status = SharedReference originalNodeId
              , duration = Nothing
              , children = []
              , metadata = HashMap.empty
              }
        pure $ case collectorNodeStack state of
          [] -> state
          (top:rest) ->
            state
              { collectorNodeStack = top { children = children top ++ [sn] } : rest
              , collectorNextId = collectorNextId state + 1
              , collectorSharedResources = collectorSharedResources state + 1
              }
      Nothing -> pure state

-------------------------------------------------------------------------------
-- Running with Diagnostics
-------------------------------------------------------------------------------

-- | Build a layer with diagnostics tracking and return only the diagnostics.
--
-- The environment is discarded and resources are released immediately.
-- Useful for inspecting layer structure without keeping resources alive.
-- Use 'withLayerDiagnostics' when you need both the environment and diagnostics.
buildLayerDiagnostics ::
  MonadUnliftIO m =>
  -- | Layer to build
  Layer m deps env ->
  -- | Dependencies
  deps ->
  m LayerDiagnostics
buildLayerDiagnostics layer deps = do
  collector <- newDiagnosticsCollector
  let interceptor = createDiagnosticsInterceptor collector
  _ <- runLayerWithInterceptor interceptor deps layer
  finalizeDiagnostics collector

-- | Run a layer with diagnostics enabled.
-- Returns both the environment and the diagnostics information.
--
-- This function runs the layer with diagnostics collection enabled, then provides
-- both the initialized environment and collected diagnostics to your continuation.
--
-- Example:
-- @
-- withLayerDiagnostics myLayer config $ \(env, diags) -> do
--   putStrLn $ renderLayerTree diags
--   -- Use the environment...
-- @
withLayerDiagnostics ::
  MonadUnliftIO m =>
  -- | Layer to run
  Layer m deps env ->
  -- | Dependencies
  deps ->
  -- | Action with environment and diagnostics
  ((env, LayerDiagnostics) -> m r) ->
  m r
withLayerDiagnostics layer deps action = do
  collector <- newDiagnosticsCollector
  let interceptor = createDiagnosticsInterceptor collector
  withLayerAndInterceptor interceptor deps layer $ \env -> do
    diags <- finalizeDiagnostics collector
    action (env, diags)

-------------------------------------------------------------------------------
-- Rendering to Terminal
-------------------------------------------------------------------------------

-- | Render the layer tree as ASCII art (compact version)
--
-- Example output:
-- @
-- Root
-- ├── ConfigLayer (0.05s)
-- └── ParallelComposition (0.3s)
--     ├── DatabaseLayer (0.2s)
--     │   └── ConnectionPool [SHARED: svc-001]
--     └── WebServerLayer (0.1s)
--         └── MetricsCollector [SHARED: svc-001]
-- @
renderLayerTree :: LayerDiagnostics -> String
renderLayerTree diags =
  unlines $
    [ "Layer Initialization Tree"
    , "═════════════════════════"
    , ""
    , "⧗ Duration: " ++ show (totalDuration diags) ++ "s"
    , "◆ Resources: " ++ show (totalResources diags)
    , "↻ Shared: " ++ show (sharedResources diags)
    , ""
    ] ++ renderNode "" True (rootNode diags)

-- | Detailed rendering with metadata
renderLayerTreeDetailed :: LayerDiagnostics -> String
renderLayerTreeDetailed diags =
  unlines $
    [ "Layer Initialization Tree (Detailed)"
    , "═════════════════════════════════════"
    , ""
    , "⧗ Duration: " ++ show (totalDuration diags) ++ "s"
    , "◆ Resources: " ++ show (totalResources diags)
    , "↻ Shared: " ++ show (sharedResources diags)
    , ""
    ] ++ renderNodeDetailed "" True (rootNode diags)

-- | Render a single node in the tree
renderNode :: String -> Bool -> LayerNode -> [String]
renderNode prefix isLast node =
  let connector = if isLast then "└── " else "├── "
      extension = if isLast then "    " else "│   "
      typeSymbol = nodeTypeSymbol (nodeType node)
      nodeInfo = typeSymbol ++ " " ++ T.unpack (nodeName node) ++ formatDuration (duration node) ++ formatStatus (status node)
      header = prefix ++ connector ++ nodeInfo
      childPrefix = prefix ++ extension
      childCount = length (children node)
      childLines = concat $ zipWith (\i child -> renderNode childPrefix (i == childCount - 1) child) [1..] (children node)
  in header : childLines

-- | Render a node with detailed metadata
renderNodeDetailed :: String -> Bool -> LayerNode -> [String]
renderNodeDetailed prefix isLast node =
  let connector = if isLast then "└── " else "├── "
      extension = if isLast then "    " else "│   "
      typeSymbol = nodeTypeSymbol (nodeType node)
      nodeInfo = typeSymbol ++ " " ++ T.unpack (nodeName node) ++ " [" ++ show (nodeType node) ++ "]"
      header = prefix ++ connector ++ nodeInfo
      childPrefix = prefix ++ extension

      -- Add metadata lines
      metaLines = if HashMap.null (metadata node)
        then []
        else map (\(k, v) -> childPrefix ++ "  ▸ " ++ T.unpack k ++ ": " ++ T.unpack v) (HashMap.toList $ metadata node)

      -- Add type info if available
      typeLines = case resourceType node of
        Nothing -> []
        Just tr -> [childPrefix ++ "  ▸ Type: " ++ show tr]

      -- Add duration
      durationLines = case duration node of
        Nothing -> []
        Just d -> [childPrefix ++ "  ⧗ " ++ show d ++ "s"]

      -- Add status
      statusLines = [childPrefix ++ "  " ++ formatStatusDetailed (status node)]

      detailLines = typeLines ++ durationLines ++ statusLines ++ metaLines

      childCount = length (children node)
      childLines = concat $ zipWith (\i child -> renderNodeDetailed childPrefix (i == childCount - 1) child) [1..] (children node)
  in (header : detailLines) ++ childLines

-- | Symbol for each node type
nodeTypeSymbol :: LayerNodeType -> String
nodeTypeSymbol ResourceNode = "◆"
nodeTypeSymbol EffectNode = "⚡"
nodeTypeSymbol ServiceNode = "◉"
nodeTypeSymbol ComposedNode = "⊕"
nodeTypeSymbol ParallelNode = "⋈"
nodeTypeSymbol SequentialNode = "⇒"

formatDuration :: Maybe Double -> String
formatDuration Nothing = ""
formatDuration (Just d) = " ⧗" ++ show d ++ "s"

formatStatus :: ResourceStatus -> String
formatStatus Initializing = " ⟳"
formatStatus Initialized = " ✓"
formatStatus (Failed err) = " ✗ [" ++ err ++ "]"
formatStatus (SharedReference refId) = " ↻ " ++ T.unpack refId

formatStatusDetailed :: ResourceStatus -> String
formatStatusDetailed Initializing = "⟳ Initializing"
formatStatusDetailed Initialized = "✓ Initialized"
formatStatusDetailed (Failed err) = "✗ Failed - " ++ err
formatStatusDetailed (SharedReference refId) = "↻ Shared reference to " ++ T.unpack refId

-- | Get a snapshot of current diagnostics without finalizing
--
-- This is useful for live rendering while initialization is still in progress.
-- Unlike 'finalizeDiagnostics', this doesn't mark the collector as finished.
snapshotDiagnostics :: MonadIO m => DiagnosticsCollector -> m LayerDiagnostics
snapshotDiagnostics (DiagnosticsCollector ref) = liftIO $ do
  state <- readMVar ref
  endTime <- getCurrentTime
  let totalDur = realToFrac (diffUTCTime endTime (collectorStartTime state))
      root = case collectorNodeStack state of
        [] -> LayerNode
          { nodeId = "empty"
          , nodeName = "Empty"
          , nodeType = ComposedNode
          , resourceType = Nothing
          , status = Initialized
          , duration = Nothing
          , children = []
          , metadata = HashMap.empty
          }
        (n:_) -> n
  pure LayerDiagnostics
    { rootNode = root
    , totalDuration = totalDur
    , totalResources = collectorTotalResources state
    , sharedResources = collectorSharedResources state
    }

-------------------------------------------------------------------------------
-- JSON Export
-------------------------------------------------------------------------------

-- | Convert diagnostics to a JSON Value
diagnosticsToJSON :: LayerDiagnostics -> Value
diagnosticsToJSON = toJSON