{-# LANGUAGE RecordWildCards, TupleSections #-}
module Main (main) where

import qualified Paths_ghc_lib as GHC

import GHC.Driver.Monad
import GHC
    ( runGhc, defaultErrorHandler
    , setTargets
    , parseModule, typecheckModule, desugarModule, coreModule
    )

import GHC.Core.TyCon
import GHC.CoreToStg.Prep
import GHC.Data.Bag
import GHC.Data.FastString
import GHC.Driver.Main (hscSimplify)
import GHC.Driver.Make
import GHC.Driver.Phases
import GHC.Driver.Session
import GHC.Driver.Types
import GHC.Hs.Instances ()
import GHC.Hs.Instances ()
import GHC.Iface.Tidy (tidyProgram)
import GHC.Iface.UpdateIdInfos
import GHC.StgToCmm.Types (CgInfos(..))
import GHC.Types.Avail
import GHC.Types.Id
import GHC.Types.Name
import GHC.Types.Name.Set
import GHC.Unit.Module.Env
import GHC.Unit.Module.Name
import GHC.Unit.State
import GHC.Unit.Types
import GHC.Utils.Error
import GHC.Utils.Outputable (ppr, showSDoc)

import qualified Data.Map as M
import Data.Time
import Control.Monad (unless)
import Data.Either (partitionEithers)
import System.Exit
import System.FilePath

import qualified GHC.Data.EnumSet as EnumSet
import qualified GHC.LanguageExtensions.Type as Ext

main :: IO ()
main = defaultErrorHandler defaultFatalMessager defaultFlushOut $ do
    libDir <- GHC.getDataDir
    runGhc (Just libDir) $ withSourceErrors $ do
        setup
        compileModule =<< prepareModule "Inline"

withSourceErrors :: (GhcMonad m) => m a -> m a
withSourceErrors = handleSourceError $ \e -> do
    printException e
    liftIO $ exitWith $ ExitFailure 1

setup :: (GhcMonad m) => m ()
setup = do
    dflags <- getSessionDynFlags
    dflags <- return $ gopt_set dflags Opt_NoTypeableBinds
    dflags <- return $ dflags
        { hscTarget = HscNothing
        , mainModIs = noMainModule
        , packageDBFlags = [PackageDB $ PkgDbPath primPkgDb, ClearPackageDBs]
        }

    dflags <- liftIO $ initUnits dflags
    modifySession $ \env -> env{ hsc_dflags = dflags }
    invalidateModSummaryCache
  where
    primPkgDb = "/home/mi/.ghcup/ghc/9.0.1/lib/ghc-9.0.1/package.conf.d"

noMainModule :: Module
noMainModule = mkModule HoleUnit $ mkModuleName "Main"

noPrelude :: HscEnv -> HscEnv
noPrelude env = env
    { hsc_dflags = let dflags = hsc_dflags env in dflags
        { extensionFlags = EnumSet.delete Ext.ImplicitPrelude $ extensionFlags dflags
        }
    }

prepareModule :: (GhcMonad m) => String -> m ModSummary
prepareModule modName = do
    let mod = mkModuleName modName
        target = resolve mod
    setTargets [target]

    -- Anything already in the provided map should be left as-is
    providers <- do
        env <- getSession
        return $ moduleNameProvidersMap . unitState . hsc_dflags $ env
    let exclude = M.keys providers

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

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

    let Just ms = lookupModuleEnv menv $ mkModule mainUnit mod
    return ms

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

compileModule
    :: (GhcMonad m)
    => ModSummary
    -> m ()
compileModule ms = do
    pmod <- parseModule ms
    tmod <- typecheckModule pmod
    dmod <- desugarModule tmod
    let mguts = coreModule dmod

    env <- getSession
    mguts' <- liftIO $ hscSimplify env [] mguts
    (cg_guts, details) <- liftIO $ tidyProgram env mguts'

    let binds = cg_binds cg_guts
        tycons = cg_tycons cg_guts
        data_tycons = filter isDataTyCon tycons
        mod = cg_module cg_guts
        modLoc = ms_location ms
    (binds, _ccs) <- liftIO $ corePrepPgm env mod modLoc binds data_tycons

    let cg_info = CgInfos
            { cgNonCafs = mempty
            , cgLFInfos = mempty
            }
        details' = updateModDetailsIdInfos (hsc_dflags env) cg_info details

    let type_env = md_types details'
        lookupVar v = lookupTypeEnv type_env =<< exportedName details' varName v
    liftIO $ case lookupVar "f" of
        Just (AnId f) -> putStrLn $ showSDoc (hsc_dflags env) . ppr $ idUnfolding f
        _ -> putStrLn "No `f' in module"

    return ()

resolve :: ModuleName -> Target
resolve mod = mkTarget $ "input" </> path <.> "src"
  where
    path = moduleNameSlashes mod

    mkTarget filePath = Target
        { targetId = TargetFile filePath (Just $ Cpp HsSrcFile)
        , targetAllowObjCode = False
        , targetContents = Nothing
        }

invalidateModSummaryCache :: (GhcMonad m) => m ()
invalidateModSummaryCache = modifySession $ \env -> env
    { hsc_mod_graph = invalidateMG (hsc_mod_graph env)
    }
  where
    invalidateMG = mapMG invalidateMS
    invalidateMS ms = ms{ ms_hs_date = addUTCTime (-1) (ms_hs_date ms) }

exportedName :: ModDetails -> NameSpace -> String -> Maybe Name
exportedName ModDetails{..} = \ns s ->
    let occ = mkOccNameFS ns (fsLit s)
    in case nameSetElemsStable . filterNameSet (\nm -> occName nm == occ) $ names of
        [nm] -> Just nm
        _ -> Nothing
  where
    names = availsToNameSet md_exports