module Main where

import qualified Distribution.PackageDescription as P
import qualified Distribution.Verbosity as Verbosity
import qualified Distribution.ReadE as ReadE
import Distribution.Simple.PackageDescription (readGenericPackageDescription)
import Distribution.PackageDescription
         (GenericPackageDescription, package, packageDescription)
import Distribution.Types.Dependency (Dependency(Dependency))
import Distribution.Types.PackageName (PackageName, unPackageName)
import Distribution.Types.PackageId (pkgName)

import qualified Options.Applicative as OP
import Shell.Utility.Exit (exitFailureMsg)

import qualified System.FilePath as FilePath
import System.FilePath ((</>))

import qualified Data.Graph.Inductive.Query.DFS as GraphQuery
import qualified Data.Graph.Inductive.Graph as Graph
import Data.Graph.Inductive.Tree (Gr)

import qualified Control.Monad.Exception.Synchronous as Exc
import qualified Control.Monad.Trans.Class as Trans
import Control.Arrow ((***))
import Control.Monad (guard, when)
import Control.Applicative (pure, (<*>), (<|>))

import qualified Data.Foldable as Fold
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.Map (Map)
import Data.Maybe (fromMaybe, maybeToList)


main :: IO ()
main = do
   (flags, cabalPaths) <- OP.execParser $ info options
   Exc.resolveT (\e -> exitFailureMsg $ "Aborted: " ++ e) $
      sortCabalFiles flags cabalPaths


data Flags =
   Flags {
      optVerbosity :: Verbosity.Verbosity,
      optInfo :: SourcePackage -> String,
      optOutputFormat :: Format,
      optBuilddir :: FilePath,
      optInstall  :: String
   }

info :: OP.Parser a -> OP.ParserInfo (a, [String])
info p =
   OP.info
      (OP.helper <*>
         OP.liftA2 (,) p (OP.many (OP.strArgument (OP.metavar "CABAL-FILE"))))
      (OP.fullDesc <>
       OP.progDesc
         "Topological sort of Cabal packages according to dependencies.")

infoMap :: Map String (SourcePackage -> String)
infoMap =
   Map.fromList $
      ("name", unPackageName . pkgNameFromDescription . description) :
      ("path", location) :
      ("dir", FilePath.takeDirectory . location) :
      []

data Format = Serial | Parallel | Makefile
   deriving (Eq, Ord, Show, Enum)

options :: OP.Parser Flags
options =
   pure Flags
   <*>
      OP.option
         (OP.eitherReader $ ReadE.runReadE Verbosity.flagToVerbosity)
         (OP.short 'v' <>
          OP.long "verbose" <>
          OP.metavar "N" <>
          OP.value Verbosity.silent <>
          OP.help "verbosity level: 0..3")
   <*>
      OP.option
         (OP.eitherReader $ \str ->
            maybe (Left $ "unknown info type " ++ str) Right $
            Map.lookup str infoMap)
         (OP.long "info" <>
          OP.metavar "KIND" <>
          OP.value location <>
          OP.help ("kind of output: " ++
                     List.intercalate ", " (Map.keys infoMap)))
   <*>
      (OP.flag' Makefile
         (OP.short 'm' <>
          OP.long "makefile" <>
          OP.help "Generate a makefile of package dependencies")
       <|>
       OP.flag Serial Parallel
         (OP.short 'p' <>
          OP.long "parallel" <>
          OP.help "Display independently buildable groups of packages"))
   <*>
      OP.strOption
         (OP.long "builddir" <>
          OP.metavar "PATH" <>
          OP.value "." <>
          OP.help "Specify the build dir to use for generated makefile")
   <*>
      OP.strOption
         (OP.long "install-cmd" <>
          OP.metavar "CMD" <>
          OP.value "cabal v1-install" <>
          OP.help "Specify the install command to use in generated makefile")


data SourcePackage =
   SourcePackage {
      location :: FilePath,
      description :: GenericPackageDescription
   }
   deriving (Show, Eq)

type DependencyGraph = Gr SourcePackage ()

sortCabalFiles :: Flags -> [FilePath] -> Exc.ExceptionalT String IO ()
sortCabalFiles flags cabalPaths =
   do pkgDescs <-
         Trans.lift $
         mapM (readGenericPackageDescription (optVerbosity flags)) cabalPaths
      when (optVerbosity flags >= Verbosity.verbose) $
         Trans.lift $
         Fold.for_ pkgDescs $ \pkgDesc -> do
            putStrLn $ unPackageName (pkgNameFromDescription pkgDesc) ++ ":"
            let deps =
                   Set.toAscList $ Set.fromList $
                   map (unPackageName . depName) $
                   allDependencies pkgDesc
            Fold.for_ deps $ \dep -> putStrLn $ "  " ++ dep
      let graph = dependencyGraph $ zipWith SourcePackage cabalPaths pkgDescs
      checkForCycles graph
      Trans.lift $
         case optOutputFormat flags of
            Makefile -> printMakefile flags $ getDeps graph
            Serial ->
               mapM_ (putStrLn . optInfo flags) $ GraphQuery.topsort' graph
            Parallel ->
               mapM_ (putStrLn . unwords . map (optInfo flags)) $
               map (GraphQuery.topsort' . flip Graph.subgraph graph) $
               GraphQuery.components graph


printMakefile :: Flags -> [(SourcePackage, [SourcePackage])] -> IO ()
printMakefile flags deps = do
    let printDep (l, ls) = putStrLn (l ++ ": " ++ unwords ls)
        stamp =
           (optBuilddir flags </>) .
           flip FilePath.replaceExtension "cstamp" . location
        allDeps = unwords (map (stamp . fst) deps)
    putStrLn (optBuilddir flags </> "%.cstamp:")
    putStrLn ("\t" ++ optInstall flags ++ " `dirname $*`")
    putStrLn "\tmkdir -p `dirname $@`"
    putStrLn "\ttouch $@"
    putStrLn ""
    putStrLn ("all: " ++ allDeps)
    putStrLn ""
    putStrLn "clean:"
    putStrLn ("\t$(RM) " ++ allDeps)
    putStrLn ""
    mapM_ (printDep . (stamp *** map stamp)) deps

getDeps :: DependencyGraph -> [(SourcePackage, [SourcePackage])]
getDeps gr =
    let c2dep :: Graph.Context SourcePackage () ->
                  (SourcePackage, [SourcePackage])
        c2dep ctx =
           (Graph.lab' ctx,
            map (Graph.lab' . Graph.context gr)
                  (Graph.pre gr . Graph.node' $ ctx))
    in  Graph.ufold (\ctx ds -> c2dep ctx : ds) [] gr

dependencyGraph :: [SourcePackage] -> DependencyGraph
dependencyGraph srcPkgs =
   let nodes = zip [0..] srcPkgs
       nodeDict =
          Map.fromList $
          zip (map (pkgNameFromDescription . description) srcPkgs) [0..]
       edges = do
          (srcNode,desc) <- nodes
          dependency <- allDependencies $ description desc
          dstNode <- maybeToList $ Map.lookup (depName dependency) nodeDict
          guard (dstNode /= srcNode)
          return (dstNode, srcNode, ())
   in  Graph.mkGraph nodes edges


checkForCycles :: (Monad m) => DependencyGraph -> Exc.ExceptionalT String m ()
checkForCycles graph =
   case getCycles graph of
      [] -> return ()
      cycles ->
         Exc.throwT $ unlines $
         "Cycles in dependencies:" :
         map (unwords . map location . nodeLabels graph) cycles

nodeLabels :: Gr a b -> [Graph.Node] -> [a]
nodeLabels graph =
   map (fromMaybe (error "node not found in graph") . Graph.lab graph)

getCycles :: Gr a b -> [[Graph.Node]]
getCycles =
   filter (\component -> case component of _:_:_ -> True; _ -> False) .
   GraphQuery.scc


allDependencies :: GenericPackageDescription -> [Dependency]
allDependencies pkg =
   P.allBuildDepends (packageDescription pkg) ++
   maybe [] (concatMap snd . flattenCondTree) (P.condLibrary pkg) ++
   concatMap (concatMap snd . flattenCondTree . snd) (P.condExecutables pkg)

flattenCondTree :: P.CondTree v c a -> [(a,c)]
flattenCondTree tree =
   (P.condTreeData tree, P.condTreeConstraints tree) :
   concatMap
      (\(P.CondBranch _ thenBranch elseBranch) ->
         flattenCondTree thenBranch ++
         maybe [] flattenCondTree elseBranch)
      (P.condTreeComponents tree)

depName :: Dependency -> PackageName
depName (Dependency name _ _) = name

pkgNameFromDescription :: GenericPackageDescription -> PackageName
pkgNameFromDescription = pkgName . package . packageDescription