{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}
module Compile.Exec (
  KernelLib,
  buildKernel,
  callKernelFun,

  -- * misc
  lineNumbers,
) where

import Control.Monad (when)
import Data.IORef
import Foreign (Ptr)
import Foreign.Ptr (FunPtr)
import System.Directory (removeDirectoryRecursive)
import System.Environment (lookupEnv)
import System.Exit (ExitCode(..))
import System.IO (hPutStrLn, stderr)
import System.IO.Error (mkIOError, userErrorType)
import System.IO.Unsafe (unsafePerformIO)
import System.Posix.DynamicLinker
import System.Posix.Temp (mkdtemp)
import System.Process (readProcessWithExitCode)


debug :: Bool
debug = False

-- The IORef wrapper is required for the finalizer to attach properly (see the 'Weak' docs)
data KernelLib = KernelLib !(IORef (FunPtr (Ptr () -> IO ())))

buildKernel :: String -> String -> IO (KernelLib, String)
buildKernel csource funname = do
  template <- (++ "/tmp.chad.") <$> getTempDir
  path <- mkdtemp template

  let outso = path ++ "/out.so"
  let args = ["-O3", "-march=native"
             ,"-shared", "-fPIC"
             ,"-std=c99", "-x", "c"
             ,"-o", outso, "-"
             ,"-Wall", "-Wextra"
             ,"-Wno-unused-variable", "-Wno-unused-but-set-variable"
             ,"-Wno-unused-parameter", "-Wno-unused-function"
             ,"-Wno-alloc-size-larger-than"  -- ideally we'd keep this, but gcc reports false positives
             ,"-Wno-maybe-uninitialized"]  -- maximum1i goes out of range if its input is empty, yes, don't complain
  (ec, gccStdout, gccStderr) <- readProcessWithExitCode "gcc" args csource

  -- Print the source before the GCC output.
  case ec of
    ExitSuccess -> return ()
    ExitFailure{} -> hPutStrLn stderr $ "[chad] Kernel compilation failed! Source: <<<\n" ++ lineNumbers csource ++ ">>>"

  case ec of
    ExitSuccess -> return ()
    ExitFailure{} -> do
      removeDirectoryRecursive path
      ioError (mkIOError userErrorType "chad kernel compilation failed" Nothing Nothing)

  numLoaded <- atomicModifyIORef' numLoadedCounter (\n -> (n+1, n+1))
  when debug $ hPutStrLn stderr $ "[chad] loading kernel " ++ path ++ " (" ++ show numLoaded ++ " total)"
  dl <- dlopen outso [RTLD_LAZY, RTLD_LOCAL]

  removeDirectoryRecursive path  -- we keep a reference anyway because we have the file open now

  ref <- newIORef =<< dlsym dl funname
  _ <- mkWeakIORef ref (do numLeft <- atomicModifyIORef' numLoadedCounter (\n -> (n-1, n-1))
                           when debug $ hPutStrLn stderr $ "[chad] unloading kernel " ++ path ++ " (" ++ show numLeft ++ " left)"
                           dlclose dl)
  return (KernelLib ref, gccStdout ++ (if null gccStdout then "" else "\n") ++ gccStderr)

foreign import ccall "dynamic"
  wrapKernelFun :: FunPtr (Ptr () -> IO ()) -> Ptr () -> IO ()

-- Ensure that keeping a reference to the returned function also keeps the 'KernelLib' alive
{-# NOINLINE callKernelFun #-}
callKernelFun :: KernelLib -> Ptr () -> IO ()
callKernelFun (KernelLib ref) arg = do
  ptr <- readIORef ref
  wrapKernelFun ptr arg

getTempDir :: IO FilePath
getTempDir =
  lookupEnv "TMPDIR" >>= \case
    Just s | not (null s) -> return s
    _ -> return "/tmp"

{-# NOINLINE numLoadedCounter #-}
numLoadedCounter :: IORef Int
numLoadedCounter = unsafePerformIO $ newIORef 0

lineNumbers :: String -> String
lineNumbers str =
  let lns = lines str
      numlines = length lns
      width = length (show numlines)
      pad s = replicate (width - length s) ' ' ++ s
  in unlines (zipWith (\i ln -> pad (show i) ++ " | " ++ ln) [1::Int ..] lns)