module Main (main) where

import qualified Language.Haskell.Exts.Parser as Parser
import qualified Language.Haskell.Exts.Syntax as Syntax
import qualified Language.Haskell.Exts.SrcLoc as SrcLoc
import Language.Haskell.Exts.Pretty (prettyPrint, )

import qualified Distribution.PackageDescription.Configuration as Config
import qualified Distribution.PackageDescription as P
import qualified Distribution.Package as Pkg
import qualified Distribution.Simple.LocalBuildInfo as LBI
import qualified Distribution.Simple.PackageIndex as PkgIdx
import qualified Distribution.Simple.Configure as Configure
import qualified Distribution.Simple.Setup as Setup
import qualified Distribution.InstalledPackageInfo as InstPkg
import qualified Distribution.ModuleName as DistModuleName
import Distribution.PackageDescription.Parse (readPackageDescription, )
import Distribution.Package (PackageName(PackageName), )
import Distribution.Simple.Utils
          (defaultPackageDesc, findModuleFiles, findFileWithExtension',
           notice, )

import qualified Distribution.Verbosity as Verbosity
import qualified Distribution.Version as Version
import qualified Distribution.ReadE as ReadE
import Distribution.Version (Version, )
import Distribution.Text (display, )


import qualified System.Environment as Env
import qualified System.IO as IO
import System.Console.GetOpt
          (ArgOrder(RequireOrder), OptDescr(Option), ArgDescr(NoArg, ReqArg),
           getOpt, usageInfo, )
import System.Exit (exitSuccess, exitFailure, )
import System.FilePath ((</>), )

import Text.Printf (printf, hPrintf, )

import qualified Control.Monad.Exception.Synchronous as Exc
import qualified Control.Monad.Trans.Class as MT

import qualified Data.NonEmpty as NonEmpty
import qualified Data.Foldable as Fold
import qualified Data.Monoid.HT as Mn
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Set as Set
import Control.Monad (when, )
import Control.Functor.HT (void, )
import Data.Maybe (mapMaybe, maybeToList, )
import Data.Set (Set, )
import Data.Foldable (foldMap, forM_, )


data Flags =
   Flags {
      flagHelp :: Bool,
      flagVerbosity :: Verbosity.Verbosity,
      flagBuildDir :: FilePath,
      flagClassifyDependencies :: Bool,
      flagPedantic :: Bool,
      flagCheckLibrary,
      flagCheckExecutables,
      flagCheckTestSuites,
      flagCheckBenchmarks :: Bool,
      flagExcludedModules :: Set Syntax.ModuleName,
      flagLoadPackageIndex :: Bool,
      flagCriticalModules :: PkgIdx.PackageIndex -> [DepAttrs] -> ModuleSet,
      flagCriticalModified :: Bool
   }

defaultFlags :: Flags
defaultFlags = Flags {
   flagHelp = False,
   flagVerbosity = Verbosity.silent,
   flagBuildDir = Setup.defaultDistPref,
   flagClassifyDependencies = False,
   flagPedantic = False,
   flagCheckLibrary = True,
   flagCheckExecutables = True,
   flagCheckTestSuites = True,
   flagCheckBenchmarks = True,
   flagExcludedModules = Set.empty,
   flagLoadPackageIndex = True,
   flagCriticalModules =
      \pkgIdx depAttrs ->
         ModuleSet $ dependentModules pkgIdx depAttrs,
   flagCriticalModified = False
   }

options :: [OptDescr (Flags -> Exc.Exceptional String Flags)]
options =
   Option ['h'] ["help"]
      (NoArg (\flags -> return $ flags{flagHelp = True}))
      "show options" :
   Option ['v'] ["verbose"]
      (ReqArg
         (\str flags ->
            fmap (\n -> flags{flagVerbosity = n}) $
            Exc.fromEither $
            ReadE.runReadE Verbosity.flagToVerbosity str)
         "N")
      "verbosity level: 0..3" :
   Option [] ["builddir"]
      (ReqArg (\str flags -> return $ flags{flagBuildDir = str}) "DIR")
      (printf "directory to look for package configuration (default %s)" $
       flagBuildDir defaultFlags) :
   Option [] ["classify-dependencies"]
      (NoArg (\flags -> return $ flags{flagClassifyDependencies = True}))
      "print diagnostics of version ranges in Build-Depends fields" :
   Option [] ["pedantic"]
      (NoArg (\flags -> return $ flags{flagPedantic = True}))
      "also check for hiding imports" :

   Option [] ["include-all"]
      (NoArg
         (\flags ->
            if flagCriticalModified flags
              then Exc.throw "include-all option must be the first amongst module set modifiers"
              else return $
                   (modifyFlagCritical
                       (const $ ComplementSet Set.empty) flags)
                   {flagLoadPackageIndex = False}))
      "check all imports, ignore package database" :
   Option [] ["include-import"]
      (ReqArg
         (\str flags -> return $
            modifyFlagCritical (insertModule (Syntax.ModuleName str)) flags)
         "MODULE")
      "check import of MODULE" :
   Option [] ["exclude-import"]
      (ReqArg
         (\str flags -> return $
            modifyFlagCritical (deleteModule (Syntax.ModuleName str)) flags)
         "MODULE")
      "ignore import of MODULE" :
   Option [] ["include-dependency"]
      (ReqArg
         (\str flags -> return $
            modifyFlagCriticalWithPkgIdx True
               (\pkgIdx -> insertModuleSet (moduleSetFromPackage pkgIdx str)) flags)
         "PKG")
      "check all imports from PKG" :
   Option [] ["exclude-dependency"]
      (ReqArg
         (\str flags -> return $
            modifyFlagCriticalWithPkgIdx True
               (\pkgIdx -> deleteModuleSet (moduleSetFromPackage pkgIdx str)) flags)
         "PKG")
      "ignore all imports from PKG" :

   Option [] ["exclude-module"]
      (ReqArg
          (\str flags ->
             return $ flags{flagExcludedModules =
                Set.insert (Syntax.ModuleName str) $
                flagExcludedModules flags})
          "MODULE")
      "do not check MODULE" :
   Option [] ["exclude-library"]
      (NoArg (\flags -> return $ flags{flagCheckLibrary = False}))
      "do not check library" :
   Option [] ["exclude-executables"]
      (NoArg (\flags -> return $ flags{flagCheckExecutables = False}))
      "do not check executables" :
   Option [] ["exclude-testsuites"]
      (NoArg (\flags -> return $ flags{flagCheckTestSuites = False}))
      "do not check testsuites" :
   Option [] ["exclude-benchmarks"]
      (NoArg (\flags -> return $ flags{flagCheckBenchmarks = False}))
      "do not check benchmarks" :
   []


modifyFlagCritical :: (ModuleSet -> ModuleSet) -> Flags -> Flags
modifyFlagCritical modify =
   modifyFlagCriticalWithPkgIdx False (const modify)

modifyFlagCriticalWithPkgIdx ::
   Bool ->
   (PkgIdx.PackageIndex -> ModuleSet -> ModuleSet) ->
   Flags -> Flags
modifyFlagCriticalWithPkgIdx loadPkgIdx modify flags =
   flags {
      flagCriticalModules =
         \pkgIdx depAttrs ->
            modify pkgIdx $ flagCriticalModules flags pkgIdx depAttrs,
      flagCriticalModified = True,
      flagLoadPackageIndex = loadPkgIdx || flagLoadPackageIndex flags
   }

moduleSetFromPackage :: PkgIdx.PackageIndex -> String -> Set Syntax.ModuleName
moduleSetFromPackage pkgIdx name =
   lookupPkgModuleSet pkgIdx (Pkg.PackageName name)


data ModuleSet =
     ModuleSet (Set Syntax.ModuleName)
   | ComplementSet (Set Syntax.ModuleName)

memberModule :: Syntax.ModuleName -> ModuleSet -> Bool
memberModule modu (ModuleSet set) = Set.member modu set
memberModule modu (ComplementSet set) = not $ Set.member modu set


insertModule :: Syntax.ModuleName -> ModuleSet -> ModuleSet
insertModule modu (ModuleSet set) = ModuleSet $ Set.insert modu set
insertModule modu (ComplementSet set) = ComplementSet $ Set.delete modu set

deleteModule :: Syntax.ModuleName -> ModuleSet -> ModuleSet
deleteModule modu (ModuleSet set) = ModuleSet $ Set.delete modu set
deleteModule modu (ComplementSet set) = ComplementSet $ Set.insert modu set


insertModuleSet :: Set Syntax.ModuleName -> ModuleSet -> ModuleSet
insertModuleSet new (ModuleSet set) = ModuleSet $ Set.union new set
insertModuleSet new (ComplementSet set) = ComplementSet $ Set.difference set new

deleteModuleSet :: Set Syntax.ModuleName -> ModuleSet -> ModuleSet
deleteModuleSet new (ModuleSet set) = ModuleSet $ Set.difference set new
deleteModuleSet new (ComplementSet set) = ComplementSet $ Set.union new set


main :: IO ()
main =
   Exc.resolveT exitFailureMsg $ do
      argv <- MT.lift Env.getArgs
      let (opts, args, errors) =
             getOpt RequireOrder options argv
      when (not (null errors)) $ Exc.throwT $ concat $ errors
      when (not (null args)) $ Exc.throwT $
         "I have no usage for the arguments " ++ show args
      flags <-
         Exc.ExceptionalT $ return $
            foldl (>>=) (return defaultFlags) opts
      when (flagHelp flags)
         (MT.lift $
          Env.getProgName >>= \programName ->
          putStrLn
             (usageInfo ("Usage: " ++ programName ++
                         " [OPTIONS]") options) >>
          exitSuccess)

      MT.lift $ run flags

data CheckFlags =
   CheckFlags {
      criticalModule :: Syntax.ModuleName -> Bool,
      checkPedantic :: Bool
   }

run :: Flags -> IO ()
run flags = do
   let verbosity = flagVerbosity flags
   pdfile <- defaultPackageDesc verbosity
   desc <-
      fmap Config.flattenPackageDescription $
      readPackageDescription verbosity pdfile

   notice verbosity "Package description"
   let classified = classifyDependencies $ P.buildDepends desc
   mapM_
      (printUpperBoundDiagnostics $ flagClassifyDependencies flags)
      classified

   pkgIdx <-
      if flagLoadPackageIndex flags
        then loadPackageIndex (flagBuildDir flags)
        else return $ error "no package index loaded"
   let modIdx = flagCriticalModules flags pkgIdx classified
   let checkFlags =
          CheckFlags {
             criticalModule = flip memberModule modIdx,
             checkPedantic = flagPedantic flags
          }

   when (flagCheckLibrary flags) $ do
      notice verbosity "Library"
      P.withLib desc $ \lib -> do
         let bi = P.libBuildInfo lib
             modules =
                excludeModules (flagExcludedModules flags) $
                   P.exposedModules lib
                   ++
                   P.otherModules bi
             sourceDirs = P.hsSourceDirs bi
         checkModules checkFlags =<< findModuleFiles sourceDirs ["hs"] modules

   when (flagCheckExecutables flags) $ do
      notice verbosity "Executables"
      P.withExe desc $ \exe -> do
         let name = P.exeName exe
         notice verbosity name
         let bi = P.buildInfo exe
             modules =
                excludeModules (flagExcludedModules flags) $
                P.otherModules bi
             sourceDirs = P.hsSourceDirs bi
         mainPath <- findMainModule sourceDirs $ P.modulePath exe
         paths <- findModuleFiles sourceDirs ["hs"] modules
         checkModules checkFlags $ maybeToList mainPath ++ paths

   when (flagCheckTestSuites flags) $ do
      notice verbosity "Test-Suites"
      P.withTest desc $ \exe -> do
         let name = P.testName exe
         notice verbosity name
         let bi = P.testBuildInfo exe
             modules =
                excludeModules (flagExcludedModules flags) $
                P.otherModules bi
             sourceDirs = P.hsSourceDirs bi
         paths <- findModuleFiles sourceDirs ["hs"] modules
         mainPath <-
            case P.testInterface exe of
               P.TestSuiteExeV10 _ path -> findMainModule sourceDirs path
               _ -> return Nothing
         checkModules checkFlags $ maybeToList mainPath ++ paths

   when (flagCheckBenchmarks flags) $ do
      notice verbosity "Benchmarks"
      P.withBenchmark desc $ \exe -> do
         let name = P.benchmarkName exe
         notice verbosity name
         let bi = P.benchmarkBuildInfo exe
             modules =
                excludeModules (flagExcludedModules flags) $
                P.otherModules bi
             sourceDirs = P.hsSourceDirs bi
         paths <- findModuleFiles sourceDirs ["hs"] modules
         mainPath <-
            case P.benchmarkInterface exe of
               P.BenchmarkExeV10 _ path -> findMainModule sourceDirs path
               _ -> return Nothing
         checkModules checkFlags $ maybeToList mainPath ++ paths

findMainModule :: [FilePath] -> FilePath -> IO (Maybe (FilePath, FilePath))
findMainModule sourceDirs path = do
   maybeMainPath <- findFileWithExtension' [""] sourceDirs path
   case maybeMainPath of
      Nothing -> do
         void $ hPrintf IO.stderr "main module %s not found" path
         return Nothing
      Just mainPath -> return $ Just mainPath

loadPackageIndex :: FilePath -> IO PkgIdx.PackageIndex
loadPackageIndex buildDir =
   fmap LBI.installedPkgs $
   Configure.getPersistBuildConfig buildDir

dependentModules ::
   PkgIdx.PackageIndex -> [DepAttrs] -> Set Syntax.ModuleName
dependentModules pkgIdx =
   foldMap (lookupPkgModuleSet pkgIdx) .
   map depPkgName .
   filter
      (\dep ->
         case depUpperBoundClass dep of
            Open -> True
            Lax _ -> True
            Generous _ _ -> True
            Tight _ -> False)

lookupPkgModuleSet ::
   PkgIdx.PackageIndex -> Pkg.PackageName -> Set Syntax.ModuleName
lookupPkgModuleSet pkgIdx name =
   Set.fromList $
   map syntaxFromDistModuleName $
   concatMap InstPkg.exposedModules $
   concatMap snd $
   PkgIdx.lookupPackageName pkgIdx name

excludeModules ::
   Set Syntax.ModuleName ->
   [DistModuleName.ModuleName] -> [DistModuleName.ModuleName]
excludeModules set =
   filter (not . flip Set.member set . syntaxFromDistModuleName)

syntaxFromDistModuleName ::
   DistModuleName.ModuleName -> Syntax.ModuleName
syntaxFromDistModuleName =
   Syntax.ModuleName . List.intercalate "." . DistModuleName.components

checkModules :: CheckFlags -> [(FilePath, FilePath)] -> IO ()
checkModules flags paths =
   forM_ paths $ \(dir,path) -> do
      let dirPath = dir </> path
      txt <- readFile dirPath
      case
         Parser.parseWithMode
            (Parser.defaultParseMode {Parser.parseFilename = dirPath}) txt of

         Parser.ParseFailed loc msg ->
            hPrintf IO.stderr "\n%s\n    %s\n" (formatSrcLoc loc) msg

         Parser.ParseOk modu -> checkModule flags modu

checkModule :: CheckFlags -> Syntax.Module -> IO ()
checkModule flags
      (Syntax.Module _loc _name _pragma _warn _export imports _decls) = do
   forM_ (filter (criticalModule flags . Syntax.importModule) imports) $ \imp -> do
      let problems =
             Mn.when (not $ strictImport imp) ["lax import"]
             ++
             (Mn.when
                 (checkPedantic flags && Fold.any fst (Syntax.importSpecs imp))
                 ["hiding import"])
             ++
             (flip map (implicitSpecs imp) $ \name ->
                 "open constructor or method list for " ++ prettyPrint name)

      when (not $ null problems) $ do
         void $
            printf "\n%s:\n    Problems encountered in import of %s:\n"
               (formatSrcLoc $ Syntax.importLoc imp)
               (unpackModuleName $ Syntax.importModule imp)
         putStr $ unlines $ map (replicate 8 ' ' ++) problems

   let conflictAbbrevs =
          Map.toAscList $
          Map.filter
             (\mods ->
                Set.size mods >= 2
                &&
                (not $ Set.null $
                 Set.filter (criticalModule flags . snd) mods)) $
          Map.fromListWith Set.union $
          mapMaybe
             (\imp ->
                fmap
                   (\impAs ->
                      (impAs,
                       Set.singleton
                          (Syntax.importLoc imp,
                           Syntax.importModule imp)))
                   (Syntax.importAs imp)) $
          imports

   forM_ conflictAbbrevs $ \(impAs, conflicts) -> do
      void $
         printf "\nMultiple modules imported with abbreviation \"%s\":\n"
            (unpackModuleName impAs)
      forM_ (Set.toAscList conflicts) $ \(loc, modu) ->
         printf "\n%s:\n    conflicting import of %s\n"
            (formatSrcLoc loc)
            (unpackModuleName modu)

formatSrcLoc :: SrcLoc.SrcLoc -> String
formatSrcLoc loc =
   printf "%s:%d:%d"
      (SrcLoc.srcFilename loc)
      (SrcLoc.srcLine loc)
      (SrcLoc.srcColumn loc)

unpackModuleName :: Syntax.ModuleName -> String
unpackModuleName (Syntax.ModuleName str) = str

strictImport :: Syntax.ImportDecl -> Bool
strictImport imp =
   Syntax.importQualified imp
   ||
   Fold.any
      (\(hide, _specs) -> not hide)
      (Syntax.importSpecs imp)

implicitSpecs :: Syntax.ImportDecl -> [Syntax.Name]
implicitSpecs imp =
   foldMap
      (\(hide, specs) -> if hide then [] else mapMaybe maybeImplicitSpec specs)
      (Syntax.importSpecs imp)

maybeImplicitSpec :: Syntax.ImportSpec -> Maybe Syntax.Name
maybeImplicitSpec spec =
   case spec of
      Syntax.IThingAll name -> Just name
      Syntax.IThingWith _ _ -> Nothing
      Syntax.IAbs _ -> Nothing
      Syntax.IVar _ -> Nothing


data DepAttrs =
   DepAttrs {
      depPkgName :: PackageName,
      depMissingUpperBounds :: [Version.LowerBound],
      depInclusiveUpperBounds :: [Version],
      depUpperBoundClass :: BoundClass
   }

data BoundClass = Open | Lax Int | Generous Int Int | Tight [Int]


printUpperBoundDiagnostics :: Bool -> DepAttrs -> IO ()
printUpperBoundDiagnostics classifyAll depAttrs =
   let warn = (,) True
       info = (,) False

       msgs =
          (flip map (depMissingUpperBounds depAttrs) $ \(Version.LowerBound ver typ) ->
              warn $
              printf "missing upper bound associated with lower bound \"%s\"" $
              display $
              case typ of
                 Version.InclusiveBound -> Version.orLaterVersion ver
                 Version.ExclusiveBound -> Version.laterVersion ver)
          ++
          (flip map (depInclusiveUpperBounds depAttrs) $ \uppBnd ->
              warn $ printf "found inclusive upper bound %s" $ display uppBnd)
          ++
          case depUpperBoundClass depAttrs of
             Open -> []
             Lax x -> [warn $ printf "upper bound %d is too lax" x]
             Generous x y ->
                [info $ printf "upper bound %d.%d requires strict imports" x y]
             Tight xs ->
                [info $
                 printf "upper bound %s is tight" $
                 List.intercalate "." $ map show xs]

       filteredMsgs =
          map snd $
          if classifyAll
            then msgs
            else filter fst msgs

   in  when (not $ null filteredMsgs) $ do
          putStrLn $ unpackPkgName $ depPkgName depAttrs
          putStrLn $ unlines $ map (replicate 4 ' ' ++) filteredMsgs

classifyDependencies :: [Pkg.Dependency] -> [DepAttrs]
classifyDependencies deps =
   flip map deps $ \(Pkg.Dependency dependName rng) ->
      let intervals = Version.asVersionIntervals rng

          maybeUpperBound Version.NoUpperBound = Nothing
          maybeUpperBound (Version.UpperBound ver bnd) = Just (ver, bnd)
          isExclusiveBound Version.ExclusiveBound = True
          isExclusiveBound Version.InclusiveBound = False
          (upperBounds, noUpperBounds) =
             ListHT.partitionMaybe (maybeUpperBound . snd) intervals
          (exclusiveUpperBounds, inclusiveUpperBounds) =
             ListHT.partition (isExclusiveBound . snd) upperBounds

          branches =
             case NonEmpty.fetch exclusiveUpperBounds of
                Nothing -> []
                Just xs ->
                   NonEmpty.minimumKey length $
                   fmap (ListHT.dropWhileRev (0==) .
                         Version.versionBranch . fst) xs

          boundClass =
             case branches of
                [] -> Open
                [x] -> Lax x
                [x,y] -> Generous x y
                _ -> Tight branches

      in  DepAttrs {
             depPkgName = dependName,
             depMissingUpperBounds = map fst noUpperBounds,
             depInclusiveUpperBounds = map fst inclusiveUpperBounds,
             depUpperBoundClass = boundClass
          }

exitFailureMsg :: String -> IO ()
exitFailureMsg msg = do
   IO.hPutStrLn IO.stderr $ "Aborted: " ++ msg
   exitFailure

unpackPkgName :: PackageName -> String
unpackPkgName (PackageName name) = name