  Copyright 2013  
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or -3.0.txt

{-# LANGUAGE OverloadedStrings #-}
module Shake.HsDeps (
    Generated(..)
    , importsOf, transitiveImportsOf
    , importsOfIO
    , importsPackagagesOf_
    , loadPackageDb, savePackageDb
) where
import qualified Control.Concurrent as Concurrent
import qualified Control.Exception as Exception
import qualified Control.Monad.Trans as Trans

import qualified Data.ByteString.Char8 as B
import qualified Data.Map as Map
import qualified Data.Maybe as Maybe
import qualified Data.Serialize as Serialize
import qualified Data.Set as Set

import qualified Development.Shake as Shake
import qualified System.Directory as Directory
import qualified System.Exit as Exit
import qualified System.FilePath as FilePath
import qualified System.IO as IO
import qualified System.Process as Process

import qualified Util.Maps as Maps
import qualified Util.Seq as Seq
import qualified Shake.Util as Util

import           Control.Monad


type Package = String
type ModuleName = B.ByteString

-- | Normally 'importsOf' and 'transitiveImportsOf' filter out files that don't
-- exist, under the assumption that they belong to external packages.  This
-- set contains exceptions to that, so they will get a need call, so they can
-- be generated.
data Generated = Generated {
    _generatedHs :: Set.Set FilePath
    -- | Try files with these extensions, in addition to .hs.
    -- E.g. [".hsc", ".chs"].
    , _generatedExtensions :: [FilePath]
    } deriving (Eq, Show)

-- | Find files of modules this module imports, in the form A/B.hs or A/B.hsc.
-- Paths that don't exist are assumed to be package imports and are omitted.
importsOf :: Generated -> Maybe [String] -> FilePath -> Shake.Action [FilePath]
importsOf generated cppFlags fn = do
    Shake.need [fn]
    Trans.liftIO (importsOfIO generated cppFlags fn)

importsOfIO :: Generated
  ^ If Just , first run CPP with these flags .
    -> FilePath -> IO [FilePath]
importsOfIO generated cppFlags fn = do
  TODO get CcDeps.includesOf so I can need them too .
    mods <- Seq.unique_sort . parseImports <$> preprocess cppFlags fn
    Maybe.catMaybes <$> mapM (fileOf generated) mods

  * 

-- | Map a module name to the package it comes from.
type PackageDb = Map.Map ModuleName Package

-- | Get local imports and package dependencies.
--
-- I thought I wound use this but wound up not wanting it.  But I'll leave the
-- basic implementation here in case I change my mind.  This still needs to
-- be integrated with 'importsOf' and 'transitiveImportsOf'.
importsPackagagesOf_ :: PackageDb
  ^ If Just , first run CPP with these flags .
    -> FilePath -> IO ([FilePath], [Package])
importsPackagagesOf_ packageDb generated cppFlags fn = do
    mods <- parseImports <$> preprocess cppFlags fn
    files <- Maybe.catMaybes <$> mapM (fileOf generated) mods
    return (files, Maybe.mapMaybe (`Map.lookup` packageDb) mods)

-- | Load cached module to package db.
loadPackageDb :: FilePath -> Shake.Action PackageDb
loadPackageDb fn = do
    Shake.need [fn]
    Trans.liftIO $ do
        result <- Serialize.decode <$> B.readFile fn
        either (Util.errorIO . (("load from " ++ fn ++ ":")++)) return result

-- | Call 'getModuleToPackage' and save its contents to the file.
-- If there are colliding modules, throw an IO exception.
savePackageDb :: [Package] -> FilePath -> IO ()
savePackageDb packages cacheFn = do
    (packageDb, collisions) <- getModuleToPackage packages
    unless (null collisions) $
        Util.errorIO $ "modules found in >1 package:\n" ++ unlines
            [B.unpack k ++ ": " ++ unwords vs | (k, vs) <- collisions]
    B.writeFile cacheFn $ Serialize.encode packageDb

getModuleToPackage :: [Package]
    -> IO (Map.Map ModuleName Package, [(ModuleName, [Package])])
    -- ^ also return any modules found under multiple packages
getModuleToPackage packages = do
    packageMods <- zip packages <$> mapM getExposedModules packages
    return $ Maps.unique2
        [(mod, package) | (package, mods) <- packageMods, mod <- mods]

getExposedModules :: Package -> IO [ModuleName]
getExposedModules package =
    parse <$> processStdout "ghc-pkg" ["field", package, "exposed-modules"]
    where
  The first word is " exposed - modules : " .
  TODO gets duplicates if you have the same package installed under
    -- different versions.
    parse = drop 1 . B.words

-- | Like 'importsOf' but transitive.  Includes the given module.
--
  TODO Technically I should run CPP on the output of hsc2hs , which means
-- this should map the module names to the appropriate .hs and 'need' it.
  Otherwise the ' # include ' that belongs to hsc2hs will get processed by CPP .
transitiveImportsOf :: Generated -> (FilePath -> Maybe [String])
    -> FilePath -> Shake.Action [FilePath]
transitiveImportsOf generated cppFlagsOf fn = do
    Shake.need [fn]
    go Set.empty [fn]
    where
    go checked (fn:fns)
        | fn `Set.member` checked = go checked fns
        | otherwise = do
            imports <- importsOf generated (cppFlagsOf fn) fn
            let checked' = Set.insert fn checked
            go checked' (fns ++ filter (`Set.notMember` checked') imports)
    go checked [] = return $ Set.toList checked

fileOf :: Generated -> ModuleName -> IO (Maybe FilePath)
fileOf generated mod
    | fn `Set.member` _generatedHs generated = return $ Just fn
    | otherwise = findExistingFile $
        fn : map (FilePath.replaceExtension fn) (_generatedExtensions generated)
    where
    fn = B.unpack $ B.map slash mod <> ".hs"
    slash c = if c == '.' then '/' else c

-- | Any line that starts with @import@ should be an import.  It's a reserved
-- word so I think that's safe?
parseImports :: B.ByteString -> [ModuleName]
parseImports = Maybe.mapMaybe (parse . B.words) . B.lines
    where
    parse ("import" : "qualified" : mod : _) = Just mod
    parse ("import" : mod : _) = Just mod
    parse _ = Nothing

  | Read the file , and preprocess with CPPHS if cppFlags are given .
preprocess :: Maybe [String] -> FilePath -> IO B.ByteString
preprocess Nothing fn = B.readFile fn
preprocess (Just flags) fn = processStdout "cpphs" ("--cpp" : flags ++ [fn])


-- * util

processStdout :: FilePath -> [String] -> IO B.ByteString
processStdout cmd args = Exception.bracket open IO.hClose B.hGetContents
    where
    open = do
        (_, Just stdout, _, _) <- loggedProcess $
            (Process.proc cmd args) { Process.std_out = Process.CreatePipe }
        return stdout

  | Like '  ' , but actually report when the binary is n't
-- found.
loggedProcess :: Process.CreateProcess -> IO (Maybe IO.Handle,
       Maybe IO.Handle, Maybe IO.Handle, Process.ProcessHandle)
loggedProcess create = do
    r@(_, _, _, pid) <- Process.createProcess create
    Concurrent.forkIO $ do
        code <- Process.waitForProcess pid
        case code of
            Exit.ExitFailure c -> IO.hPutStrLn IO.stderr $
                "subprocess " ++ show (binaryOf create) ++ " failed: "
                ++ if c == 127 then "binary not found" else show c
            _ -> return ()
    return r
    where
    binaryOf create = case Process.cmdspec create of
        Process.RawCommand fn _ -> fn
        Process.ShellCommand cmd -> takeWhile (/=' ') cmd

findExistingFile :: [FilePath] -> IO (Maybe FilePath)
findExistingFile (fn:fns) =
    Util.ifM (Directory.doesFileExist fn) (return (Just fn))
        (findExistingFile fns)
findExistingFile [] = return Nothing

 This program is distributed under the terms of the GNU General Public
 License 3.0, see COPYING or -3.0.txt
# LANGUAGE OverloadedStrings #
 | Normally 'importsOf' and 'transitiveImportsOf' filter out files that don't
 exist, under the assumption that they belong to external packages.  This
 set contains exceptions to that, so they will get a need call, so they can
 be generated.
 | Try files with these extensions, in addition to .hs.
 E.g. [".hsc", ".chs"].
 | Find files of modules this module imports, in the form A/B.hs or A/B.hsc.
 Paths that don't exist are assumed to be package imports and are omitted.
 | Map a module name to the package it comes from.
 | Get local imports and package dependencies.

 I thought I wound use this but wound up not wanting it.  But I'll leave the
 basic implementation here in case I change my mind.  This still needs to
 be integrated with 'importsOf' and 'transitiveImportsOf'.
 | Load cached module to package db.
 | Call 'getModuleToPackage' and save its contents to the file.
 If there are colliding modules, throw an IO exception.
 ^ also return any modules found under multiple packages
 different versions.
 | Like 'importsOf' but transitive.  Includes the given module.

 this should map the module names to the appropriate .hs and 'need' it.
 | Any line that starts with @import@ should be an import.  It's a reserved
 word so I think that's safe?
 * util
 found.
  Copyright 2013  

module Shake.HsDeps (
    Generated(..)
    , importsOf, transitiveImportsOf
    , importsOfIO
    , importsPackagagesOf_
    , loadPackageDb, savePackageDb
) where
import qualified Control.Concurrent as Concurrent
import qualified Control.Exception as Exception
import qualified Control.Monad.Trans as Trans

import qualified Data.ByteString.Char8 as B
import qualified Data.Map as Map
import qualified Data.Maybe as Maybe
import qualified Data.Serialize as Serialize
import qualified Data.Set as Set

import qualified Development.Shake as Shake
import qualified System.Directory as Directory
import qualified System.Exit as Exit
import qualified System.FilePath as FilePath
import qualified System.IO as IO
import qualified System.Process as Process

import qualified Util.Maps as Maps
import qualified Util.Seq as Seq
import qualified Shake.Util as Util

import           Control.Monad


type Package = String
type ModuleName = B.ByteString

data Generated = Generated {
    _generatedHs :: Set.Set FilePath
    , _generatedExtensions :: [FilePath]
    } deriving (Eq, Show)

importsOf :: Generated -> Maybe [String] -> FilePath -> Shake.Action [FilePath]
importsOf generated cppFlags fn = do
    Shake.need [fn]
    Trans.liftIO (importsOfIO generated cppFlags fn)

importsOfIO :: Generated
  ^ If Just , first run CPP with these flags .
    -> FilePath -> IO [FilePath]
importsOfIO generated cppFlags fn = do
  TODO get CcDeps.includesOf so I can need them too .
    mods <- Seq.unique_sort . parseImports <$> preprocess cppFlags fn
    Maybe.catMaybes <$> mapM (fileOf generated) mods

  * 

type PackageDb = Map.Map ModuleName Package

importsPackagagesOf_ :: PackageDb
  ^ If Just , first run CPP with these flags .
    -> FilePath -> IO ([FilePath], [Package])
importsPackagagesOf_ packageDb generated cppFlags fn = do
    mods <- parseImports <$> preprocess cppFlags fn
    files <- Maybe.catMaybes <$> mapM (fileOf generated) mods
    return (files, Maybe.mapMaybe (`Map.lookup` packageDb) mods)

loadPackageDb :: FilePath -> Shake.Action PackageDb
loadPackageDb fn = do
    Shake.need [fn]
    Trans.liftIO $ do
        result <- Serialize.decode <$> B.readFile fn
        either (Util.errorIO . (("load from " ++ fn ++ ":")++)) return result

savePackageDb :: [Package] -> FilePath -> IO ()
savePackageDb packages cacheFn = do
    (packageDb, collisions) <- getModuleToPackage packages
    unless (null collisions) $
        Util.errorIO $ "modules found in >1 package:\n" ++ unlines
            [B.unpack k ++ ": " ++ unwords vs | (k, vs) <- collisions]
    B.writeFile cacheFn $ Serialize.encode packageDb

getModuleToPackage :: [Package]
    -> IO (Map.Map ModuleName Package, [(ModuleName, [Package])])
getModuleToPackage packages = do
    packageMods <- zip packages <$> mapM getExposedModules packages
    return $ Maps.unique2
        [(mod, package) | (package, mods) <- packageMods, mod <- mods]

getExposedModules :: Package -> IO [ModuleName]
getExposedModules package =
    parse <$> processStdout "ghc-pkg" ["field", package, "exposed-modules"]
    where
  The first word is " exposed - modules : " .
  TODO gets duplicates if you have the same package installed under
    parse = drop 1 . B.words

  TODO Technically I should run CPP on the output of hsc2hs , which means
  Otherwise the ' # include ' that belongs to hsc2hs will get processed by CPP .
transitiveImportsOf :: Generated -> (FilePath -> Maybe [String])
    -> FilePath -> Shake.Action [FilePath]
transitiveImportsOf generated cppFlagsOf fn = do
    Shake.need [fn]
    go Set.empty [fn]
    where
    go checked (fn:fns)
        | fn `Set.member` checked = go checked fns
        | otherwise = do
            imports <- importsOf generated (cppFlagsOf fn) fn
            let checked' = Set.insert fn checked
            go checked' (fns ++ filter (`Set.notMember` checked') imports)
    go checked [] = return $ Set.toList checked

fileOf :: Generated -> ModuleName -> IO (Maybe FilePath)
fileOf generated mod
    | fn `Set.member` _generatedHs generated = return $ Just fn
    | otherwise = findExistingFile $
        fn : map (FilePath.replaceExtension fn) (_generatedExtensions generated)
    where
    fn = B.unpack $ B.map slash mod <> ".hs"
    slash c = if c == '.' then '/' else c

parseImports :: B.ByteString -> [ModuleName]
parseImports = Maybe.mapMaybe (parse . B.words) . B.lines
    where
    parse ("import" : "qualified" : mod : _) = Just mod
    parse ("import" : mod : _) = Just mod
    parse _ = Nothing

  | Read the file , and preprocess with CPPHS if cppFlags are given .
preprocess :: Maybe [String] -> FilePath -> IO B.ByteString
preprocess Nothing fn = B.readFile fn
preprocess (Just flags) fn = processStdout "cpphs" ("--cpp" : flags ++ [fn])



processStdout :: FilePath -> [String] -> IO B.ByteString
processStdout cmd args = Exception.bracket open IO.hClose B.hGetContents
    where
    open = do
        (_, Just stdout, _, _) <- loggedProcess $
            (Process.proc cmd args) { Process.std_out = Process.CreatePipe }
        return stdout

  | Like '  ' , but actually report when the binary is n't
loggedProcess :: Process.CreateProcess -> IO (Maybe IO.Handle,
       Maybe IO.Handle, Maybe IO.Handle, Process.ProcessHandle)
loggedProcess create = do
    r@(_, _, _, pid) <- Process.createProcess create
    Concurrent.forkIO $ do
        code <- Process.waitForProcess pid
        case code of
            Exit.ExitFailure c -> IO.hPutStrLn IO.stderr $
                "subprocess " ++ show (binaryOf create) ++ " failed: "
                ++ if c == 127 then "binary not found" else show c
            _ -> return ()
    return r
    where
    binaryOf create = case Process.cmdspec create of
        Process.RawCommand fn _ -> fn
        Process.ShellCommand cmd -> takeWhile (/=' ') cmd

findExistingFile :: [FilePath] -> IO (Maybe FilePath)
findExistingFile (fn:fns) =
    Util.ifM (Directory.doesFileExist fn) (return (Just fn))
        (findExistingFile fns)
findExistingFile [] = return Nothing
