module Main where

import qualified GF
import Test.Tasty.Bench

import Control.Exception (catch)
import Control.Monad (filterM, unless, when)
import Data.Char (isAlphaNum, toLower, toUpper)
import System.Directory
       ( createDirectoryIfMissing
       , doesDirectoryExist
       , doesFileExist
       , getCurrentDirectory
       , removeFile
       )
import System.Environment (lookupEnv, withArgs)
import System.FilePath ((</>), (<.>), splitSearchPath)
import System.FilePath (takeBaseName)
import System.IO.Error (isDoesNotExistError)

data BenchConfig = BenchConfig
  { benchLabel      :: String
  , benchOutputName :: String
  , benchGfoDir     :: FilePath
  , benchPgfDir     :: FilePath
  , benchModules    :: [FilePath]
  , benchImportDirs :: [FilePath]
  }

main :: IO ()
main = do
  configs <- readBenchConfigs
  mapM_ validateBenchConfig configs
  mapM_ prepareOutputDirs configs
  defaultMain
    [ bgroup "GF compiler"
      [ bench (benchLabel config) $ whnfIO (runCompiler config)
      | config <- configs
      ]
    ]

readBenchConfigs :: IO [BenchConfig]
readBenchConfigs = do
  cwd <- getCurrentDirectory
  rglDir <- lookupEnvDefault "GF_BENCH_RGL_DIR" (".." </> "gf-rgl")
  let benchDir = cwd </> "profile" </> "rgl-bench"
      defaultImportDirs = [rglDir </> "src" </> "api"]
  outputName <- lookupEnvDefault "GF_BENCH_NAME" "RGLBench"
  gfoDir <- lookupEnvDefault "GF_BENCH_GFO_DIR" (benchDir </> "gfo")
  pgfDir <- lookupEnvDefault "GF_BENCH_PGF_DIR" (benchDir </> "pgf")
  importDirs <- lookupEnvList "GF_BENCH_IMPORT_DIRS" defaultImportDirs
  customModules <- lookupEnvList "GF_BENCH_MODULES" []
  return $
    if null customModules
      then map (defaultBenchConfig rglDir outputName gfoDir pgfDir importDirs)
               defaultRGLBenchmarks
      else customBenchConfigs outputName gfoDir pgfDir importDirs customModules

defaultRGLBenchmarks :: [(String, String, FilePath)]
defaultRGLBenchmarks =
  [ ("English", "english", "english" </> "LangEng.gf")
  , ("German",  "german",  "german"  </> "LangGer.gf")
  , ("French",  "french",  "french"  </> "LangFre.gf")
  , ("Russian", "russian", "russian" </> "LangRus.gf")
  , ("Finnish", "finnish", "finnish" </> "LangFin.gf")
  ]

defaultBenchConfig
  :: FilePath -> String -> FilePath -> FilePath -> [FilePath]
  -> (String, String, FilePath)
  -> BenchConfig
defaultBenchConfig rglDir outputName gfoDir pgfDir importDirs (language, slug, modulePath) =
  BenchConfig
    { benchLabel = "RGL " ++ language ++ " compile/link"
    , benchOutputName = outputNameFor outputName slug
    , benchGfoDir = gfoDir </> slug
    , benchPgfDir = pgfDir </> slug
    , benchModules = [rglDir </> "src" </> modulePath]
    , benchImportDirs = importDirs
    }

customBenchConfigs :: String -> FilePath -> FilePath -> [FilePath] -> [FilePath] -> [BenchConfig]
customBenchConfigs outputName gfoDir pgfDir importDirs modules =
  [ BenchConfig
    { benchLabel = "RGL " ++ takeBaseName modulePath ++ " compile/link"
    , benchOutputName = caseOutputName modulePath
    , benchGfoDir = caseDir gfoDir modulePath
    , benchPgfDir = caseDir pgfDir modulePath
    , benchModules = [modulePath]
    , benchImportDirs = importDirs
    }
  | modulePath <- modules
  ]
  where
    multiple = length modules > 1

    caseOutputName modulePath
      | multiple  = outputNameFor outputName (takeBaseName modulePath)
      | otherwise = outputName

    caseDir dir modulePath
      | multiple  = dir </> slugify (takeBaseName modulePath)
      | otherwise = dir

lookupEnvDefault :: String -> String -> IO String
lookupEnvDefault key fallback = do
  value <- lookupEnv key
  return $ case value of
    Just xs | not (null xs) -> xs
    _                      -> fallback

lookupEnvList :: String -> [FilePath] -> IO [FilePath]
lookupEnvList key fallback = do
  value <- lookupEnv key
  return $ case value of
    Just xs | not (null xs) -> parsePathList xs
    _                      -> fallback

parsePathList :: String -> [FilePath]
parsePathList =
  filter (not . null) . concatMap splitSearchPath . words

validateBenchConfig :: BenchConfig -> IO ()
validateBenchConfig config = do
  when (null (benchModules config)) $
    fail "GF_BENCH_MODULES did not contain any input modules."
  missingModules <- filterM (fmap not . doesFileExist) (benchModules config)
  unless (null missingModules) $
    fail $ "Missing GF benchmark module(s): " ++ unwords missingModules
  missingImportDirs <- filterM (fmap not . doesDirectoryExist) (benchImportDirs config)
  unless (null missingImportDirs) $
    fail $ "Missing GF benchmark import dir(s): " ++ unwords missingImportDirs

prepareOutputDirs :: BenchConfig -> IO ()
prepareOutputDirs config = do
  createDirectoryIfMissing True (benchGfoDir config)
  createDirectoryIfMissing True (benchPgfDir config)

runCompiler :: BenchConfig -> IO ()
runCompiler config = do
  prepareOutputDirs config
  removeFileIfExists (benchPgfDir config </> benchOutputName config <.> "pgf")
  withArgs (compilerArgs config) GF.main

compilerArgs :: BenchConfig -> [String]
compilerArgs config =
  [ "--make"
  , "--quiet"
  , "--src"
  , "--preproc=mkPresent"
  , "--gf-lib-path=" ++ benchGfoDir config
  , "--gfo-dir=" ++ benchGfoDir config
  , "--output-dir=" ++ benchPgfDir config
  , "--name=" ++ benchOutputName config
  ]
  ++ concatMap (\dir -> ["-i", dir]) (benchImportDirs config)
  ++ benchModules config

removeFileIfExists :: FilePath -> IO ()
removeFileIfExists path =
  removeFile path `catch` \err ->
    unless (isDoesNotExistError err) (ioError err)

outputNameFor :: String -> String -> String
outputNameFor outputName slug = outputName ++ capitalize (filter isAlphaNum slug)

capitalize :: String -> String
capitalize []     = []
capitalize (x:xs) = toUpper x : xs

slugify :: String -> String
slugify = map clean
  where
    clean c
      | isAlphaNum c = toLower c
      | otherwise    = '-'