--
-- 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 { printHelp      :: Bool,
                         nodeStyle      :: NodeStyle,
                         colorCoord     :: ColorCoord,
                         convertToGraph :: Bool,
                         showModules    :: Bool,
                         outputFile     :: Maybe FilePath }
             deriving Show

defaultOptions :: Options
defaultOptions = Options { printHelp      = False,
                           nodeStyle      = Verbose,
                           colorCoord     = NoColor,
                           convertToGraph = True,
                           showModules    = False,
                           outputFile     = Nothing }


options :: [ OptDescr (Options -> Options) ]
options = 
  [ Option "b" ["brief"]       (NoArg  (\opt -> opt {nodeStyle  = Brief}      ))  "show only the names of call centers",
    Option "c" ["colorcalls"]  (NoArg  (\opt -> opt {colorCoord = ColorCalls} ))  "colorize by cost center call count",
    Option "t" ["colortime"]   (NoArg  (\opt -> opt {colorCoord = ColorTicks} ))  "colorize by cost center time",
    Option "a" ["colorallocs"] (NoArg  (\opt -> opt {colorCoord = ColorAllocs}))  "colorize by cost center memory allocations",
    Option "u" ["ungraph"]     (NoArg  (\opt -> opt {convertToGraph = False}  ))  "convert the call graph to a call tree",
    Option "m" ["modules"]     (NoArg  (\opt -> opt {showModules = True}      ))  "group cost centers by module",
    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 program\n\
              \with the +RTS -px -RTS option) into a dot file.  The dot file describes the call graph\n\
              \of the program and can be rendered in a variety of formats using graphviz's dot program.\n\n\
              \By default, the cost center tree is collapsed to a graph.  To see the cost center tree, i.e.,\n\
              \cost centers with multiple parents appearing as distinct children, use the -u (or --ungraph)\n\
              \option\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 { graphColorCoord  = colorCoord opts,
                                             graphConvertTree = convertToGraph opts,
                                             graphNodeStyle   = nodeStyle 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