import           Control.Monad                          (when)
import           Data.Char                              (toUpper)
import           Data.Dynamic                           (fromDynamic)
import           Data.List                              (intercalate, isPrefixOf, isSuffixOf)
import           Data.Maybe                             (catMaybes)
import           Data.Monoid                            (Monoid (mappend, mempty, mconcat))
import qualified DynFlags                       as G    (defaultDynFlags, PackageFlag (ExposePackage, HidePackage, IgnorePackage))
import qualified GHC                            as G    (GhcMonad, packageFlags, getModuleGraph, ms_mod, setContext, moduleName,
                                                         moduleNameString, targetAllowObjCode, targetContents, ghcLink, flags,
                                                         hscTarget, load, modInfoTyThings, getModuleInfo, setTargets,
                                                         dynCompileExpr, runGhc, defaultErrorHandler, getSessionDynFlags,
                                                         setSessionDynFlags, ghcMode, GhcMode (CompManager),
                                                         LoadHowMuch (LoadAllTargets), SuccessFlag (Failed, Succeeded),
                                                         TyThing (AnId), GhcLink (LinkInMemory), TargetId (TargetFile),
                                                         Target (Target, targetId), HscTarget (HscInterpreted),
                                                         DynFlag (Opt_ImplicitImportQualified))
import qualified GHC.Paths                      as G    (libdir)
import qualified MonadUtils                     as G    (liftIO)
import qualified Name                           as G    (nameOccName)
import qualified OccName                        as G    (occNameString)
import           System.Console.GetOpt                  (usageInfo, getOpt, OptDescr (Option), ArgOrder (ReturnInOrder),
                                                         ArgDescr (ReqArg, NoArg))
import           System.Directory                       (getDirectoryContents)
import           System.Environment                     (getArgs)
import           System.Exit                            (exitWith, exitFailure, exitSuccess, ExitCode (ExitSuccess, ExitFailure))
import           System.IO                              (hPutStrLn, stderr)
import           Test.QuickCheck                        (Arbitrary (arbitrary, coarbitrary), variant, vector)
import qualified Var                            as G    (varName)


{-

This is how to embed QuickCheck properties in your code: make them top-level declarations whose names start with `prop_`.

You don't need to include each property explicitly in a list; that would make it too easy to add a property but forget to run it.

Instead, `hstest` looks through your program and builds the list of properties to run itself.

-}

prop_Count_Monoid_leftIdentity, prop_Count_Monoid_rightIdentity :: Count -> Bool
prop_Count_Monoid_associative :: Count -> Count -> Count -> Bool
prop_Count_Monoid_leftIdentity x    = mappend mempty x == x
prop_Count_Monoid_rightIdentity x   = mappend x mempty == x
prop_Count_Monoid_associative x y z = mappend x (mappend y z) == mappend (mappend x y) z


data TestResult = Failed [String]
                | Passed
                | Exhausted Int
                | Unparsed String

reportResult n (Unparsed s)    = concat ["While testing ", n, ":\n", s]
reportResult n (Exhausted c)   = concat ["Exhausted ", n, " after ", show c, " tests\n"]
reportResult n  Passed         = concat ["Passed ", n, "\n"]
reportResult n (Failed params) = unlines (concat ["Failed ", n, ":"] : map f params)
  where f param = concat [" *  ", param]

shouldExplicitlyReportResult Passed = False
shouldExplicitlyReportResult _      = True


data Count = Count [Int] deriving (Eq, Show)

instance Monoid Count
  where mempty = Count (replicate cResultIndicies 0)
        mappend (Count xs) (Count ys) = Count (zipWith (+) xs ys)

instance Arbitrary Count
  where arbitrary = fmap Count $ vector cResultIndicies
        coarbitrary (Count xs) z = foldr variant z xs

countResultI i = Count $ map f [0 .. cResultIndicies - 1]
  where f i' = if i == i' then 1 else 0

reportCount s' (Count cs) = concat [s', ": ", if null s then "0 tests!" else s]
  where s = ucfirst $ intercalate ", " . catMaybes . zipWith f ss $ cs
        f _            0 = Nothing
        f (s, (s', _)) 1 = Just $ concat [s, " 1 ", s']
        f (s, (_, s')) c = Just $ concat [s, " ", show c, " ", s']
        ss = [("failed", sProperties),
              ("couldn't parse results of", sProperties),
              ("couldn't compile", sFiles),
              ("exhausted arguments while checking", sProperties),
              ("passed", sProperties)]
        sFiles = ("file", "files")
        sProperties = ("property", "properties")
        ucfirst ""       = ""
        ucfirst (ch : s) = toUpper ch : s

putCountLn s c = putStrLn (reportCount s c) >> return c

iFailedProperty : iUnparsedProperty : iUncompiledFile : -- FAIL
        iExhaustedProperty :                            -- not proven
        iPassedProperty :                               -- WIN
        cResultIndicies : _ = [0..]                     -- must be grouped like this for `exitCodeFor` to work

exitCodeFor (Count cs) = case length $ takeWhile (== 0) cs of
                                c | c < iExhaustedProperty -> ExitFailure 2
                                c | c < iPassedProperty    -> ExitFailure 1
                                _                          -> ExitSuccess

countResult (Failed _)    = countResultI iFailedProperty
countResult  Passed       = countResultI iPassedProperty
countResult (Exhausted _) = countResultI iExhaustedProperty
countResult (Unparsed _)  = countResultI iUnparsedProperty


runTest :: G.GhcMonad m => String -> m Count
runTest nTest = do      result <- fmap (maybe wrongTestResult divineTestResult . fromDynamic) (G.dynCompileExpr expr)
                        when (shouldExplicitlyReportResult result) (G.liftIO $ putStr $ reportResult nTest result)
                        return (countResult result)
  where -- `expr` originally copied from Test.QuickCheck, copyright © Koen Claessen <koen@chalmers.se>, licenced under BSD3
        expr = concat ["let gen = Test.QuickCheck.evaluate (Test.QuickCheck.property ", nTest, ") in ",
                       "let f rnd0 cPassed cMissed stamps",
                             "| cPassed == 100  = \"Passed\"",
                             "| cMissed == 1000 = \"Exhausted\"",
                             "| otherwise       = let (rnd1, rnd2) = System.Random.split rnd0 in ",
                                                 "let result = Test.QuickCheck.generate ((cPassed `div` 2) + 3) rnd2 gen in ",
                                                 "case Test.QuickCheck.ok result of",
                                                     "{Nothing    -> f rnd1 cPassed (1 + cMissed) stamps;",
                                                      "Just True  -> f rnd1 (1 + cPassed) cMissed ",
                                                                        "(Test.QuickCheck.stamp result : stamps);",
                                                      "Just False -> unlines (\"Failed\" : Test.QuickCheck.arguments result)} in ",
                       "f (System.IO.Unsafe.unsafePerformIO System.Random.newStdGen) 0 0 []"]
        divineTestResult :: String -> TestResult
        divineTestResult "Passed"                       = Passed
        divineTestResult "Exhausted"                    = Exhausted 1000
        divineTestResult s | "Failed" `isPrefixOf` s    = Failed (tail $ lines s)
                           | otherwise                  = Unparsed s
        wrongTestResult = Failed ["Was expecting property to be of type Testable a => a"]

runTests :: Flags -> String -> IO Count
runTests flags nf = G.runGhc (Just G.libdir) (G.defaultErrorHandler G.defaultDynFlags init) >>= putCountLn nf
  where init = do       -- have to get and then set dynamic flags even if I don't want to change them
                        -- somehow this initialises fields I don't want to care about
                        dynFlags <- G.getSessionDynFlags
                        G.setSessionDynFlags dynFlags {G.ghcMode = G.CompManager, G.ghcLink = G.LinkInMemory,
                                                       G.hscTarget = G.HscInterpreted,
                                                       G.flags = G.Opt_ImplicitImportQualified : G.flags dynFlags,
                                                       G.packageFlags = packageFlags ++ G.packageFlags dynFlags}
                        G.setTargets targets
                        G.load G.LoadAllTargets >>= loaded
        loaded G.Failed = return (countResultI iUncompiledFile)
        loaded G.Succeeded = G.getModuleGraph >>= loadedModule . map G.ms_mod
        loadedModule [onlyModule]
              = G.getModuleInfo onlyModule >>= performTests . fmap (catMaybes . map nTestFromTyThing . G.modInfoTyThings)
          where performTests Nothing = error "Was expecting module to be loaded"
                performTests (Just ns) = G.setContext [onlyModule] [] >> fmap mconcat (mapM runTest ns)
        loadedModule mods = error (concat ["loadedModule was expecting only one module but got ", show (length mods), ": ",
                                           intercalate ", " $ map (G.moduleNameString . G.moduleName) mods])
        nTestFromTyThing (G.AnId identity) = if "prop_" `isPrefixOf` n then Just n else Nothing
          where n = G.occNameString . G.nameOccName $ G.varName identity
        nTestFromTyThing _ = Nothing
        targets = [G.Target {G.targetId = G.TargetFile nf Nothing, G.targetAllowObjCode = False, G.targetContents = Nothing}]
        packageFlags = map f (pkgsFromFlags flags)
          where f (ExposePkg n) = G.ExposePackage n
                f (HidePkg n)   = G.HidePackage n
                f (IgnorePkg n) = G.IgnorePackage n


cmdLineDescr = [Option "" ["help"]           (NoArg $ withHelpOption True)                  "display this help",
                Option "" ["expose-package"] (ReqArg (withPkg . ExposePkg) "PACKAGE")       "expose a package",
                Option "" ["hide-package"]   (ReqArg (withPkg . HidePkg) "PACKAGE")         "hide a package",
                Option "" ["ignore-package"] (ReqArg (withPkg . IgnorePkg) "PACKAGE")       "ignore a package"]

usageMsg = usageInfo "Usage: hstest [FLAGS] [SOURCE FILES]" cmdLineDescr

getCommandLineOptions wrap = getArgs >>= act . getOpt (ReturnInOrder wrap) cmdLineDescr
  where act (opts, _, [])   = return opts
        act (_,    _, errs) = hPutStrLn stderr (unlines errs ++ usageMsg) >> exitFailure

data Options = Options {helpFlagFromOptions :: Bool,
                        nfsFromOptions :: [String],
                        flagsFromOptions :: Flags}
noOptions = Options False [] defaultFlags
withHelpOption flag opts = opts {helpFlagFromOptions = flag}
withNfOption nf opts = opts {nfsFromOptions = nf : nfsFromOptions opts}
withNfsOption nfs opts = opts {nfsFromOptions = nfs}
alterFlags f opts = opts {flagsFromOptions = f (flagsFromOptions opts)}

data Flags = Flags {pkgsFromFlags :: [PkgFlag]}
data PkgFlag = ExposePkg String
             | HidePkg String
             | IgnorePkg String
defaultFlags = Flags []
withPkg pkg = alterFlags $ \flags -> flags {pkgsFromFlags = pkg : pkgsFromFlags flags}

main = fmap (foldr ($) noOptions) (getCommandLineOptions withNfOption) >>= defaultNfs >>= act >>= exitWith . exitCodeFor
  where act Options {helpFlagFromOptions = True} = putStrLn usageMsg >> exitSuccess
        act opts @ Options {nfsFromOptions = [nf]} = runTests (flagsFromOptions opts) nf
        act opts = mapM (runTests $ flagsFromOptions opts) (nfsFromOptions opts) >>= putCountLn "Total" . mconcat
        defaultNfs opts @ Options {nfsFromOptions = []}
              = fmap ((`withNfsOption` opts) . filter (".hs" `isSuffixOf`)) (getDirectoryContents ".")
        defaultNfs opts = return opts