content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# OPTIONS_GHC -Wall #-}
  
  import Control.Applicative
  import Control.Monad
  import Data.Bool
  import Data.Foldable
  import Data.Function
  import Data.List
  import qualified Data.Map as Map
  import Data.Maybe
  import Data.Monoid
  import qualified Data.Set as Set
  import Data.String
  import Distribution.Simple hiding (Language (..), Module (..))
  import Distribution.Simple.BuildPaths
  import Distribution.Simple.LocalBuildInfo
  import Distribution.Types.BuildInfo
  import Distribution.Types.Library
  import Distribution.Types.PackageDescription
  import Distribution.Utils.Path
  import Language.Haskell.Exts (Comment (..), Decl (TypeSig), Language (..), Module (..), Name (..), ParseMode (..), SrcSpanInfo, associateHaddock, fromParseResult, parseFileWithComments)
  import qualified Language.Haskell.Exts as Exts
  import System.Directory
  import System.FilePath
  import Prelude
  
  collectTests :: FilePath -> [FilePath] -> IO [(String, String, String, String)]
  collectTests pkgRoot roots =
    concat <$> traverse (findAllTests . (<> "/Test")) roots
    where
      findAllTests :: FilePath -> IO [(String, String, String, String)]
      findAllTests root = do
        paths <- findPaths root
        concat <$> traverse (findModuleTests root) (filter (not . noise) paths)
  
      -- don't touch emacs auto-save or backup files
      noise :: FilePath -> Bool
      noise (last . splitDirectories -> fn) = autosave fn || backup fn
        where
          backup = (== '~') . last
          autosave ('.' : '#' : _) = True
          autosave ('#' : _) = True -- (if you want to make this pattern more strict: there is also a '#' at the end.)
          autosave _ = False
  
      findModuleTests :: FilePath -> FilePath -> IO [(String, String, String, String)]
      findModuleTests root path = do
        let modl = "Test." <> toModule root path
        tests <- collectTestsInModule pkgRoot path
        pure $ map (\(testName, summary, full) -> (modl, testName, summary, full)) tests
  
      toModule :: FilePath -> FilePath -> String
      toModule root = map setDot . dropExtension . makeRelative root
        where
          setDot '/' = '.'
          setDot c = c
  
      findPaths :: FilePath -> IO [FilePath]
      findPaths d = do
        isDir <- doesDirectoryExist d
        if isDir
          then do
            entries <- listDirectory d
            concat <$> traverse (findPaths . (d </>)) entries
          else pure [d]
  
  stripHaddock :: String -> String
  stripHaddock = \case
    ' ' : '|' : ' ' : xs -> xs
    ' ' : '|' : xs -> xs
    ' ' : '^' : ' ' : xs -> xs
    ' ' : '^' : xs -> xs
    ' ' : xs -> xs
    xs -> xs
  
  collectDescription :: [String] -> (String, String)
  collectDescription ls =
    case uncons $ stripHaddock <$> ls of
      Nothing -> ("", "")
      Just (summary, rest) -> (summary, unlines (dropWhile null rest))
  
  collectTestsInModule :: FilePath -> FilePath -> IO [(String, String, String)]
  collectTestsInModule pkgRoot fn = do
    -- associateHaddock requires all comments that we want to stick onto a test
    -- should be in the Haddock style, otherwise they won't make it through the parser.
    res <-
      associateHaddock . fromParseResult
        <$> parseFileWithComments
          -- Haskell2010 is the closest we have to getting haskell-src-exts to
          -- playing nicely with GHC2021. One annoying feature it can't handle is
          -- ImportQualifiedPost, so all of our tests have to use traditional
          -- qualified import syntax.
          (ParseMode fn Haskell2010 extensions False False Nothing False)
          absolutePath
    case res of
      Module _ _ _ _ decs ->
        pure $
          decs >>= \case
            TypeSig _ names _ -> mapMaybe testName names
            _ -> []
      _ -> error "XmlPage and XmlHybrid handling not set up. Please fix me!"
    where
      extractComment :: Comment -> String
      extractComment (Comment _ _ s) = s
  
      testName :: Name (SrcSpanInfo, [Comment]) -> Maybe (String, String, String)
      testName name =
        let (n', comments) =
              case name of
                Ident (_, cs) n -> (n, extractComment <$> cs)
                Symbol (_, cs) n -> (n, extractComment <$> cs)
         in if "test" `isPrefixOf` n'
              then
                let (summary, rest) = collectDescription comments
                 in pure (n', summary, rest)
              else Nothing
  
      absolutePath = pkgRoot </> fn
  
      -- All of the haskell-src-exts supported extensions that we are using.
      -- Several that are in the cabal file couldn't be directly copied over,
      -- but they aren't causing trouble at the moment.
      -- ImportQualifiedPost is an important one we use elsewhere in this repo
      -- that we can't use in `integration` as haskell-src-exts doesn't support
      -- it currently.
      extensions =
        [ Exts.EnableExtension Exts.BangPatterns,
          Exts.EnableExtension Exts.BlockArguments,
          Exts.EnableExtension Exts.ConstraintKinds,
          Exts.EnableExtension Exts.DataKinds,
          Exts.EnableExtension Exts.DefaultSignatures,
          Exts.EnableExtension Exts.DeriveFunctor,
          Exts.EnableExtension Exts.DeriveGeneric,
          Exts.EnableExtension Exts.DeriveTraversable,
          Exts.EnableExtension Exts.DerivingStrategies,
          Exts.EnableExtension Exts.DerivingVia,
          Exts.EnableExtension Exts.EmptyCase,
          Exts.EnableExtension Exts.FlexibleContexts,
          Exts.EnableExtension Exts.FlexibleInstances,
          Exts.EnableExtension Exts.FunctionalDependencies,
          Exts.EnableExtension Exts.GADTs,
          Exts.EnableExtension Exts.GeneralizedNewtypeDeriving,
          Exts.EnableExtension Exts.InstanceSigs,
          Exts.EnableExtension Exts.KindSignatures,
          Exts.EnableExtension Exts.LambdaCase,
          Exts.EnableExtension Exts.MultiParamTypeClasses,
          Exts.EnableExtension Exts.MultiWayIf,
          Exts.EnableExtension Exts.NamedFieldPuns,
          Exts.EnableExtension Exts.OverloadedLabels,
          Exts.EnableExtension Exts.PackageImports,
          Exts.EnableExtension Exts.PatternSynonyms,
          Exts.EnableExtension Exts.PolyKinds,
          Exts.EnableExtension Exts.QuasiQuotes,
          Exts.EnableExtension Exts.RankNTypes,
          Exts.EnableExtension Exts.RecordWildCards,
          Exts.EnableExtension Exts.ScopedTypeVariables,
          Exts.EnableExtension Exts.StandaloneDeriving,
          Exts.EnableExtension Exts.TupleSections,
          Exts.EnableExtension Exts.TypeApplications,
          Exts.EnableExtension Exts.TypeFamilies,
          Exts.EnableExtension Exts.TypeFamilyDependencies,
          Exts.EnableExtension Exts.TypeOperators,
          Exts.EnableExtension Exts.UndecidableInstances,
          Exts.EnableExtension Exts.ViewPatterns
        ]
  
  testHooks :: UserHooks -> UserHooks
  testHooks hooks =
    hooks
      { buildHook = \p l h f -> generate p l >> buildHook hooks p l h f,
        haddockHook = \p l h f -> generate p l >> haddockHook hooks p l h f,
        replHook = \p l h f args -> generate p l >> replHook hooks p l h f args
      }
    where
      generate p l = do
        let cname = CLibName LMainLibName
        let roots = case lookupComponent p cname of
              Just (CLib lib) -> map getSymbolicPath (hsSourceDirs (libBuildInfo lib))
              _ -> []
        for_ (Map.lookup cname (componentNameMap l)) $ \compBIs -> do
          for_ compBIs $ \compBI -> do
            let dest = autogenComponentModulesDir l compBI </> "RunAllTests.hs"
            tests <- collectTests (dataDir p) roots
            let modules = Set.toList (Set.fromList (map (\(m, _, _, _) -> m) tests))
            createDirectoryIfMissing True (takeDirectory dest)
            writeFile
              dest
              ( unlines
                  [ "module RunAllTests where",
                    "import Testlib.PTest",
                    "import Prelude",
                    "import Control.Monad.Trans.Writer",
                    unlines (map ("import qualified " <>) modules),
                    "mkAllTests :: IO [Test]",
                    "mkAllTests = execWriterT $ do",
                    unlines (map (\(m, n, s, f) -> "  yieldTests " <> unwords [show m, show n, show s, show f, m <> "." <> n]) tests)
                  ]
              )
            pure ()
  
  main :: IO ()
  main = do
    defaultMainWithHooks (testHooks simpleUserHooks)
  