{-# LANGUAGE ScopedTypeVariables, OverloadedStrings #-}
{-# LANGUAGE RecordWildCards, BlockArguments, TupleSections #-}

module Main (main) where

import qualified Paths_ghc_lib as GHC

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.Driver.Types
import Data.Time
import GHC.Utils.Fingerprint
import GHC.Iface.Make (mkIface_)
import GHC.Tc.Types
import qualified Data.Map as M
import Control.Monad
import Control.Monad.State
import Data.Maybe
import GHC.Utils.Outputable hiding ((<>))

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

        let resolveTarget mod = do
                -- Pretend this does something more interesting
                return Target
                    { targetId = TargetFile ("input/fixity/" <> moduleNameSlashes mod <> ".hs") Nothing
                    , targetAllowObjCode = False
                    , targetContents = Nothing
                    }

        loadDependencies resolveTarget $ mkModuleName "Standalone"
        liftIO $ putStrLn "Success with Standalone.hs"
        loadDependencies resolveTarget $ mkModuleName "Importer"
        liftIO $ putStrLn "Success with Importer.hs"

-- This is very custom, because in my real use case I can't let GHC
-- load modules on its own from files.
loadDependencies :: forall m. (GhcMonad m) => (ModuleName -> m Target) -> ModuleName -> m ()
loadDependencies resolveTarget mod = evalStateT (go mod) []
  where
    go :: ModuleName -> StateT [ModuleName] m ()
    go mod = do
        done <- gets (mod `elem`)
        unless done do
            modify (mod:)
            ms <- lift do
                target <- resolveTarget mod
                addTarget target
                mg <- depanal [] False
                return $ fromMaybe (error "mgLookupModule") $
                    mgLookupModule mg (mkModule mainUnit mod)
            let deps = [mod | (_, L _ mod) <- ms_textual_imps ms]
            mapM_ go deps
            lift $ loadModule ms

setup :: (GhcMonad m) => m ()
setup = do
    dflags <- getSessionDynFlags
    dflags <- return $ gopt_set dflags Opt_NoTypeableBinds
    dflags <- return $ dflags{ hscTarget = HscNothing }
    dflags <- liftIO $ initUnits dflags
    modifySession $ \env -> env{ hsc_dflags = dflags }
    modifySession $ invalidateModSummaryCache

loadModule :: (GhcMonad m) => ModSummary -> m ()
loadModule ms = do
    liftIO $ putStrLn $ unwords ["Loading module", moduleNameString mod]
    tmod <- typecheckModule =<< parseModule ms
    env <- getSession
    let (tcg, details) = tm_internals_ tmod
        iface = mkModIface env tcg details
        mod_info = HomeModInfo iface details Nothing
    liftIO $ putStrLn $ unwords ["Fixities:", showPpr (hsc_dflags env) $ mi_fixities iface]
    modifySession $
        extendHpt mod_info .
        registerModule (mkModule mainUnit mod)
  where
    mod = moduleName . ms_mod $ ms

mkModIface :: HscEnv -> TcGblEnv -> ModDetails -> ModIface
mkModIface env tcg mod_details = fill partial
  where
    empty = emptyFullModIface this_mod

    fill :: PartialModIface -> ModIface
    fill partial = partial
        { mi_decls = map (fingerprint0,) (mi_decls partial)
        , mi_final_exts = mi_final_exts empty
        }

    partial = mkIface_ env
        this_mod hsc_src
        used_th deps rdr_env
        fix_env warns hpc_info
        (imp_trust_own_pkg imports) safe_mode usages
        doc_hdr' doc_map arg_map
        mod_details

    -- Copypasta from GHC.Iface.Make.mkIfaceTc
    this_mod = tcg_mod tcg
    hsc_src = HsSrcFile
    used_th = False
    deps = noDependencies -- TODO: dep tracking
    rdr_env = tcg_rdr_env tcg
    fix_env = tcg_fix_env tcg
    warns = tcg_warns tcg
    hpc_info = emptyHpcInfo False
    imports = tcg_imports tcg
    safe_mode = Sf_Ignore
    usages = [] -- TODO: dep tracking
    doc_hdr' = Nothing
    doc_map = emptyDeclDocMap
    arg_map = emptyArgDocMap

registerModule :: Module -> HscEnv -> HscEnv
registerModule mod env = env
    { hsc_dflags = let dflags = hsc_dflags env in dflags
        { unitState = addModules (unitState dflags)
        }
    }
  where
    addModules us = us
        { moduleNameProvidersMap = M.insert (moduleName mod) (M.singleton mod modOrig) $ moduleNameProvidersMap us
        }
    modOrig = ModOrigin (Just True) [] [] True

extendHpt :: HomeModInfo -> HscEnv -> HscEnv
extendHpt hmi env = env{ hsc_HPT = hpt', hsc_type_env_var = Nothing }
  where
    mod = mi_module . hm_iface $ hmi
    hpt' = addToHpt (hsc_HPT env) (moduleName mod) hmi

invalidateModSummaryCache :: HscEnv -> HscEnv
invalidateModSummaryCache 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) }