{-# LANGUAGE OverloadedStrings #-}

module Skeletest.Internal.Preprocessor (
  processFile,
) where

import Control.Monad (guard)
import Data.Char (isDigit, isLower, isUpper)
import Data.List (sort)
import Data.Maybe (mapMaybe)
import Data.Text (Text)
import Data.Text qualified as Text
import System.Directory (doesDirectoryExist, listDirectory)
import System.FilePath (makeRelative, splitExtensions, takeDirectory, (</>))
import UnliftIO.Exception (throwIO)

import Skeletest.Internal.Constants (mainFileSpecsListIdentifier)
import Skeletest.Internal.Error (SkeletestError (..))

-- | Preprocess the given Haskell file. See Main.hs
processFile :: FilePath -> Text -> IO Text
processFile :: FilePath -> Text -> IO Text
processFile FilePath
path Text
file = do
  file' <- if Text -> Bool
isMain Text
file then FilePath -> Text -> IO Text
updateMainFile FilePath
path Text
file else Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
file
  pure
    . addLine pluginPragma
    . addLine linePragma
    $ file'
  where
    addLine :: a -> a -> a
addLine a
line a
f = a
line a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
"\n" a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
f
    quoted :: a -> a
quoted a
s = a
"\"" a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
s a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
"\""

    pluginPragma :: Text
pluginPragma = Text
"{-# OPTIONS_GHC -fplugin=Skeletest.Internal.Plugin #-}"
    linePragma :: Text
linePragma =
      -- this is needed to tell GHC to use original path in error messages
      Text
"{-# LINE 1 " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
forall {a}. (Semigroup a, IsString a) => a -> a
quoted (FilePath -> Text
Text.pack FilePath
path) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" #-}"

isMain :: Text -> Bool
isMain :: Text -> Bool
isMain Text
file =
  case (Text -> Maybe Text) -> [Text] -> [Text]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Text -> Maybe Text
getModuleName ([Text] -> [Text]) -> [Text] -> [Text]
forall a b. (a -> b) -> a -> b
$ Text -> [Text]
Text.lines Text
file of
    -- there was a module line
    [Text
name] -> Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"Main"
    -- there were no module lines, it's the main module
    [] -> Bool
True
    -- something else? just silently ignore it
    [Text]
_ -> Bool
False
  where
    getModuleName :: Text -> Maybe Text
getModuleName Text
s =
      case Text -> [Text]
Text.words Text
s of
        Text
"module" : Text
name : [Text]
_ -> Text -> Maybe Text
forall a. a -> Maybe a
Just Text
name
        [Text]
_ -> Maybe Text
forall a. Maybe a
Nothing

updateMainFile :: FilePath -> Text -> IO Text
updateMainFile :: FilePath -> Text -> IO Text
updateMainFile FilePath
path Text
file = do
  modules <- FilePath -> IO [(FilePath, Text)]
findTestModules FilePath
path
  either throwIO pure $
    pure file
      >>= insertImports modules
      >>= pure . addSpecsList modules

-- | Find all test modules using the given path to the Main module.
--
-- >>> findTestModules "test/Main.hs"
-- ["My.Module.Test1", "My.Module.Test2", ...]
findTestModules :: FilePath -> IO [(FilePath, Text)]
findTestModules :: FilePath -> IO [(FilePath, Text)]
findTestModules FilePath
path = (FilePath -> Maybe (FilePath, Text))
-> [FilePath] -> [(FilePath, Text)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe FilePath -> Maybe (FilePath, Text)
toTestModule ([FilePath] -> [(FilePath, Text)])
-> IO [FilePath] -> IO [(FilePath, Text)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO [FilePath]
listDirectoryRecursive FilePath
testDir
  where
    testDir :: FilePath
testDir = FilePath -> FilePath
takeDirectory FilePath
path

    toTestModule :: FilePath -> Maybe (FilePath, Text)
toTestModule FilePath
fp = do
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (FilePath
fp FilePath -> FilePath -> Bool
forall a. Eq a => a -> a -> Bool
/= FilePath
path)
      (fpNoExt, ".hs") <- (FilePath, FilePath) -> Maybe (FilePath, FilePath)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((FilePath, FilePath) -> Maybe (FilePath, FilePath))
-> (FilePath, FilePath) -> Maybe (FilePath, FilePath)
forall a b. (a -> b) -> a -> b
$ FilePath -> (FilePath, FilePath)
splitExtensions FilePath
fp
      guard ("Spec" `Text.isSuffixOf` Text.pack fpNoExt)
      name <- moduleNameFromPath $ Text.pack $ makeRelative testDir fpNoExt
      pure (fp, name)

    moduleNameFromPath :: Text -> Maybe Text
moduleNameFromPath = ([Text] -> Text) -> Maybe [Text] -> Maybe Text
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Text -> [Text] -> Text
Text.intercalate Text
".") (Maybe [Text] -> Maybe Text)
-> (Text -> Maybe [Text]) -> Text -> Maybe Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Maybe Text) -> [Text] -> Maybe [Text]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Text -> Maybe Text
validateModuleName ([Text] -> Maybe [Text])
-> (Text -> [Text]) -> Text -> Maybe [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => Text -> Text -> [Text]
Text -> Text -> [Text]
Text.splitOn Text
"/"

    -- https://www.haskell.org/onlinereport/syntax-iso.html
    -- large { small | large | digit | ' }
    validateModuleName :: Text -> Maybe Text
validateModuleName Text
name = do
      (first, rest) <- Text -> Maybe (Char, Text)
Text.uncons Text
name
      guard $ isUpper first
      guard $ Text.all (\Char
c -> Char -> Bool
isUpper Char
c Bool -> Bool -> Bool
|| Char -> Bool
isLower Char
c Bool -> Bool -> Bool
|| Char -> Bool
isDigit Char
c Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\'') rest
      pure name

addSpecsList :: [(FilePath, Text)] -> Text -> Text
addSpecsList :: [(FilePath, Text)] -> Text -> Text
addSpecsList [(FilePath, Text)]
testModules Text
file =
  [Text] -> Text
Text.unlines
    [ Text
file
    , Text
mainFileSpecsListIdentifier Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" :: [(FilePath, String, Spec)]"
    , Text
mainFileSpecsListIdentifier Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" = " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [(Text, Text, Text)] -> Text
renderSpecList [(Text, Text, Text)]
specsList
    ]
  where
    specsList :: [(Text, Text, Text)]
specsList =
      [ (Text -> Text
forall {a}. (Semigroup a, IsString a) => a -> a
quote (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ FilePath -> Text
Text.pack FilePath
fp, Text -> Text
forall {a}. (Semigroup a, IsString a) => a -> a
quote Text
modName, Text
modName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
".spec")
      | (FilePath
fp, Text
modName) <- [(FilePath, Text)]
testModules
      ]
    quote :: a -> a
quote a
s = a
"\"" a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
s a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
"\""
    renderSpecList :: [(Text, Text, Text)] -> Text
renderSpecList [(Text, Text, Text)]
xs = Text
"[" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> (Text -> [Text] -> Text
Text.intercalate Text
", " ([Text] -> Text)
-> ([(Text, Text, Text)] -> [Text]) -> [(Text, Text, Text)] -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, Text, Text) -> Text) -> [(Text, Text, Text)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Text, Text) -> Text
forall {a}. (Semigroup a, IsString a) => (a, a, a) -> a
renderSpecInfo) [(Text, Text, Text)]
xs Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"]"
    renderSpecInfo :: (a, a, a) -> a
renderSpecInfo (a
fp, a
name, a
spec) = a
"(" a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
fp a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
", " a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
name a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
", " a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
spec a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
")"

-- | Add imports after the Skeletest.Main import, which should always be present in the Main module.
insertImports :: [(FilePath, Text)] -> Text -> Either SkeletestError Text
insertImports :: [(FilePath, Text)] -> Text -> Either SkeletestError Text
insertImports [(FilePath, Text)]
testModules Text
file =
  let ([Text]
pre, [Text]
post) = (Text -> Bool) -> [Text] -> ([Text], [Text])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break Text -> Bool
isSkeletestImport ([Text] -> ([Text], [Text])) -> [Text] -> ([Text], [Text])
forall a b. (a -> b) -> a -> b
$ Text -> [Text]
Text.lines Text
file
   in if [Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Text]
post
        then SkeletestError -> Either SkeletestError Text
forall a b. a -> Either a b
Left (SkeletestError -> Either SkeletestError Text)
-> SkeletestError -> Either SkeletestError Text
forall a b. (a -> b) -> a -> b
$ Text -> SkeletestError
CompilationError Text
"Could not find Skeletest.Main import in Main module"
        else Text -> Either SkeletestError Text
forall a. a -> Either SkeletestError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> Either SkeletestError Text)
-> ([Text] -> Text) -> [Text] -> Either SkeletestError Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Text
Text.unlines ([Text] -> Either SkeletestError Text)
-> [Text] -> Either SkeletestError Text
forall a b. (a -> b) -> a -> b
$ [Text]
pre [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> [Text]
importTests [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> [Text]
post
  where
    isSkeletestImport :: Text -> Bool
isSkeletestImport Text
line =
      case Text -> [Text]
Text.words Text
line of
        Text
"import" : Text
"Skeletest.Main" : [Text]
_ -> Bool
True
        [Text]
_ -> Bool
False

    importTests :: [Text]
importTests =
      [ Text
"import qualified " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name
      | (FilePath
_, Text
name) <- [(FilePath, Text)]
testModules
      ]

{----- Helpers -----}

listDirectoryRecursive :: FilePath -> IO [FilePath]
listDirectoryRecursive :: FilePath -> IO [FilePath]
listDirectoryRecursive FilePath
fp = ([[FilePath]] -> [FilePath]) -> IO [[FilePath]] -> IO [FilePath]
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([FilePath] -> [FilePath]
forall a. Ord a => [a] -> [a]
sort ([FilePath] -> [FilePath])
-> ([[FilePath]] -> [FilePath]) -> [[FilePath]] -> [FilePath]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[FilePath]] -> [FilePath]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat) (IO [[FilePath]] -> IO [FilePath])
-> ([FilePath] -> IO [[FilePath]]) -> [FilePath] -> IO [FilePath]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FilePath -> IO [FilePath]) -> [FilePath] -> IO [[FilePath]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (FilePath -> IO [FilePath]
go (FilePath -> IO [FilePath])
-> (FilePath -> FilePath) -> FilePath -> IO [FilePath]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FilePath
fp FilePath -> FilePath -> FilePath
</>)) ([FilePath] -> IO [FilePath]) -> IO [FilePath] -> IO [FilePath]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< FilePath -> IO [FilePath]
listDirectory FilePath
fp
  where
    go :: FilePath -> IO [FilePath]
go FilePath
child = do
      isDir <- FilePath -> IO Bool
doesDirectoryExist FilePath
child
      if isDir
        then listDirectoryRecursive child
        else pure [child]