module Main (main) where

import qualified Paths_ghc_lib as GHC
import Paths_hello

import GHC hiding (loadModule)
import GHC.Driver.Monad
import GHC.Driver.Session
import GHC.Unit.State
import GHC.Unit.Types
import GHC.Unit.Module.Name
import GHC.Unit.Module.Env
import GHC.Driver.Types
import GHC.Driver.Make
import GHC.Tc.Types
import qualified Data.Map as M
import Control.Monad
import Data.Maybe
import GHC.Iface.Make
import GHC.Utils.Error
import GHC.Data.Bag
import Data.Either
import System.FilePath
import GHC.Data.FastString

myUnitId :: UnitId
myUnitId = UnitId $ fsLit "my-unit"

myUnit :: Unit
myUnit = RealUnit (Definite myUnitId)

main :: IO ()
main = do
    libdir <- GHC.getDataDir
    runGhc (Just libdir) $ do
        setup

        let fileTarget mod = Target
                { targetId = TargetFile ("input/lookupVers1/" <> moduleNameSlashes mod <> ".hs") Nothing
                , targetAllowObjCode = False
                , targetContents = Nothing
                }

        setTargets []
        let prepare unit mod = do
                addTarget $ fileTarget mod

                (errs, mss) <- do
                    env <- getSession
                    env <- return $ setHomeUnit (toUnitId unit) env
                    liftIO $ partitionEithers <$> downsweep env [] [] False
                reportErrors errs

                let menv = mkModuleEnv [(ms_mod ms, ms) | ms <- mss]

                return $ fromMaybe (error $ unwords ["lookupModuleEnv", unitString unit, moduleNameString mod]) $
                    lookupModuleEnv menv $ mkModule unit mod

        let load ms = do
                -- Loading dependencies has changed the unitState, so we
                -- need to refresh the one stored in the ModSummary
                ms <- do
                    us <- unitState . hsc_dflags <$> getSession
                    return $ withUnitState us ms
                (iface, details) <- prepareSource ms
                modifySession $ registerModule iface details

        ms1 <- prepare mainUnit $ mkModuleName "Top"
        ms2 <- prepare myUnit $ mkModuleName "MyPrim"
        load ms2
        load ms1

        liftIO $ putStrLn "All done."

reportErrors :: (GhcMonad m) => [ErrorMessages] -> m ()
reportErrors errs = do
    errs <- return $ unionManyBags errs
    unless (isEmptyBag errs) $ throwErrors errs

prepareSource
    :: (GhcMonad m)
    => ModSummary
    -> m (ModIface, ModDetails)
prepareSource ms = do
    pmod <- parseModule ms
    tmod <- typecheckModule pmod

    let (tcg, details) = tm_internals_ tmod
    env <- getSession
    iface <- liftIO $ mkIfaceTc env Sf_Ignore details tcg
    return (iface, details)

setup :: (GhcMonad m) => m ()
setup = do
    baseDir <- liftIO getDataDir
    let pkgDb = baseDir </> "ghc-packages/package.conf.d"

    dflags <- getSessionDynFlags
    dflags <- return $ gopt_set dflags Opt_NoTypeableBinds
    dflags <- return $ dflags{ hscTarget = HscNothing }
    dflags <- return $ dflags{ packageDBFlags = [PackageDB $ PkgDbPath pkgDb, ClearPackageDBs] }
    dflags <- liftIO $ initUnits dflags
    modifySession $ \env -> env{ hsc_dflags = dflags }

setHomeUnit :: UnitId -> HscEnv -> HscEnv
setHomeUnit unitId env = env
    { hsc_dflags = (hsc_dflags env)
        { homeUnitId = unitId
        }
    }

modifyUnitState :: (UnitState -> UnitState) -> HscEnv -> HscEnv
modifyUnitState f env = env
    { hsc_dflags = let dflags = hsc_dflags env in dflags
        { unitState = f (unitState dflags)
        }
    }

registerModule :: ModIface -> ModDetails -> HscEnv -> HscEnv
registerModule iface details = extendHPT . addModule
  where
    mod_info = HomeModInfo iface details Nothing

    mod = mi_module iface
    modOrig = ModOrigin (Just True) [] [] True

    addModule = modifyUnitState $ \us -> us
        { moduleNameProvidersMap = M.insert (moduleName mod) (M.singleton mod modOrig) $ moduleNameProvidersMap us
        }

    extendHPT env = env{ hsc_HPT = hpt, hsc_type_env_var = Nothing }
      where
        hpt = addToHpt (hsc_HPT env) (moduleName mod) mod_info


withUnitState :: UnitState -> ModSummary -> ModSummary
withUnitState us ms = ms
    { ms_hspp_opts = (ms_hspp_opts ms)
        { unitState = us
        }
    }