module Test.Tasty.Auto (findTests, showTestDriver) where

import Data.Function (on)
import Data.List (find, isPrefixOf, isSuffixOf, nub, intersperse, groupBy, sortOn, dropWhileEnd)
import Data.Maybe (fromJust)
import System.Directory (getDirectoryContents, doesDirectoryExist)
import Data.Traversable (for)
import System.FilePath ((</>), takeDirectory, pathSeparator, dropExtension)
import Data.Monoid (Endo(..))
import Data.Foldable (fold)

data Generator = Generator
  { genPrefix :: String
  , genImport, genClass :: ShowS
  , genSetup :: Test -> ShowS }

data Test = Test { testModule, testFunction :: String }

str :: String -> ShowS
str = (++)

sp, nl :: ShowS
sp = (' ':)
nl = ('\n':)

tr :: Char -> Char -> String -> String
tr a b = map $ \c -> if c == a then b else c

name, fn :: Test -> ShowS
name = shows . tr '_' ' ' . tail . dropWhile (/= '_') . testFunction
fn t = str (testModule t) . ('.':) . str (testFunction t)

generators :: [Generator]
generators =
  [ Generator { genPrefix = "prop_"
              , genImport = str "import qualified Test.Tasty.QuickCheck as QC\n"
              , genClass  = id
              , genSetup  = \t -> str "pure $ QC.testProperty " . name t . sp . fn t }
  , Generator { genPrefix = "scprop_"
              , genImport = str "import qualified Test.Tasty.SmallCheck as SC\n"
              , genClass  = id
              , genSetup  = \t -> str "pure $ SC.testProperty " . name t . sp . fn t }
  , Generator { genPrefix = "case_"
              , genImport = str "import qualified Test.Tasty.HUnit as HU\n"
              , genClass  = str "class TestCase a where testCase :: String -> a -> IO T.TestTree\n\
                                \instance TestCase (IO ())                      where testCase n = pure . HU.testCase      n\n\
                                \instance TestCase (IO String)                  where testCase n = pure . HU.testCaseInfo  n\n\
                                \instance TestCase ((String -> IO ()) -> IO ()) where testCase n = pure . HU.testCaseSteps n\n"
              , genSetup  = \t -> str "testCase " . name t . sp . fn t }
  , Generator { genPrefix = "spec_"
              , genImport = str "import qualified Test.Tasty.Hspec as HS\n"
              , genClass  = id
              , genSetup  = \t -> str "HS.testSpec " . name t . sp . fn t }
  , Generator { genPrefix = "test_"
              , genImport = id
              , genClass  = str "class TestGroup a where testGroup :: String -> a -> IO T.TestTree\n\
                                \instance TestGroup T.TestTree        where testGroup _ a = pure a\n\
                                \instance TestGroup [T.TestTree]      where testGroup n a = pure $ T.testGroup n a\n\
                                \instance TestGroup (IO T.TestTree)   where testGroup _ a = a\n\
                                \instance TestGroup (IO [T.TestTree]) where testGroup n a = T.testGroup n <$> a\n"
              , genSetup  = \t -> str "testGroup " . name t . sp . fn t } ]

testFileSuffixes :: [String]
testFileSuffixes = (++) <$> ["Spec", "Test"] <*> [".lhs", ".hs"]

getGenerator :: Test -> Generator
getGenerator t = fromJust $ find ((`isPrefixOf` testFunction t) . genPrefix) generators

getGenerators :: [Test] -> [Generator]
getGenerators = map head . groupBy  ((==) `on` genPrefix) . sortOn genPrefix . map getGenerator

showImports :: [String] -> ShowS
showImports = foldEndo . map (\m -> str "import qualified " . str m . nl) . nub

showSetup :: Test -> ShowS -> ShowS
showSetup t var = str "  " . var . str " <- " . genSetup (getGenerator t) t . nl

foldEndo :: (Functor f, Foldable f) => f (a -> a) -> (a -> a)
foldEndo = appEndo . fold . fmap Endo

ingredientImport :: String -> String
ingredientImport = init . dropWhileEnd (/= '.')

ingredients :: [String] -> ShowS
ingredients is = foldEndo (map (\i -> str i . (':':)) is) . str "T.defaultIngredients"

showTestDriver :: String -> [String] -> FilePath -> [Test] -> ShowS
showTestDriver modname is src ts =
  let gs = getGenerators ts; vars = map (str . ('t':) . show) [(0::Int)..] in
    str "{-# LINE 1 " . shows src . str " #-}\n\
        \{-# LANGUAGE FlexibleInstances #-}\n\
        \module " . str modname . str " (main, ingredients, tests) where\n\
        \import Prelude\n\
        \import qualified Test.Tasty as T\n\
        \import qualified Test.Tasty.Ingredients as T\n"
  . foldEndo (map genImport gs)
  . showImports (map ingredientImport is ++ map testModule ts)
  . foldEndo (map genClass gs)
  . str "tests :: IO T.TestTree\n\
        \tests = do\n"
  . foldEndo (zipWith showSetup ts vars)
  . str "  pure $ T.testGroup " . shows src . str " ["
  . foldEndo (intersperse (',':) $ zipWith (curry snd) ts vars)
  . str "]\n"
  . str "ingredients :: [T.Ingredient]\n\
        \ingredients = " . ingredients is . str "\n\
        \main :: IO ()\n\
        \main = tests >>= T.defaultMainWithIngredients ingredients\n"

filesBySuffix :: FilePath -> [String] -> IO [FilePath]
filesBySuffix dir suffixes = do
  entries <- filter (\s -> head s /= '.') <$> getDirectoryContents dir
  found <- for entries $ \entry -> do
    let dir' = dir </> entry
    exists <- doesDirectoryExist dir'
    if exists then map (entry </>) <$> filesBySuffix dir' suffixes else pure []
  pure $ filter (\x -> any (`isSuffixOf` x) suffixes) entries ++ concat found

findTests :: FilePath -> IO [Test]
findTests src = do
  let dir = takeDirectory src
  files <- filesBySuffix dir testFileSuffixes
  concat <$> traverse (\f -> extractTests f <$> readFile (dir </> f)) files

mkTest :: FilePath -> String -> Test
mkTest = Test . tr pathSeparator '.' . dropExtension

extractTests :: FilePath -> String -> [Test]
extractTests file =
    map (mkTest file) . nub
  . filter (\n -> any ((`isPrefixOf` n) . genPrefix) generators)
  . map fst . concatMap lex . lines