--
-- Prof2Dot.hs
--
-- Main module for prof2dot, which converts cost center profiles
-- generated by the -px flag into GraphViz's dot format.
--
-- Gregory Wright, 28 July 2007
--
-- Copyright (c) 2007, 2008 Antiope Associates LLC, all rights reserved.
--

module Main where


import Data.List
import Data.Maybe
import System.Console.GetOpt
import System
import System.FilePath
import System.IO

import Grapher
import ParseProfile


data Options =
        Options { debug         :: Bool,
                  printHelp     :: Bool,
                  nodeStyle     :: NodeStyle,
                  annotateEdges :: Bool,
                  colorCoord    :: ColorCoord,
                  showModules   :: Bool,
                  outputFile    :: Maybe FilePath }
        deriving Show

defaultOptions :: Options
defaultOptions =
        Options { debug         = False,
                  printHelp     = False,
                  nodeStyle     = Verbose,
                  annotateEdges = False,
                  colorCoord    = ColorCalls,
                  showModules   = False,
                  outputFile    = Nothing }


options :: [ OptDescr (Options -> Options) ]
options = 
  [ Option "b" ["brief"]       (NoArg  (\opt -> opt {nodeStyle  = Brief}      ))  "show only the module and cost center name",
    Option "e" ["edges"]       (NoArg  (\opt -> opt {annotateEdges = True}    ))  "annotate edges with (call counts, ticks, allocations)",
    Option "c" ["colorcalls"]  (NoArg  (\opt -> opt {colorCoord = ColorCalls} ))  "colorize by cost center call count (default)",
    Option "t" ["colortime"]   (NoArg  (\opt -> opt {colorCoord = ColorTicks} ))  "colorize by cost center ticks",
    Option "a" ["colorallocs"] (NoArg  (\opt -> opt {colorCoord = ColorAllocs}))  "colorize by cost center memory allocations",
    Option "m" ["modules"]     (NoArg  (\opt -> opt {showModules = True}      ))  "group cost centers by module",
    Option "n" ["nocolor"]     (NoArg  (\opt -> opt {colorCoord = NoColor}    ))  "no colorization",
    Option "d" ["debug"]       (NoArg  (\opt -> opt {debug = True}            ))  "unpruned graph with debug information",
    Option "?" ["help"]        (NoArg  (\opt -> opt {printHelp = True}        ))  "print this help message and exit",
    Option "o" ["output"]      (ReqArg (\arg opt -> opt {outputFile = Just arg})  "filename") "output filename (default: stdout)" ]


usageHeader :: String
usageHeader = "prof2dot converts a ghc runtime profile (generated by running a ghc-compiled\n\
              \program with the +RTS -px -RTS option) into a dot file.  The dot file\n\
              \describes the call graph of the program and can be rendered in a variety of\n\
              \formats using graphviz's dot program.\n\n\
              \The default input is stdin, the default output is stdout.\n\n\
              \Usage: prof2dot [OPTIONS] inputfile"

  
parseOptions :: [ String ] -> IO (Options, Maybe FilePath)
parseOptions args = do
  case getOpt Permute options args of
    (opts, files, []  ) -> do
                               let
                                   opts' = foldl' (flip id) defaultOptions opts
                                   file  = if null files then Nothing else Just (head files)
                               return (opts', file)
    (_,    _,     errs) -> ioError (userError (concat errs ++ usageInfo usageHeader options))


main :: IO ()
main = do
  args              <- getArgs
  (opts, inputFile) <- parseOptions args

  if printHelp opts
      then putStrLn (usageInfo usageHeader options)
      else do
        ohdl          <- if isJust (outputFile opts)
                             then openFile (fromJust (outputFile opts)) WriteMode
                             else return stdout
        (ihdl, iName) <- if isJust inputFile
                             then do
                               ih <- openFile (fromJust inputFile) ReadMode
                               return (ih, fromJust inputFile)
                             else return (stdin, "stdin")
        pfil <- hGetContents ihdl

        let
                grafOptions = GraphOptions { graphDebug         = debug opts,
                                             graphColorCoord    = colorCoord opts,
                                             graphNodeStyle     = nodeStyle opts,
                                             graphAnnotateEdges = annotateEdges opts,
                                             graphShowModules   = showModules opts }

                grafName = if isJust inputFile
                               then takeBaseName (fromJust inputFile)
                               else "stdin"
                prof = parseProfile iName pfil
                graf = if isJust prof
                           then profToDot grafOptions grafName (fromJust prof)
                           else error "parse error. Are you sure the input was generated by running with the +RTS -px -RTS option?"

        hPutStr ohdl graf
        hClose  ohdl