{-# LANGUAGE ScopedTypeVariables #-}

import Control.Monad
import Data.Maybe
import Data.List (sortBy)
import Data.Ord
import System.IO (stdout)

import Options.Applicative
import Prettyprinter as PP
import Prettyprinter.Render.Terminal as PP
import qualified System.Console.ANSI

import Text.Regex.TDFA
import Text.Regex.TDFA.Common (Regex)
import Text.Regex.TDFA.Text ()

import GhcDump.ToHtml
import GhcDump.Pretty
import GhcDump.Util
import GhcDump.Ast

data Column a = Col { colWidth :: Int
                    , colHeader :: String
                    , colGet :: a -> Doc AnsiStyle
                    }

type Table a = [Column a]

renderTable :: forall a. Table a -> [a] -> Doc AnsiStyle
renderTable cols xs =
         row (PP.annotate PP.bold . pretty . colHeader)
    <$$> vcat [ row (flip colGet x) | x <- xs ]
  where
    row :: (Column a -> Doc AnsiStyle) -> Doc AnsiStyle
    row toCell = go cols
      where
        go :: [Column a] -> Doc AnsiStyle
        go []           = mempty
        go [col]        = align $ toCell col
        go (col : rest) = fillBreak (colWidth col) (align $ toCell col) PP.<+> go rest

filterBindings :: Regex -> Module -> Module
filterBindings re m =
    m { moduleTopBindings = mapMaybe filterTopBinding $ moduleTopBindings m }
  where
    filterTopBinding b'@(NonRecTopBinding b _ _)
      | nameMatches b  = Just b'
      | otherwise      = Nothing
    filterTopBinding (RecTopBinding bs)
      | not $ null bs' = Just $ RecTopBinding bs'
      | otherwise      = Nothing
      where bs' = filter (\(b,_,_) -> nameMatches b) bs

    nameMatches :: Binder -> Bool
    nameMatches b = matchTest re (binderUniqueName b)

renderIOTerm :: PP.LayoutOptions -> Doc AnsiStyle -> IO ()
renderIOTerm layoutOpts doc = do
    supportsANSI <- System.Console.ANSI.hSupportsANSI System.IO.stdout
    let doc' = if supportsANSI
                 then doc
                 else PP.unAnnotate doc
    PP.renderIO stdout $ PP.layoutPretty layoutOpts doc'

modes :: Parser (IO ())
modes = subparser
     $ mode "show" showMode (progDesc "print Core")
    <> mode "list-bindings" listBindingsMode (progDesc "list top-level bindings, their sizes, and types")
    <> mode "summarize" summarizeMode (progDesc "summarize multiple dump files")
  where
    mode name f opts = command name (info (helper <*> f) opts)

    dumpFile :: Parser FilePath
    dumpFile = argument str (metavar "DUMP FILE" <> help "CBOR dump file")

    filterCond :: Parser (Module -> Module)
    filterCond =
        fmap (maybe id filterBindings)
        $ option (str >>= fmap Just . makeRegexM')
                 (short 'f' <> long "filter" <> value Nothing <> help "filter bindings by name")
      where
        makeRegexM' = makeRegexM :: String -> ReadM Regex

    bindingsSort :: Parser (Module -> Module)
    bindingsSort =
        option (str >>= readSortField)
               (long "sort" <> short 's' <> value id
                <> help "Sort by (accepted values: none, terms, types, coercions, type)")
      where
        readSortField "none"      = return $ id
        readSortField "terms"     = return $ onBinds $ sortBy (flip $ comparing $ csTerms . getStats)
        readSortField "types"     = return $ onBinds $ sortBy (flip $ comparing $ csTypes . getStats)
        readSortField "coercions" = return $ onBinds $ sortBy (flip $ comparing $ csCoercions . getStats)
        readSortField "type"      = return $ onBinds $ sortBy (comparing $ binderType . unBndr . getBinder)
        readSortField f           = fail $ "unknown sort field "++f

        onBinds :: ([(Binder, CoreStats, Expr)] -> [(Binder, CoreStats, Expr)])
                -> Module -> Module
        onBinds f mod = mod { moduleTopBindings = [RecTopBinding $ f $ moduleBindings mod] }

    prettyOpts :: Parser PrettyOpts
    prettyOpts =
        PrettyOpts
          <$> switch (short 'u' <> long "show-uniques" <> help "Show binder uniques")
          <*> switch (short 'i' <> long "show-idinfo" <> help "Show IdInfo of bindings")
          <*> switch (short 'T' <> long "show-let-types" <> help "Show type signatures for let-bound binders")
          <*> switch (short 'U' <> long "show-unfoldings" <> help "Show unfolding templates")

    showMode =
        run <$> filterCond <*> bindingsSort <*> prettyOpts <*> html <*> dumpFile
      where
        html = switch (short 'H' <> long "html" <> help "Render to HTML")
        run filterFn sortBindings opts html fname = do
            dump <- sortBindings . filterFn <$> GhcDump.Util.readDump fname
            if html
              then writeFile "out.html" $ show $ topBindingsToHtml (moduleTopBindings dump)
              else renderIOTerm PP.defaultLayoutOptions $ pprModule opts dump

    listBindingsMode =
        run <$> filterCond <*> bindingsSort <*> prettyOpts <*> dumpFile
      where
        run filterFn sortBindings opts fname = do
            dump <- sortBindings . filterFn <$> GhcDump.Util.readDump fname
            let table = [ Col 30 "Name"   (pprBinder opts . getBinder)
                        , Col 6  "Terms"  (pretty . csTerms . getStats)
                        , Col 6  "Types"  (pretty . csTypes . getStats)
                        , Col 6  "Coerc." (pretty . csCoercions . getStats)
                        , Col 3000 "Type"  (pprType opts . binderType . unBndr . getBinder)
                        ]
            let layoutOpts = PP.defaultLayoutOptions { layoutPageWidth = Unbounded }
            renderIOTerm layoutOpts $ vcat
              [ pretty (modulePhase dump)
              , renderTable table (moduleBindings dump)
              ]

    summarizeMode =
        run <$> some dumpFile
      where
        run fnames = do
            mods <- mapM (\fname -> do mod <- readDump fname
                                       return (fname, mod)) fnames
            let totalSize :: Module -> CoreStats
                totalSize = foldMap getStats . moduleBindings
            let table = [ Col 35 "Name" (pretty . fst)
                        , Col 8  "Terms" (pretty . csTerms . totalSize . snd)
                        , Col 8  "Types" (pretty . csTypes . totalSize . snd)
                        , Col 8  "Coerc." (pretty . csCoercions . totalSize . snd)
                        , Col 35 "Previous phase" (pretty . modulePhase . snd)
                        ]
            print $ renderTable table mods


getBinder (b,_,_) = b
getStats (_,s,_) = s
getRHS (_,_,e) = e

main :: IO ()
main = join $ execParser $ info (helper <*> modes) mempty