{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}

module Main where

import Control.Monad
import qualified Data.ByteString.Char8 as B8
import Data.Char (isDigit)
import Data.Function
import qualified Data.HashTable.IO as H
import Data.IORef
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.List (sortOn)
import qualified Data.Text as T
import Data.Word
import Graphics.Vega.VegaLite hiding (name)
import OpenTelemetry.Common
import OpenTelemetry.EventlogStreaming_Internal
import System.Environment
import Text.Printf

type HashTable k v = H.BasicHashTable k v

data PerOperationStats = PerOperationStats
  { count :: !Int,
    min_ns :: !Word64,
    max_ns :: !Word64,
    total_ns :: !Word64
  }

data MetricStats = MetricStats
  { max_threads :: !Int,
    -- | Per-capability
    total_alloc_bytes :: !(IntMap Int),
    max_live_bytes :: !Int
  }

main :: IO ()
main = do
  args <- getArgs
  case args of
    [path] -> do
      (opCounts :: H.CuckooHashTable T.Text PerOperationStats) <- H.new
      metricStats <- newIORef (MetricStats 0 IntMap.empty 0)
      let span_exporter =
            Exporter
              ( \sps -> do
                  forM_ sps $ \sp -> do
                    let duration = spanFinishedAt sp - spanStartedAt sp
                    H.mutate
                      opCounts
                      (spanOperation sp)
                      ( \m ->
                          ( Just $
                              maybe
                                (PerOperationStats 1 duration duration duration)
                                ( \PerOperationStats {..} ->
                                    PerOperationStats
                                      (count + 1)
                                      (min min_ns duration)
                                      (max max_ns duration)
                                      (total_ns + duration)
                                )
                                m,
                            ()
                          )
                      )
                  pure ExportSuccess
              )
              (pure ())
      metric_exporter <-
        aggregated $
          Exporter
            ( \metrics -> do
                forM_ metrics $ \(AggregatedMetric (CaptureInstrument _ name) (MetricDatapoint _ value)) ->
                  modifyIORef metricStats $ \s -> case splitCapability (B8.unpack name) of
                    (_, "threads") -> s {max_threads = max value (max_threads s)}
                    (Just cap, "heap_alloc_bytes") -> s {total_alloc_bytes = IntMap.insert cap value (total_alloc_bytes s)}
                    (_, "heap_live_bytes") -> s {max_live_bytes = max value (max_live_bytes s)}
                    _ -> s
                pure ExportSuccess
            )
            (pure ())
      exportEventlog span_exporter metric_exporter path
      leaderboard <- sortOn (total_ns . snd) <$> H.toList opCounts
      printf "Count\tTot ms\tMin ms\tMax ms\tOperation\n"
      printf "-----\t------\t------\t------\t---------\n"
      mapM_
        ( \(op, PerOperationStats {..}) ->
            printf
              "%d\t%d\t%d\t%d\t%s\n"
              count
              (total_ns `div` 1000000)
              (min_ns `div` 1000000)
              (max_ns `div` 1000000)
              (T.unpack op)
        )
        leaderboard

      putStrLn "---"

      MetricStats {max_threads, total_alloc_bytes, max_live_bytes} <- readIORef metricStats
      printf "Max threads: %v\n" max_threads
      putStrLn "Total allocations:"
      _ <-
        IntMap.traverseWithKey
          (\cap bytes -> printf "  * Capability %v: %vMB\n" cap (bytes `div` 1000000))
          total_alloc_bytes
      printf "Max live: %vMB\n" (max_live_bytes `div` 1000000)

      putStrLn "---"

      let vega_visualization =
            toVegaLite
              [ title "Total duration of operation" [],
                vega_dat,
                mark Bar [],
                vega_enc
              ]

          vega_enc =
            []
              & position Y [PName "op", PmType Nominal, PAxis [AxTitle "operation"]]
              & position X [PName "dur", PmType Quantitative, PAxis [AxTitle "duration in nanoseconds"]]
              & encoding
          vega_dat =
            []
              & dataColumn "op" (Strings (map fst leaderboard))
              & dataColumn "dur" (Numbers (map (fromIntegral . total_ns . snd) leaderboard))
              & dataFromColumns []

      toHtmlFile "eventlog-summary.html" vega_visualization
      putStrLn "Generated report: eventlog-summary.html"
    _ -> do
      putStrLn "Usage:"
      putStrLn ""
      putStrLn "  eventlog-summary <program.eventlog>"

-- | Parse capability number out of event/metric name. If it fails, returns
-- (Nothing, the original string)
splitCapability :: String -> (Maybe Int, String)
splitCapability fullName@('c' : 'a' : 'p' : '_' : rest) =
  case span isDigit rest of
    (numStr, '_' : name) -> (Just (read numStr), name)
    _ -> (Nothing, fullName)
splitCapability name = (Nothing, name)