{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NondecreasingIndentation #-}
{-# LANGUAGE TupleSections #-}
module GHC.Driver.Session.Units (initMake, initMulti) where

-- The official GHC API
import qualified GHC
import GHC              (parseTargetFiles, Ghc, GhcMonad(..))

import GHC.Driver.Env
import GHC.Driver.Errors
import GHC.Driver.Errors.Types
import GHC.Driver.Phases
import GHC.Driver.Session
import GHC.Driver.Ppr
import GHC.Driver.Pipeline  ( oneShot, compileFile )
import GHC.Driver.Config.Diagnostic

import GHC.Unit.Env
import GHC.Unit (UnitId)
import GHC.Unit.Home.PackageTable
import qualified GHC.Unit.Home.Graph as HUG
import GHC.Unit.State  ( emptyUnitState )
import qualified GHC.Unit.State as State

import GHC.Types.SrcLoc
import GHC.Types.SourceError

import GHC.Utils.Misc
import GHC.Utils.Panic
import GHC.Utils.Outputable as Outputable
import GHC.Utils.Monad       ( liftIO, mapMaybeM )
import GHC.Data.Maybe

import System.IO
import System.Exit
import System.FilePath
import Control.Monad
import Data.List ( partition, (\\) )
import qualified Data.Set as Set
import Prelude
import GHC.ResponseFile (expandResponse)
import Data.Bifunctor
import GHC.Data.Graph.Directed
import qualified Data.List.NonEmpty as NE

-- Strip out any ["+RTS", ..., "-RTS"] sequences in the command string list.
removeRTS :: [String] -> [String]
removeRTS ("+RTS" : xs)  =
  case dropWhile (/= "-RTS") xs of
    [] -> []
    (_ : ys) -> removeRTS ys
removeRTS (y:ys)         = y : removeRTS ys
removeRTS []             = []

initMake :: [(String,Maybe Phase)] -> Ghc [(String, Maybe Phase)]
initMake srcs  = do
    let (hs_srcs, non_hs_srcs) = partition isHaskellishTarget srcs

    hsc_env <- GHC.getSession

    -- if we have no haskell sources from which to do a dependency
    -- analysis, then just do one-shot compilation and/or linking.
    -- This means that "ghc Foo.o Bar.o -o baz" links the program as
    -- we expect.
    if (null hs_srcs)
       then liftIO (oneShot hsc_env NoStop srcs) >> return []
       else do

    o_files <- mapMaybeM (\x -> liftIO $ compileFile hsc_env NoStop x)
                 non_hs_srcs
    dflags <- GHC.getSessionDynFlags
    let dflags' = dflags { ldInputs = map (FileOption "") o_files
                                      ++ ldInputs dflags }
    _ <- GHC.setSessionDynFlags dflags'
    return hs_srcs

initMulti :: NE.NonEmpty String
          -> (DynFlags -> [(String,Maybe Phase)] -> [String] -> [String] -> IO ())
          -- ^ Function to lint initMulti DynFlags and sources.
          -- In GHC, this is instanced to @checkOptions@.
          -> Ghc ([(String, Maybe UnitId, Maybe Phase)])
initMulti unitArgsFiles lintDynFlagsAndSrcs = do
  hsc_env <- GHC.getSession
  let logger = hsc_logger hsc_env
  initial_dflags <- GHC.getSessionDynFlags

  dynFlagsAndSrcs <- forM unitArgsFiles $ \f -> do
    when (verbosity initial_dflags > 2) (liftIO $ print f)
    args <- liftIO $ expandResponse [f]
    (dflags2, fileish_args, warns) <- parseDynamicFlagsCmdLine logger initial_dflags (map (mkGeneralLocated f) (removeRTS args))
    handleSourceError (\e -> do
       GHC.printException e
       liftIO $ exitWith (ExitFailure 1)) $ do
         liftIO $ printOrThrowDiagnostics logger (initPrintConfig dflags2) (initDiagOpts dflags2) (GhcDriverMessage <$> warns)

    let (dflags3, srcs, objs) = parseTargetFiles dflags2 (map unLoc fileish_args)
        dflags4 = offsetDynFlags dflags3

    let (hs_srcs, non_hs_srcs) = partition isHaskellishTarget srcs

    -- This is dubious as the whole unit environment won't be set-up correctly, but
    -- that doesn't matter for what we use it for (linking and oneShot)
    let dubious_hsc_env = hscSetFlags dflags4 hsc_env
    -- if we have no haskell sources from which to do a dependency
    -- analysis, then just do one-shot compilation and/or linking.
    -- This means that "ghc Foo.o Bar.o -o baz" links the program as
    -- we expect.
    if (null hs_srcs)
       then liftIO (oneShot dubious_hsc_env NoStop srcs) >> return (dflags4, [])
       else do

    o_files <- mapMaybeM (\x -> liftIO $ compileFile dubious_hsc_env NoStop x)
                 non_hs_srcs
    let dflags5 = dflags4 { ldInputs = map (FileOption "") o_files
                                      ++ ldInputs dflags4 }

    liftIO $ lintDynFlagsAndSrcs dflags5 srcs objs []

    pure (dflags5, hs_srcs)

  let
    unitDflags = NE.map fst dynFlagsAndSrcs
    srcs = NE.map (\(dflags, lsrcs) -> map (uncurry (,Just $ homeUnitId_ dflags,)) lsrcs) dynFlagsAndSrcs
    (hs_srcs, _non_hs_srcs) = unzip (map (partition (\(file, _uid, phase) -> isHaskellishTarget (file, phase))) (NE.toList srcs))

  checkDuplicateUnits initial_dflags (NE.toList (NE.zip unitArgsFiles unitDflags))

  (initial_home_graph, mainUnitId) <- liftIO $ createUnitEnvFromFlags unitDflags
  let home_units = HUG.allUnits initial_home_graph

  home_unit_graph <- forM initial_home_graph $ \homeUnitEnv -> do
    let cached_unit_dbs = homeUnitEnv_unit_dbs homeUnitEnv
        hue_flags = homeUnitEnv_dflags homeUnitEnv
        dflags = homeUnitEnv_dflags homeUnitEnv
    (dbs,unit_state,home_unit,mconstants) <- liftIO $ State.initUnits logger hue_flags cached_unit_dbs home_units

    updated_dflags <- liftIO $ updatePlatformConstants dflags mconstants
    emptyHpt <- liftIO $ emptyHomePackageTable
    pure $ HomeUnitEnv
      { homeUnitEnv_units = unit_state
      , homeUnitEnv_unit_dbs = Just dbs
      , homeUnitEnv_dflags = updated_dflags
      , homeUnitEnv_hpt = emptyHpt
      , homeUnitEnv_home_unit = Just home_unit
      }

  checkUnitCycles initial_dflags home_unit_graph

  let dflags = homeUnitEnv_dflags $ HUG.unitEnv_lookup mainUnitId home_unit_graph
  unitEnv <- assertUnitEnvInvariant <$> (liftIO $ initUnitEnv mainUnitId home_unit_graph (ghcNameVersion dflags) (targetPlatform dflags))
  let final_hsc_env = hsc_env { hsc_unit_env = unitEnv }

  GHC.setSession final_hsc_env

  -- if we have no haskell sources from which to do a dependency
  -- analysis, then just do one-shot compilation and/or linking.
  -- This means that "ghc Foo.o Bar.o -o baz" links the program as
  -- we expect.
  if (null hs_srcs)
      then do
        liftIO $ hPutStrLn stderr $ "Multi Mode can not be used for one-shot mode."
        liftIO $ exitWith (ExitFailure 1)
      else do

{-
  o_files <- liftIO $ mapMaybeM
                (\(src, uid, mphase) ->
                  compileFile (hscSetActiveHomeUnit (ue_unitHomeUnit (fromJust uid) unitEnv) final_hsc_env) NoStop (src, mphase)
                )
                (concat non_hs_srcs)
                -}

  -- MP: This should probably modify dflags for each unit?
  --let dflags' = dflags { ldInputs = map (FileOption "") o_files
  --                                  ++ ldInputs dflags }
  return $ concat hs_srcs

checkUnitCycles :: DynFlags -> HUG.HomeUnitGraph -> Ghc ()
checkUnitCycles dflags graph = processSCCs (HUG.hugSCCs graph)
  where

    processSCCs [] = return ()
    processSCCs (AcyclicSCC _: other_sccs) = processSCCs other_sccs
    processSCCs (CyclicSCC uids: _) = throwGhcException $ CmdLineError $ showSDoc dflags (cycle_err uids)


    cycle_err uids =
      hang (text "Units form a dependency cycle:")
           2
           (one_err uids)

    one_err uids = vcat $
                    (map (\uid -> text "-" <+> ppr uid <+> text "depends on") start)
                    ++ [text "-" <+> ppr final]
      where
        start = init uids
        final = last uids

-- | Check that we don't have multiple units with the same UnitId.
checkDuplicateUnits :: DynFlags -> [(FilePath, DynFlags)] -> Ghc ()
checkDuplicateUnits dflags flags =
  unless (null duplicate_ids)
         (throwGhcException $ CmdLineError $ showSDoc dflags multi_err)

  where
    uids = map (second homeUnitId_) flags
    deduplicated_uids = ordNubOn snd uids
    duplicate_ids = Set.fromList (map snd uids \\ map snd deduplicated_uids)

    duplicate_flags = filter (flip Set.member duplicate_ids . snd) uids

    one_err (fp, home_uid) = text "-" <+> ppr home_uid <+> text "defined in" <+> text fp

    multi_err =
      hang (text "Multiple units with the same unit-id:")
           2
           (vcat (map one_err duplicate_flags))


offsetDynFlags :: DynFlags -> DynFlags
offsetDynFlags dflags =
  dflags { hiDir = c hiDir
         , objectDir  = c objectDir
         , stubDir = c stubDir
         , hieDir  = c hieDir
         , dumpDir = c dumpDir  }

  where
    c f = augment_maybe (f dflags)

    augment_maybe Nothing = Nothing
    augment_maybe (Just f) = Just (augment f)
    augment f | isRelative f, Just offset <- workingDirectory dflags = offset </> f
              | otherwise = f


createUnitEnvFromFlags :: NE.NonEmpty DynFlags -> IO (HomeUnitGraph, UnitId)
createUnitEnvFromFlags unitDflags = do
  unitEnvList <- forM unitDflags $ \dflags -> do
    emptyHpt <- emptyHomePackageTable
    let newInternalUnitEnv =
          HUG.mkHomeUnitEnv emptyUnitState Nothing dflags emptyHpt Nothing
    return (homeUnitId_ dflags, newInternalUnitEnv)
  let activeUnit = fst $ NE.head unitEnvList
  return (HUG.hugFromList (NE.toList unitEnvList), activeUnit)