module Stg.GHC.Backend where

-- Compiler
import GHC
import GHC.Paths ( libdir )
import GHC.Platform ( platformOS, osSubsectionsViaSymbols )
import GHC.Driver.CodeOutput
import GHC.Driver.Hooks
import GHC.Driver.Main
import GHC.Driver.Phases
import GHC.Driver.Pipeline
import GHC.Driver.Session
--import GHC.Driver.Types
import GHC.Utils.Error
import GHC.Utils.Outputable
import GHC.Builtin.Names (rOOT_MAIN)
import GHC.Unit.State

-- Stg Types
import GHC.Data.FastString
import GHC.Stg.FVs
import GHC.Stg.Lint
import GHC.Stg.Syntax
import GHC.Stg.Unarise
import GHC.Types.CostCentre
import GHC.Unit.Types
import GHC.Types.Name.Set
import GHC.Data.Stream (Stream)
import qualified GHC.Data.Stream as Stream

-- Core Passes
import GHC.Cmm
import GHC.Cmm.Info (cmmToRawCmm )
import GHC.StgToCmm (codeGen)
import GHC.Types.Unique.Supply ( mkSplitUniqSupply, initUs_ )
import GHC.StgToCmm.Types (CgInfos (..))

import Control.Monad.Trans
import Control.Monad

import Data.Binary
import qualified Data.Set as Set

import qualified Stg.Syntax as C
--import qualified Stg.GHC.Convert_9_2 as C

import Data.List (isSuffixOf)
import System.FilePath

-------------------------------------------------------------------------------
-- Module
-------------------------------------------------------------------------------

modl :: Module
modl = rOOT_MAIN

-------------------------------------------------------------------------------
-- Compilation
-------------------------------------------------------------------------------

data Backend = NCG | LLVM


compileToObject :: Backend -> Unit -> ModuleName -> ForeignStubs -> [TyCon] -> [StgTopBinding] -> FilePath -> IO ()
compileToObject backend unitId modName stubs tyCons topBinds_simple outputName = do
  runGhc (Just libdir) $ compileToObjectM backend unitId modName stubs tyCons topBinds_simple outputName

compileToObjectM :: Backend -> Unit -> ModuleName -> ForeignStubs -> [TyCon] -> [StgTopBinding] -> FilePath -> Ghc ()
compileToObjectM backend unitId modName stubs tyCons topBinds_simple outputName = do
  dflags <- getSessionDynFlags

  let ccs       = emptyCollectedCCs :: CollectedCCs
      hpc       = emptyHpcInfo False

      this_mod  = mkModule unitId modName :: Module

      -- backend
      (target, link, outAsmFName) = case backend of
        LLVM  -> (HscLlvm, LlvmOpt, outputName -<.> ".ll")
        NCG   -> (HscAsm, As False, outputName -<.> ".s")

      beforeStgbinFName = outputName -<.> ".before_unarise.stgbin"
      afterStgbinFName  = outputName -<.> ".after_unarise.stgbin"

  liftIO $ do
    pure ()
{-
    putStrLn "==== Writing before unarise STG ===="
    encodeFile beforeStgbinFName $ C.cvtModule [] [] "whole program stg before unarise" mainUnit (mkModuleName ":Main") topBinds_simple NoStubs []
-}
{-
    putStrLn "==== Lint STG ===="
    lintStgTopBindings dflags this_mod True "Manual" topBinds_simple
-}
  --us <- liftIO $ mkSplitUniqSupply 'u'
  --let topBinds = unarise us topBinds_simple
  let topBinds = topBinds_simple
{-
  liftIO $ do
    putStrLn "==== Writing after unarise STG ===="
    encodeFile afterStgbinFName $ C.cvtModule [] [] "whole program stg after unarise" mainUnit (mkModuleName ":Main") topBinds NoStubs []
-}
  -- Compile
  dflags <- getSessionDynFlags
  pkgs <- setSessionDynFlags $
    dflags { hscTarget = target, ghcLink = NoLink }
    `gopt_set`  Opt_KeepSFiles
    `gopt_set`  Opt_KeepLlvmFiles
--    `dopt_set`  Opt_D_dump_cmm
--    `dopt_set`  Opt_D_dump_cmm_raw
--    `dopt_set`  Opt_D_dump_cmm_from_stg
{-
    `dopt_set`  Opt_D_dump_timings
    `gopt_set`  Opt_DoStgLinting
    `gopt_set`  Opt_DoCmmLinting
-}
    `gopt_set`  Opt_SplitSections -- HINT: linker based dead code elimination
    `gopt_set`  Opt_NoModpak
    `gopt_set`  Opt_NoGhcStgapp

  env <- getSession
  liftIO $ do
    newGen dflags env outAsmFName this_mod stubs tyCons ccs topBinds hpc
    oneShot env StopLn [(outAsmFName, Just link)]
  pure ()


compileProgram :: Backend -> Bool -> [String] -> [String] -> [String] -> [String] -> ForeignStubs -> [TyCon] -> [StgTopBinding] -> IO ()
compileProgram backend noHsMain incPaths libPaths ldOpts clikeFiles stubs tyCons topBinds_simple = runGhc (Just libdir) $ do
  dflags <- getSessionDynFlags

  liftIO $ do

    -- save exported external STG
    let stgFName = "whole_program_original.stgbin"
    putStrLn $ "writing " ++ stgFName
    --encodeFile stgFName $ C.cvtModule "whole-program-stg" mainUnit (mkModuleName ":Main") Nothing topBinds_simple NoStubs []

    putStrLn "==== STG ===="
--    putStrLn $ showSDoc dflags $ pprStgTopBindings topBinds_simple
    putStrLn "==== Lint STG ===="
    lintStgTopBindings dflags modl False "Manual" topBinds_simple

  us <- liftIO $ mkSplitUniqSupply 'g'
  let topBinds = unarise us topBinds_simple

  liftIO $ do
    let stgFName = "whole_program_unarised.stgbin"
    putStrLn $ "writing " ++ stgFName
    --encodeFile stgFName $ C.cvtModule "whole-program-stg" mainUnit (mkModuleName ":Main") Nothing topBinds NoStubs []

  -- construct STG program manually
  -- TODO: specify the following properly
{-
type CollectedCCs
  = ( [CostCentre]       -- local cost-centres that need to be decl'd
    , [CostCentreStack]  -- pre-defined "singleton" cost centre stacks
    )
-}
  let ccs       = emptyCollectedCCs :: CollectedCCs
      hpc       = emptyHpcInfo False

  -- backend
  let
    outFname  = "out.ll"

    (target, link) = case backend of
      LLVM  -> (HscLlvm, LlvmOpt)
      NCG   -> (HscAsm, As False)

  -- WORKAROUND: filter out rts includes
  let incPathsFixed = [p | p <- incPaths, not (isSuffixOf "rts-1.0/include" p)]

  -- Compile & Link
  dflags <- getSessionDynFlags
  setSessionDynFlags $
    (if noHsMain then flip gopt_set Opt_NoHsMain else id) $
    dflags
      { hscTarget     = target
      , ghcLink       = LinkBinary
      , libraryPaths  = libraryPaths dflags ++ libPaths
      , ldInputs      = ldInputs dflags ++ map Option ldOpts
      , includePaths  = addQuoteInclude (includePaths dflags) incPathsFixed
      }
    `gopt_set`  Opt_KeepSFiles
    `gopt_set`  Opt_KeepLlvmFiles
--    `dopt_set`  Opt_D_dump_cmm
--    `dopt_set`  Opt_D_dump_cmm_raw
--    `dopt_set`  Opt_D_dump_cmm_from_stg
    `dopt_set`  Opt_D_dump_timings
    `gopt_set`  Opt_DoStgLinting
    `gopt_set`  Opt_DoCmmLinting
    `gopt_set`  Opt_SplitSections -- HINT: linker based dead code elimination
    `gopt_set`  Opt_NoModpak
    `gopt_set`  Opt_NoGhcStgapp
  dflags <- getSessionDynFlags
  let pkgs = preloadUnits (unitState dflags)

  let libSet = Set.fromList ["rts"] -- "rts", "ghc-prim-cbits", "base-cbits", "integer-gmp-cbits"]
  dflags <- getSessionDynFlags
  let ignored_pkgs  = [IgnorePackage p |  p <- map (unpackFS . unitIdFS) pkgs, Set.notMember p libSet]
      my_pkgs       = [ExposePackage p (PackageArg p)  (ModRenaming True []) | p <- Set.toList libSet]
  setSessionDynFlags $ dflags { ignorePackageFlags = ignored_pkgs, packageFlags = my_pkgs }
  dflags <- getSessionDynFlags

  env <- getSession
  liftIO $ do
    (_, stub_c_exists, _, _) <- newGen dflags env outFname modl stubs tyCons ccs topBinds hpc
    let stubC = case stub_c_exists of
          Nothing -> []
          Just f  -> [(f, Nothing)]
    oneShot env StopLn $ (outFname, Just link) : [(c, Nothing) | c <- clikeFiles] ++ stubC
  pure ()

-------------
-- from GHC
-------------

newGen :: DynFlags
       -> HscEnv
       -> FilePath
       -> Module
       -> ForeignStubs
       -> [TyCon]
       -> CollectedCCs
       -> [StgTopBinding]
       -> HpcInfo
       -> IO (FilePath, Maybe FilePath, [(ForeignSrcLang, FilePath)], CgInfos)
newGen dflags hsc_env output_filename this_mod foreign_stubs data_tycons cost_centre_info stg_binds hpc_info = do
  -- TODO: add these to parameters
  let location = ModLocation
        { ml_hs_file  = Just $ output_filename -<.> ".hs"
        , ml_hi_file  = output_filename -<.> ".hi"
        , ml_obj_file = output_filename -<.> ".o"
        , ml_hie_file = output_filename -<.> ".hie"
        }
      foreign_files = []
      dependencies = [] -- only used for HscC target

  cmms <- {-# SCC "StgToCmm" #-}
                  doCodeGen hsc_env this_mod data_tycons
                      cost_centre_info
                      stg_binds hpc_info

  ------------------  Code output -----------------------
  rawcmms0 <- {-# SCC "cmmToRawCmm" #-}
            lookupHook (\x -> cmmToRawCmmHook x)
              (\dflg _ -> cmmToRawCmm dflg) dflags dflags (Just this_mod) cmms

  let dump a = do
        unless (null a) $
          dumpIfSet_dyn dflags Opt_D_dump_cmm_raw "Raw Cmm" FormatCMM (ppr a)
        return a
      rawcmms1 = Stream.mapM dump rawcmms0

  (output_filename, (stub_h_exists, stub_c_exists), stub_groups, foreign_fps, cg_infos)
      <- {-# SCC "codeOutput" #-}
        codeOutput dflags this_mod output_filename location
        foreign_stubs foreign_files dependencies rawcmms1
  return (output_filename, stub_c_exists, foreign_fps, cg_infos)
