-- ghc-mod: Making Haskell development *more* fun
-- Copyright (C) 2015  Daniel Gröber <dxld ÄT darkboxed DOT org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU Affero General Public License for more details.
--
-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.

{-# LANGUAGE ScopedTypeVariables, RecordWildCards #-}
module Language.Haskell.GhcMod.HomeModuleGraph (
   GmModuleGraph(..)
 , ModulePath(..)
 , mkFileMap
 , mkModuleMap
 , mkMainModulePath
 , findModulePath
 , findModulePathSet
 , fileModuleName
 , canonicalizeModulePath
 , homeModuleGraph
 , updateHomeModuleGraph
 , canonicalizeModuleGraph
 , reachable
 , moduleGraphToDot
 ) where

import DriverPipeline
import DynFlags
import ErrUtils
import Exception
import Finder
import GHC
import HscTypes

import Control.Arrow ((&&&))
import Control.Applicative
import Control.Monad
import Control.Monad.Trans.Maybe (MaybeT(..), runMaybeT)
import Control.Monad.State.Strict (execStateT)
import Control.Monad.State.Class
import Data.Maybe
import Data.Monoid as Monoid
import Data.Map  (Map)
import qualified Data.Map  as Map
import Data.Set (Set)
import qualified Data.Set as Set
import System.FilePath
import System.Directory
import System.IO
import Prelude

import Language.Haskell.GhcMod.Logging
import Language.Haskell.GhcMod.Logger
import Language.Haskell.GhcMod.Monad.Types
import Language.Haskell.GhcMod.Types
import Language.Haskell.GhcMod.Utils (withMappedFile)
import Language.Haskell.GhcMod.Gap (parseModuleHeader)

-- | Turn module graph into a graphviz dot file
--
-- @dot -Tpng -o modules.png modules.dot@
moduleGraphToDot :: GmModuleGraph -> String
moduleGraphToDot GmModuleGraph { gmgGraph } =
    "digraph {\n" ++ concatMap edges (Map.toList graph) ++ "}\n"
 where
   graph = Map.map (Set.mapMonotonic mpPath)
         $ Map.mapKeysMonotonic mpPath gmgGraph
   edges :: (FilePath, (Set FilePath)) -> String
   edges (f, sf) =
       concatMap (\f' -> "    \""++ f ++"\" -> \""++ f' ++"\"\n") (Set.toList sf)

data S = S {
      sErrors   :: [(ModulePath, ErrorMessages)],
      sWarnings :: [(ModulePath, WarningMessages)],
      sGraph    :: GmModuleGraph
}

defaultS :: S
defaultS = S [] [] mempty

putErr :: MonadState S m
       => (ModulePath, ErrorMessages) -> m ()
putErr e = do
  s <- get
  put s { sErrors = e:sErrors s}

putWarn :: MonadState S m
       => (ModulePath, ErrorMessages) -> m ()
putWarn w = do
  s <- get
  put s { sWarnings = w:sWarnings s}

gmgLookupMP :: MonadState S m => ModulePath -> m (Maybe (Set ModulePath))
gmgLookupMP k = (Map.lookup k . gmgGraph . sGraph) `liftM` get

graphUnion :: MonadState S m => GmModuleGraph -> m ()
graphUnion gmg = do
  s <- get
  put s { sGraph = sGraph s `mappend` gmg }

reachable :: Set ModulePath -> GmModuleGraph -> Set ModulePath
reachable smp0 GmModuleGraph {..} = go smp0
 where
   go smp = let
       δsmp = Set.unions $
                collapseMaybeSet . flip Map.lookup gmgGraph <$> Set.toList smp
       smp' = smp `Set.union` δsmp
    in if smp == smp' then smp' else go smp'

pruneUnreachable :: Set ModulePath -> GmModuleGraph -> GmModuleGraph
pruneUnreachable smp0 gmg@GmModuleGraph {..} = let
    r = reachable smp0 gmg
  in
    GmModuleGraph {
      gmgGraph = Map.filterWithKey (\k _ -> k `Set.member` r) gmgGraph
    }

collapseMaybeSet :: Maybe (Set a) -> Set a
collapseMaybeSet = maybe Set.empty id

homeModuleGraph :: (IOish m, Gm m)
    => HscEnv -> Set ModulePath -> m GmModuleGraph
homeModuleGraph env smp = updateHomeModuleGraph env mempty smp smp

mkMainModulePath :: FilePath -> ModulePath
mkMainModulePath = ModulePath (mkModuleName "Main")

findModulePath :: HscEnv -> ModuleName -> IO (Maybe ModulePath)
findModulePath env mn = do
    fmap (ModulePath mn) <$> find env mn

findModulePathSet :: HscEnv -> [ModuleName] -> IO (Set ModulePath)
findModulePathSet env mns = do
    Set.fromList . catMaybes <$> findModulePath env `mapM` mns

find :: MonadIO m => HscEnv -> ModuleName -> m (Maybe FilePath)
find env mn = liftIO $ do
  res <- findHomeModule env mn
  case res of
   -- TODO: handle SOURCE imports (hs-boot stuff): addBootSuffixLocn loc
    Found loc@ModLocation { ml_hs_file = Just _ } _mod ->
        return $ normalise <$> ml_hs_file loc
    _ -> return Nothing


canonicalizeModulePath :: ModulePath -> IO ModulePath
canonicalizeModulePath (ModulePath mn fp) = ModulePath mn <$> canonicalizePath fp

canonicalizeModuleGraph :: MonadIO m => GmModuleGraph -> m GmModuleGraph
canonicalizeModuleGraph GmModuleGraph {..} = liftIO $ do
    GmModuleGraph . Map.fromList <$> mapM fmg (Map.toList gmgGraph)
 where
   fmg (mp, smp) = liftM2 (,) (canonicalizeModulePath mp) (Set.fromList <$> mapM canonicalizeModulePath (Set.toList smp))


updateHomeModuleGraph :: (IOish m, Gm m)
                      => HscEnv
                      -> GmModuleGraph
                      -> Set ModulePath -- ^ Initial set of modules
                      -> Set ModulePath -- ^ Updated set of modules
                      -> m GmModuleGraph
updateHomeModuleGraph env GmModuleGraph {..} smp sump = do
    -- TODO: It would be good if we could retain information about modules that
    -- stop to compile after we've already successfully parsed them at some
    -- point. Figure out a way to delete the modules about to be updated only
    -- after we're sure they won't fail to parse .. or something. Should probably
    -- push this whole prune logic deep into updateHomeModuleGraph'
   (pruneUnreachable smp . sGraph) `liftM` runS (updateHomeModuleGraph' env sump)
 where
   runS = flip execStateT defaultS { sGraph = graph' }
   graph' = GmModuleGraph {
       gmgGraph = Set.foldr Map.delete gmgGraph sump
    }

mkFileMap :: Set ModulePath -> Map FilePath ModulePath
mkFileMap smp = Map.fromList $ map (mpPath &&& id) $ Set.toList smp

mkModuleMap :: Set ModulePath -> Map ModuleName ModulePath
mkModuleMap smp = Map.fromList $ map (mpModule &&& id) $ Set.toList smp

updateHomeModuleGraph'
    :: forall m. (MonadState S m, IOish m, Gm m)
    => HscEnv
    -> Set ModulePath     -- ^ Initial set of modules
    -> m ()
updateHomeModuleGraph' env smp0 = do
    go `mapM_` Set.toList smp0
 where
   go :: ModulePath -> m ()
   go mp = do
     msmp <- gmgLookupMP mp
     case msmp of
       Just _ -> return ()
       Nothing -> do
           smp <- collapseMaybeSet `liftM` step mp

           graphUnion GmModuleGraph {
               gmgGraph = Map.singleton mp smp
            }

           mapM_ go (Set.toList smp)

   step :: ModulePath -> m (Maybe (Set ModulePath))
   step mp = runMaybeT $ do
       (dflags, ppsrc_fn) <- MaybeT preprocess'
       src <- liftIO $ readFile ppsrc_fn
       imports mp src dflags
    where
      preprocess' :: m (Maybe (DynFlags, FilePath))
      preprocess' = do
        let fn = mpPath mp
        ep <- preprocessFile env fn
        case ep of
          Right (_, x) -> return $ Just x
          Left errs -> do
            -- TODO: Remember these and present them as proper errors if this is
            -- the file the user is looking at.
            gmLog GmWarning ("preprocess " ++ show fn) $ Monoid.mempty $+$ (vcat $ map text errs)
            return Nothing


   imports :: ModulePath -> String -> DynFlags -> MaybeT m (Set ModulePath)
   imports mp@ModulePath {..} src dflags =
       case parseModuleHeader src dflags mpPath of
         Left err -> do
           putErr (mp, err)
           mzero

         Right (ws, lmdl) -> do
           putWarn (mp, ws)
           let HsModule {..} = unLoc lmdl
               mns = map (unLoc . ideclName)
                   $ filter (isNothing . ideclPkgQual)
                   $ map unLoc hsmodImports
           -- TODO: handle package qualifier "this"
           liftIO $ Set.fromList . catMaybes <$> mapM (findModulePath env) mns

preprocessFile :: (IOish m, GmEnv m, GmState m) =>
  HscEnv -> FilePath -> m (Either [String] ([String], (DynFlags, FilePath)))
preprocessFile env file =
  withLogger' env $ \setDf -> do
    withMappedFile file $ \fn -> do
      let env' = env { hsc_dflags = setDf (hsc_dflags env) }
      liftIO $ preprocess env' (fn, Nothing)

fileModuleName :: (IOish m, GmEnv m, GmState m) =>
  HscEnv -> FilePath -> m (Either [String] (Maybe ModuleName))
fileModuleName env fn = do
    let handler = liftIO . handle (\(_ :: SomeException) -> return $ Right Nothing)
    ep <- preprocessFile env fn
    case ep of
      Left errs -> do
        return $ Left errs
      Right (_warns, (dflags, procdFile)) -> leftM (errBagToStrList env) =<< handler (do
        src <- readFile procdFile
        case parseModuleHeader src dflags procdFile of
          Left errs -> return $ Left errs
          Right (_, lmdl) -> do
            let HsModule {..} = unLoc lmdl
            return $ Right $ unLoc <$> hsmodName)
  where
    leftM f = either (return . Left <=< f) (return . Right)