import System.Directory
import System.Environment
import System.FilePath.Posix
import System.Process
import Data.Time
import Text.Regex
import Control.Monad
import Data.Maybe
import Data.List
import Data.Function (on)
import System.Console.GetOpt
import System.Exit (ExitCode(..), exitWith)
import System.IO (hFlush, hPutStrLn, stderr, stdout)

data Flag = Help
        | Recursive
        | CSV
        | Bench
        | Range [Int]
        | TimeOut Int
        deriving (Eq, Show)

options :: [OptDescr Flag]
options = [ Option ['h'] ["help"] (NoArg Help)
                   "Show this help message",
            Option ['r'] ["recursive"] (NoArg Recursive)
                   "Traverse recursively",
            Option ['c'] ["csv"] (NoArg CSV)
                    "Print a CSV file",
            Option ['n'] ["range"] (ReqArg (\_ -> Range []) "R")
                    "Only test specific range",
            Option ['t'] ["timeout"] (ReqArg (\_ -> TimeOut 0) "secs")
                    "Only run each test at most this time, not recommended with -b",
            Option ['b'] ["bench"] (NoArg Bench)
                    "Perform a deep benchmark, recommended only with -n"]

parseArgs :: IO ()
parseArgs = do
    argv <- getArgs
    case parse argv of
        (opts, [], [])                      -> lenny opts ["."]
        (opts, dirs, [])
            | Help `elem` opts              -> help
            | CSV `elem` opts && length dirs > 1  -> help
            | otherwise                     -> lenny opts dirs
        (_,[],_)                            -> help
        (_,_,errs)                          -> die errs
  where
    parse argv = getOpt Permute options argv
    header     = "Usage: lenny [-h] [-r] [-b] [-n range] [-t timeout] \n" ++
                    "[dir | -c dir | dir1 dir2]"
    info       = usageInfo header options
    dump       = hPutStrLn stderr
    die errs   = dump (concat errs ++ info) >> exitWith (ExitFailure 1)
    help       = dump info                  >> exitWith ExitSuccess

main = parseArgs

lenny :: [Flag] -> [FilePath] -> IO ()
lenny opts [dir] = do
    bins <- getBins dir
    timings <- lennyTime bins
    if CSV `elem` opts then
        lennyCSV timings
    else
        lennyTerm timings

lenny opts (x:y:_) = print "Unimplemented =("

-- Traverse directory x
getBins :: FilePath -> IO [(FilePath, Int)]
getBins target = do
    -- Find files to be inspected
    setCurrentDirectory target
    files <- getDirectoryContents target
    dir <- getCurrentDirectory
    let abs = map ((dir ++ "/") ++) files
    noDirs <- filterM doesFileExist abs

    -- Determine which files are possible lenny
    -- binaries
    let problemNums = map problemNr noDirs
        filesNums = zipWith (,) noDirs problemNums
        validFiles = filter (isJust . snd) filesNums
        convert = (read . head . fromJust) :: Maybe [String] -> Int
        fileNr = map (\(f,n) -> (f, (convert n))) validFiles 

    -- Determine if these files are executable
    -- Sort these that are in order of problem nr
    permissions <- mapM getPermissions (map fst fileNr)
    let filesPerm = zipWith (,) fileNr permissions
        exe = filter (executable . snd) filesPerm
        binaries = map fst exe
        sortedBins = sortBy (compare `on` snd) binaries 

    return sortedBins

problemNr :: FilePath -> Maybe [String]
problemNr s = matchRegex reg s
    where
    reg = mkRegex "^.*/[A-Za-z]*0*([1-9]+)(\\..*)?$"
    -- ^.*/ directory structure
    -- [A-Za-z]* beginning of filename
    -- 0* optional zero padding
    -- ([1-9]+) the problem number
    -- (\\..*)?$ optional file suffix
 
lennyTime :: [(FilePath,Int)] -> IO [((FilePath,Int), (String, NominalDiffTime))]
lennyTime bins = mapM timing bins
    
timing (bin, num) = do
    start <- getCurrentTime
    output <- readProcess bin [] ""
    stop <- getCurrentTime
    let timing = diffUTCTime stop start
        out = head . lines $ output
    return ((bin,num),(out,timing))

-- Formatting / Printing functions
    
-- Prints out the results as a CSV file
-- Only keeps the first occurence of the
-- problem file
lennyCSV :: [((FilePath, Int),(String,NominalDiffTime))] -> IO ()
lennyCSV bins = do
    let uniqBins = nubBy (\((_,a),_) ((_,b),_) -> a == b) bins
    putStrLn $ intercalate "\n" . map formatCSV $ uniqBins 

formatCSV ((_,n),(_,t)) = (show n) ++ "," ++ ((init . show) t)

-- Prints out the results prettyprinted
-- Allow multiple occurences of the
-- same euler problem
lennyTerm :: [((FilePath, Int),(String,NominalDiffTime))] -> IO ()
lennyTerm bins = do
    mapM_ prettyPrint bins
    let timings = map (snd . snd) bins
        total = sum timings
        numOfTests = genericLength bins
    putStrLn ""
    putStrLn $ "Total running time: " ++ (show total)
    putStrLn $ "Average : " ++ (show (total / numOfTests))

prettyPrint :: ((FilePath, Int),(String,NominalDiffTime)) -> IO ()
prettyPrint ((bin, num), (out,t)) = do
    -- pretty print
    let filename = takeFileName bin
        r = mkRegex (show num) 
        colorized = subRegex r filename ("\x1b[31m" ++ (show num) ++ "\x1b[0m")
    putStr $ colorized ++ " -> "
    putStrLn $ "\x1b[32m" ++ out ++ "\x1b[0m" ++ " : " ++ (show t)
