module Main (main) where

import Control.Monad (forM_, unless)
import Data.Char (isSpace)
import Data.List (intercalate, isPrefixOf)
import Data.Maybe (listToMaybe, mapMaybe)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Semigroup (Max(Max), (<>), Sum(Sum), mappend, mconcat)
import Data.Set (Set)
import qualified Data.Set as S
import Data.Tree (Tree(..), Forest, flatten)
import System.Environment (getArgs)
import System.Exit (exitFailure)
import System.FilePath ((</>), (<.>), dropExtension)
import System.IO (hPutStr, hPutStrLn, stderr, withFile, IOMode(WriteMode))
import Text.Encoding.Z (zDecodeString, zEncodeString)

import Paths_prof2pretty (getDataFileName)

type SrcLoc  = (Int, Int)
type SrcSpan = (SrcLoc, SrcLoc)
type FileSpan = (String, SrcSpan)

decodeSCC :: String -> Maybe FileSpan
decodeSCC s = case (reverse . take 5 . reverse . split '-') s of
  z:is -> case mapMaybe readMay is of
    [l0,c0,l1,c1] -> Just (zDecodeString z, ((l0, c0), (l1, c1)))
    _ -> Nothing
  _ -> Nothing

notNull :: [a] -> Bool
notNull = not . null

deepMapForest :: (Forest a -> Forest a) -> Tree a -> Tree a
deepMapForest f (Node label forest) = Node label (f (map (deepMapForest f) forest))

split :: Eq a => a -> [a] -> [[a]]
split _ [] = []
split x xs = let (ys, zs) = break (x ==) xs in ys : split x (drop 1 zs)

readMay :: Read a => String -> Maybe a
readMay s = case reads s of
  [(x, "")] -> Just x
  _ -> Nothing

preprocess :: String -> ([String], [String])
preprocess
  = (concatMap (take 1 . words) *** (dropWhile null . dropWhile notNull . dropWhile null))
  . break null . drop 9 . lines

data Entry = Entry
  { eSCC     :: String
  , eTicks   :: !Integer
  , eBytes   :: !Integer
  }
  deriving (Show)

parseLine :: String -> (Int, Entry)
parseLine s = 
  let (indent, rest) = span isSpace s
      [scc,_,_,_,_,_,_,_,ticks,bytes] = words rest
  in  (length indent, Entry scc (read ticks) (read bytes))

type Entries = [(Int, Entry)]

collate :: Entries -> Tree Entry
collate input =
  let ([[result]], []) = collate' (-1) [] input
  in  deepMapForest reverse result

collate' :: Int -> [Forest Entry] -> Entries -> ([Forest Entry], Entries)
collate' pdepth stack [] = (adjust pdepth stack , [])
collate' pdepth stack ((depth, entry) : rest) =
  let siblings : parents = adjust (pdepth - depth) stack
      node = Node{ rootLabel = entry, subForest = [] }
  in  collate' depth ((node : siblings) : parents) rest

adjust :: Int -> [Forest Entry] -> [Forest Entry]
adjust n t = case compare n 0 of
  LT -> adjust (n + 1) ([] : t)
  EQ ->                      t
  GT -> adjust (n - 1) (case t of
    (ts:(t':ts'):ts'') -> (t'{ subForest = subForest t' ++ ts } : ts') : ts''
    _ -> error $ "prof2pretty.adjust: " ++ show (n, t))

parseEntries :: [String] -> Tree Entry
parseEntries = collate . map parseLine

type Proportion = ((Double, Double), (Double, Double))

proportion :: Tree Entry -> Tree (Entry, Proportion)
proportion t =
  let u@Node{ rootLabel = (_, (Sum totalTicks, Sum totalBytes)) } = up t
  in  down totalTicks totalBytes totalTicks totalBytes u
  where
    me e = (Sum (eTicks e), Sum (eBytes e))
    up Node{ rootLabel = e, subForest = [] } =
          Node{ rootLabel = (e, me e), subForest = [] }
    up Node{ rootLabel = e, subForest = ts } =
      let us = map up ts
          them = mconcat (map (snd . rootLabel) us)
      in  Node{ rootLabel = (e, me e `mappend` them), subForest = us }
    down totalTicks totalBytes parentTicks parentBytes
      Node{ rootLabel = (e, (Sum ticks, Sum bytes)), subForest = fs } =
        let ticky = (eTicks e /// totalTicks, ticks /// parentTicks)
            bytey = (eBytes e /// totalBytes, bytes /// parentBytes)
            local = (ticky, bytey)
        in  Node
              { rootLabel = (e, local)
              , subForest = map (down totalTicks totalBytes ticks bytes) fs
              }

(///) :: Integer -> Integer -> Double
a /// b = fromIntegral a / if b > 0 then fromIntegral b else 1

(***) :: (a -> c) -> (b -> d) -> (a, b) -> (c, d)
f *** g = \(a, b) -> (f a, g b)

second :: (b -> d) -> (a, b) -> (a, d)
second g = id *** g

compose :: [a -> a] -> a -> a
compose = foldr (.) id

type Profile = ([String], Map String Proportion)

collect :: Tree (Entry, Proportion) -> Map String Proportion
collect = fmap unM2 . M.fromListWith (<>) . map (eSCC *** m2) . flatten
  where
    unM (Sum a, Max b) = (a, b)
    unM2 (a, b) = (unM a, unM b)
    m2 (a, b) = (m a, m b)
    m (a, b) = (Sum a, Max b)

parseProfile :: String -> Profile
parseProfile = second (collect . proportion . parseEntries) . preprocess

characterize :: String -> [[String]]
characterize = map (map (:[])) . lines

decharacterize :: [[String]] -> String
decharacterize = unlines . map concat

xmlSafe :: String -> String
xmlSafe = concatMap x
  where
    x '&' = "&amp;"
    x '<' = "&lt;"
    x '>' = "&gt;"
    x  c  = [c]

wrapLine :: Int -> [a] -> [a] -> [[[a]]] -> [[[a]]]
wrapLine line pre post
  = prependAt (line,        1) pre
  . appendAt  (line, maxBound) post

prependAt :: (Int, Int) -> [a] -> [[[a]]] -> [[[a]]]
prependAt loc str = at2 loc (str ++)

appendAt :: (Int, Int) -> [a] -> [[[a]]] -> [[[a]]]
appendAt (l, c) str
  | c == 0    = appendAt (l - 1, maxBound) str
  | otherwise = at2      (l, c - 1)    (++ str)

at2 :: (Int, Int) -> ([a] -> [a]) -> [[[a]]] -> [[[a]]]
at2 (line, col) f = at1 line (at1 col f)

at1 :: Int -> ([a] -> [a]) -> [[a]] -> [[a]]
at1 n f xs = case splitAt (n - 1) xs of
  (pre, [     ]) -> pre ++ [f []]
  (pre, at:post) -> pre ++ [f at] ++ post

annotate :: [String] -> String -> String -> String
annotate prof file =
    decharacterize . (foldr f `flip` prof) . map (map xmlSafe) . characterize
  where
    f scc = case decodeSCC scc of
      Nothing -> id
      Just (file', (start@(sl, _), end'@(el', ec')))
        | file /= file' -> id
        | sl == el  -> prependAt start open . appendAt end close
        | otherwise -> prependAt start open . appendAt (sl, maxBound) close
                    .  compose [ wrapLine l open close | l <- [sl + 1 .. el - 1] ]
                    .  prependAt (el, 1) open . appendAt end close
        where
          end@(el, _)
            | ec' == 1  = (el' - 1, maxBound)
            | otherwise = end'
          open  = "<span id=\"" ++ scc ++ "\">"
          close = "</span>"

readProfile :: String -> IO Profile
readProfile f = parseProfile `fmap` readFile f

profileFiles :: Profile -> Set String
profileFiles = S.fromList . mapMaybe (fmap fst . decodeSCC) . M.keys . snd

profileSCCs :: Profile -> [String]
profileSCCs = M.keys . snd

addLineNumbers :: String -> String -> String
addLineNumbers file = unlines . map (\(n, s) ->
  "<span class=\"line-number\">" ++ replicate (6 - length (show n)) ' '
  ++ "<a id=\"" ++ zEncodeString file ++ "-" ++ show n ++ "\" "
  ++ "href=\"#" ++ zEncodeString file ++ "-" ++ show n ++ "\">"
  ++ show n ++ "</a>  </span>" ++ s) . zip [(1 :: Integer)..] . lines

main :: IO ()
main = do
  args <- getArgs
  let standalone = "--standalone" `elem` args
      Just source = listToMaybe . (++ ["."]) . map (drop (length sourcePrefix)) . filter (sourcePrefix `isPrefixOf`) $ args
        where sourcePrefix = "--source="
      mprofF = listToMaybe . filter (not . ("--" `isPrefixOf`)) $ args
  profF <- case mprofF of
    Just p -> return p
    _ -> hPutStrLn stderr "usage: prof2pretty [--standalone] [--source='/path/to/src'] executable.prof" >> exitFailure
  prof <- readProfile profF
  withFile (profF <.> "html") WriteMode $ \h -> do
    let files = S.toList $ profileFiles prof
        title = xmlSafe (dropExtension profF)
        profFJSFile = profF ++ ".js"
        profFJS =
          "var profile =\n{" ++ intercalate "\n,"
            [ show s ++ ":" ++ show [a, b, c, d]
            | (s, ((a, b), (c, d))) <- M.toList (snd prof)
            , a + b + c + d > 0
            ] ++ "\n};\n"
    unless standalone $ writeFile profFJSFile profFJS
    profCSSFile <- getDataFileName "prof2pretty.css"
    profJSFile <- getDataFileName "prof2pretty.js"
    jqueryUICSSFile <- getDataFileName "jquery-ui-1.8.21.custom.css"
    jqueryJSFile <- getDataFileName "jquery-1.7.2.min.js"
    jqueryUIJSFile <- getDataFileName "jquery-ui-1.8.21.custom.min.js"
    let readFile' = if standalone then readFile else const (return "")
    profCSS <- readFile' profCSSFile
    profJS <- readFile' profJSFile
    jqueryUICSS <- readFile' jqueryUICSSFile
    jqueryJS <- readFile' jqueryJSFile
    jqueryUIJS <- readFile' jqueryUIJSFile
    let css file text
          | standalone = "<style type=\"text/css\">/* <![CDATA[ */\n" ++ text ++ "\n/* ]]> */</style>"
          | otherwise  = "<link type=\"text/css\" href=\"" ++ xmlSafe file ++ "\" rel=\"stylesheet\" />"
        js file text
          | standalone = "<script type=\"text/javascript\">/* <![CDATA[ */\n" ++ text ++ "\n/* ]]> */</script>"
          | otherwise  = "<script type=\"text/javascript\" src=\"" ++ xmlSafe file ++ "\"></script>"
    hPutStr h . unlines $
      [ "<!DOCTYPE html>\n<html><head><title>" ++ title ++ "</title>"
      , css jqueryUICSSFile jqueryUICSS
      , js jqueryJSFile jqueryJS
      , js jqueryUIJSFile jqueryUIJS
      , css profCSSFile profCSS
      , js profFJSFile profFJS
      , js profJSFile profJS
      , "</head><body onload=\"go()\">"
      , "<div id=\"ui\">\n<div class=\"left\"><p>mode</p><form id=\"ui-mode\">"
      , "<input type=radio id=\"ui-ticks\" name=\"ui-mode\" /><label for=\"ui-ticks\">ticks</label>"
      , "<input type=radio id=\"ui-bytes\" name=\"ui-mode\" /><label for=\"ui-bytes\">bytes</label>"
      , "</form></div>"
      , "<div class=\"left\"><p>limit</p><div id=\"ui-limit\"></div></div>"
      , "<div class=\"left\"><p>gamma</p><div id=\"ui-gamma\"></div></div>"
      , "</div>\n<div id=\"main\">\n<h1>" ++ title ++ "</h1>\n<h2>Summary</h2><ul>"
      ] ++
      [ case decodeSCC scc of
          Nothing ->
            "<li><a id=\"top-" ++ xmlSafe scc ++ "\">" ++ xmlSafe scc ++ "</a></li>"
          Just (file, ((sl, sc), (el, ec))) ->
            "<li><a id=\"top-" ++ xmlSafe scc ++ "\" href=\"#" ++ xmlSafe scc ++ "\">" ++
            xmlSafe file ++ ":" ++ show sl ++ ":" ++ show sc ++ "-" ++ show el ++ ":" ++ show ec ++
            "</a></li>"
      | scc <- fst prof
      ] ++ [ "</ul>", "<h2>Details</h2>" ]
    forM_ files $ \file -> do
      hPutStr h $ "<div id=\"" ++ zEncodeString file ++ "\"><h3>" ++ xmlSafe file ++ "</h3>\n<pre>"
      hPutStr h . addLineNumbers file . annotate (profileSCCs prof) file =<< readFile (source </> file)
      hPutStr h $ "</pre></div>\n"
    hPutStrLn h "</div>\n</body></html>"