{-# OPTIONS -Wall -Werror #-}
-- Requires GHC 6.10.4

{-
    TODO:

      - get coreOutFile from a command-line flag

-}

module Main where
import Control.Monad
import Data.List
import Data.Maybe
import System.Console.GetOpt
import System.Directory
import System.Environment
import System.FilePath

import Packages
import GHC hiding (showModule)
import HscTypes
import Outputable
import DynFlags hiding (Option(..))
import Module
import StaticFlagParser

import qualified Language.Core.Core as F
import Language.Core.Check
import Language.Core.ElimDeadCode 
import Language.Core.ParsecParser
import Language.Core.Merge
import Language.Core.Prims

import Text.ParserCombinators.Parsec

import Config

main :: IO ()
main = do
  args <- getArgs
  
-- Add -fext-core as a static argument
  parseStaticFlags staticArgs

  case getOpt Permute options args of
    (opts, [fn], _) ->   runGhc (Just ghcTopDir) $ do
      let pkgRoot = getPkgRootDir opts
      let coreOutFile = getCoreOutFile opts

-- This code generates Core for all modules in the "main" package.
      dflags <- getSessionDynFlags
      (newflags,_,_) <- parseDynamicFlags dflags (dynamicArgs fn)
      setSessionDynFlags newflags
      guessTarget fn Nothing >>= addTarget
      load LoadAllTargets

-- This code figures out what packages we depend on, and makes a list of
-- all the modules in those packages

      mg <- getModuleGraph
      let thisPackage_core_files = map ((flip replaceExtension) ".hcr")
             (catMaybes (map (ml_hs_file . ms_location) mg))
      debug $ "coreFiles = " ++ show thisPackage_core_files
      dflags1 <- getSessionDynFlags 
      (_,pkgDeps) <- liftIO $ initPackages dflags1
      debug ("packages to link in = " ++ show (map packageIdString pkgDeps))

-- pkgDeps now contains a list of all the packages our program depends on.
-- Next we need to:
-- -- * snarf all modules (both exposed and hidden) out of all the packages in pkgs
-- -- * link everything into one module
-- -- * kill dead code

      extDeps <- packageDbModules True


     -- This ASSUMES you've *already* generated Core for all but the main package.

-- extDeps is now a list of all modules available in all packages that are in scope.

-- If something is in the "main" package then its Core file name is included in
-- thisPackage_core_files.
-- If something is in any other package, we use the appropriate packageRoot argument
-- to find the sources for it.

-- This next line:
-- * takes extDeps (the list of modules in packages in scope)
-- * filters it to only include those that are in packages that the current module depends on
      let depsInOtherPackages = filter ((`elem` pkgDeps) . modulePackageId)
-- The primitive module doesn't have an implementation, so we exclude it
                                  (filter (/= primModule) extDeps)
      debug ("dependencies in non-main packages: " ++ showModules depsInOtherPackages)

      otherPackageCoreFiles <- liftIO $ mapM (moduleNameToCoreFile pkgRoot) depsInOtherPackages
      let allCoreFiles = otherPackageCoreFiles ++ thisPackage_core_files

      debug ("allCoreFiles = " ++ show allCoreFiles)
      debug ("thisPackage_coreFiles = " ++ show thisPackage_core_files)

      -- And then this code does the actual linking:
      modules <- liftIO $ mapM ((liftM requireOK) . parseCore) allCoreFiles
      let single = merge [] modules
      let deadKilled@(F.Module mn _ _) = elimDeadCode False single
      liftIO $ writeFile coreOutFile (show deadKilled)
      liftIO $ when typecheckCore (case checkModule initialEnv deadKilled of
           OkC _ ->
	     do putStrLn $ "Check succeeded! for " ++ show mn
                return ()
	   FailC s -> error ("Typechecking failed: " ++ s))
    _                  -> usageError

staticArgs :: [Located String]
staticArgs = map noLoc ["-fext-core"]

dynamicArgs :: FilePath -> [Located String]
dynamicArgs path = map noLoc ["-i"++takeDirectory path, "-fforce-recomp", "-O",
-- We just compile to ext-core; we don't try to generate an executable
  "-c",
 -- This -package-conf flag says "Look at the following packages for the standard library
 -- instead of the ones built into your own GHC install." This is what allows finding the
 -- GHC libraries that were compiled with the "special" flags.
 -- The actual location of the package.conf file is set at configure time.
   "-package-conf " ++ pkgConf
  ]

sp :: Outputable a => a -> String
sp = showSDoc . ppr

options :: [OptDescr Flag]
options = [Option [] ["package-root"] (ReqArg PkgRootDir "DIRECTORY")
              ("Directory d in which to find source code for installed packages;\n"
           ++ "if package foo is installed, I will expect d/foo to contain sources for\n"
           ++ "foo.\n"
           ++ "If you supply more than one package-root, I will search them in order."),
           Option ['o'] ["out"] (ReqArg OutFile "FILEPATH")
              ("Output file for the linked External Core module")]

getCoreOutFile :: [Flag] -> FilePath
getCoreOutFile fs =
  case find isCoreOutFile fs of
     Just (OutFile outf) -> outf
     _   -> error "You must supply a Core output file with -o."

getPkgRootDir :: [Flag] -> [FilePath]
getPkgRootDir fs =
  case pkgRootDirs fs of
      [] -> error "You must provide at least one package root directory with --package-root=/path/to/Core/library/sources." 
      ps -> ps
    where pkgRootDirs :: [Flag] -> [FilePath]
          pkgRootDirs []                     = []
          pkgRootDirs ((PkgRootDir fp):rest) = fp:pkgRootDirs rest
          pkgRootDirs (_:rest)               = pkgRootDirs rest

isCoreOutFile :: Flag -> Bool
isCoreOutFile (OutFile _) = True
isCoreOutFile _           = False

-- Regrettably, I was forced to include multiple pkgRoots because
-- when GHC turns an .hsc file into an .hcr file (in the library tree),
-- the .hcr file ends up living in libraries/base/dist/build/...
moduleNameToCoreFile :: [FilePath] -> Module -> IO FilePath
moduleNameToCoreFile pkgRoots m = do
  -- try both with and without the package id
  let leafLeaf = moduleNameSlashes (moduleName m) <.> ".hcr"
  let possibleLeaves = [sp (modulePackageId m) </> leafLeaf, leafLeaf]

  let possibleFiles = concatMap (\ root -> (map (\ leaf -> root </> leaf) possibleLeaves)) pkgRoots

  existing <- mapM (\ f -> do
                exists <- doesFileExist f
                return (f, exists)) possibleFiles
  maybe
    (error $ "I couldn't find a file for " ++ showModule m ++ "; I tried " ++ show possibleFiles)
    (return . fst)
    (find snd existing)

data Flag = PkgRootDir FilePath | OutFile FilePath 
  deriving Eq

usageError :: a
usageError = error (usageInfo "link" options)

requireOK :: Either ParseError F.Module -> F.Module
requireOK (Left err) = error (show err)
requireOK (Right m)  = m

showModules :: [Module] -> String
showModules = concat . (intersperse "\n") . (map showModule)

showModule :: Module -> String
showModule = showSDoc . pprModule

-- Would be nice not to have to define this
primModule :: Module
primModule = mkModule primPackageId (mkModuleName "GHC.Prim")

debug :: String -> Ghc ()
debug s | dEBUG = liftIO (putStrLn s)
debug _ = return ()

dEBUG :: Bool
dEBUG = False

typecheckCore :: Bool
typecheckCore = False