{-# LANGUAGE PatternSynonyms #-}
module Main where

import Prelude hiding (splitAt)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Short as BSS
import qualified Data.Sequence as Seq
import Data.Sequence (Seq, (<|), pattern (:<|))
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import System.Environment
import System.Exit

import Brick
import Brick.Widgets.List
import Graphics.Vty.Attributes
import Graphics.Vty.Input.Events (Event(..), Key(..))

import Parse
import Prof


data Name = NList
  deriving (Show, Eq, Ord)

-- Row indent subtree
data Row = Row Int (ProfTree Bool)

data AppState
    = AppState { aList :: GenericList Name Seq Row }

initAppState :: Prof () -> AppState
initAppState doc =
    let Prof tree = fmap (const False) doc
    in AppState (list NList (Seq.singleton (Row 0 tree)) 1)

rowEntry :: Row -> Entry
rowEntry (Row _ (Node e _ _)) = e
rowEntry (Row _ (Leaf e)) = e

buildSubtree :: Int -> ProfTree Bool -> [Row]
buildSubtree indent tree@(Leaf _) = [Row indent tree]
buildSubtree indent tree@(Node _ False _) = [Row indent tree]
buildSubtree indent tree@(Node _ True ch) =
    Row indent tree : concatMap (buildSubtree (indent + 1)) ch

draw :: AppState -> [Widget Name]
draw app = [renderList func True (aList app)]
  where func selected row@(Row indent _) =
          let selfunc1 = if selected then withAttr (attrName "selectedRow") else id
              selfunc2 = if selected then withAttr (attrName "selected") else id
              entry = rowEntry row
              TimeShare inhTime inhAlloc = eInherited entry
          in selfunc1 $
              hBox [txt (T.replicate indent (T.singleton ' '))
                   ,selfunc2 (txt (mconcat
                       [T.decodeUtf8 (BSS.fromShort (eCS entry))
                       ,T.pack "  "
                       ,T.decodeUtf8 (BSS.fromShort (eModule entry))]))
                   ,vLimit 1 (fill ' ')
                   ,txt (T.pack (show inhTime ++ " " ++ show inhAlloc))]

handleEvent :: AppState -> BrickEvent Name e -> EventM Name (Next AppState)
handleEvent app (VtyEvent (EvKey (KChar 'q') _)) = halt app
handleEvent app (VtyEvent event) = do
    list' <- handleListEventVi fallback event (aList app)
    continue (app { aList = list' })
  where
    fallback (EvKey (KChar ' ') []) l =
        case listSelectedElement l of
          Just (_, Row _ (Leaf _)) -> return l
          Just (index, Row indent (Node entry False ch)) ->
              let (preelts, _ :<| postelts) = splitAt index (listElements l)
                  mid = Seq.fromList (buildSubtree indent (Node entry True ch))
              in return (listReplace (preelts <> mid <> postelts) (Just index) l)
          Just (index, Row indent (Node entry True ch)) ->
              let (preelts, elts) = splitAt index (listElements l)
                  sublen = length (buildSubtree indent (Node entry True ch))
                  postelts = Seq.drop sublen elts
              in return (listReplace (preelts <> (Row indent (Node entry False ch) <| postelts)) (Just index) l)
          Nothing -> return l
    fallback e l = handleListEvent e l
handleEvent app _ = continue app

main :: IO ()
main = do
    args <- getArgs
    inputFile <- case args of
      ["-v"] -> do
          putStrLn "ghc-prof-reader by Tom Smeding"
          putStrLn "Usage: ghc-prof-reader <file.prof>"
          exitSuccess
      [fname] -> BS.readFile fname
      _ -> die "Usage: ghc-prof-reader <file.prof>"

    putStrLn "Reading..."
    doc <- case parseProf inputFile of
      Right doc -> return doc
      Left err -> die err
    doc `seq` putStrLn "Imported"

    _ <- defaultMain
        (App { appDraw = draw
             , appChooseCursor = \_ l -> if null l then Nothing else Just (head l)
             , appHandleEvent = handleEvent
             , appStartEvent = return
             , appAttrMap = \_ -> attrMap defAttr
                    [(attrName "selected", withStyle defAttr standout)
                    ,(attrName "selectedRow", withStyle defAttr underline)] })
        (initAppState doc)
    return ()