{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeSynonymInstances #-}

module Development.Shake.ATS ( -- * Shake Rules
                               cgen
                             , cgenPretty
                             , cleanATS
                             , atsBin
                             , atsLex
                             , ghcExport
                             -- * Actions
                             , patsHome
                             -- Types
                             , Version (..)
                             ) where

import           Control.Monad
import           Control.Monad.IO.Class
import           Data.Bool                  (bool)
import           Data.Either                (fromRight)
import           Data.List                  (intercalate)
import           Data.Maybe                 (fromMaybe)
import           Data.Semigroup             (Semigroup (..))
import qualified Data.Text.Lazy             as TL
import           Development.Shake
import           Development.Shake.FilePath
import           Language.ATS
import           Language.ATS.Generate
import           System.Directory           (copyFile, createDirectoryIfMissing)
import           System.Exit                (ExitCode (ExitSuccess))

newtype Version = Version [Integer]
    deriving (Eq)

instance Show Version where
    show (Version is) = intercalate "." (show <$> is)

genATS :: FilePath -> FilePath -> Rules ()
genATS src target =
    target %> \out -> liftIO $ do
        contents <- readFile src
        writeFile out (generateATS contents)

ghcExport :: String -> Rules ()
ghcExport m =
    [ m ++ "_stub.h", m ++ ".o" ] &%> \[_,out] ->
        command [Cwd (takeDirectory out)] "ghc" ["-c", "-O", takeBaseName out -<.> "hs" ]

-- | The directory @~/.atspkg@
pkgHome :: Action String
pkgHome = fromMaybe "/usr/local/" <$> mh
    where mh = fmap (++ "/.atspkg/") <$> getEnv "HOME"

-- | The directory that will be @PATSHOME@.
patsHome :: Version -> Action String
patsHome v = fmap (++ (show v ++ "/")) pkgHome

atsCommand :: CmdResult r => Version -> Version -> String -> String -> Action r
atsCommand v v' sourceFile out = do
    h <- patsHome v'
    let home = h ++ "lib/ats2-postiats-" ++ show v
    let atsArgs = [EchoStderr False, AddEnv "PATSHOME" home]
        patsc = home ++ "/bin/patsopt"
    command atsArgs patsc ["--output", out, "-dd", sourceFile, "-cc"]

gcFlag :: Bool -> String
gcFlag False = "-DATS_MEMALLOC_LIBC"
gcFlag True  = "-DATS_MEMALLOC_GCBDW"

-- Copy source files to the appropriate place. This is necessary because
-- @#include@s in ATS are fucked up.
copySources :: Version -> Version -> [FilePath] -> Action ()
copySources v v' sources =
    forM_ sources $ \dep -> do
        h <- patsHome v'
        let home = h ++ "lib/ats2-postiats-" ++ show v
        liftIO $ createDirectoryIfMissing True (home ++ "/" ++ takeDirectory dep)
        liftIO $ copyFile dep (home ++ "/" ++ dep)

patsHomeLocs :: Int -> String
patsHomeLocs n = intercalate ":" $ (<> ".atspkg/contrib") . ("./" <>) <$> g
    where g = [ join $ replicate i "../" | i <- [1..n] ]

atsBin :: String -- ^ C compiler we should use
       -> [String] -- ^ Flags to pass to the C compiler
       -> Version -- ^ Library version
       -> Version -- ^ Compiler version
       -> Bool -- ^ Whether to use the garbage collector
       -> [String] -- ^ A list of libraries against which to link
       -> String -- ^ Source file
       -> [String] -- ^ Haskell dependencies to track
       -> [(String, String)] -- ^ Files to be run through @hs2ats@.
       -> String -- ^ Binary target
       -> Rules ()
atsBin cc cFlags v v' gc libs sourceFile hs atg out = do

    unless (null atg) $
        mapM_ (uncurry genATS) atg

    unless (null hs) $
        mapM_ ghcExport ((-<.> "") <$> hs)

    out %> \_ -> do
        sources <- transitiveDeps [sourceFile]
        h <- patsHome v'
        h' <- pkgHome
        let home = h ++ "lib/ats2-postiats-" ++ show v
        need (sources ++ hs)
        copySources v v' sources
        cmd_ ["mkdir", "-p", dropDirectory1 out]
        path <- fromMaybe "" <$> getEnv "PATH"
        let toLibs = fmap ("-l" <>)
        let libs' = bool libs ("gc" : libs) gc
        command
            [EchoStderr False, AddEnv "PATSHOME" home, AddEnv "PATH" (home ++ "/bin:" ++ path), AddEnv "PATSHOMELOCS" $ patsHomeLocs 5]
            (home ++ "/bin/patscc")
            ([sourceFile, "-atsccomp", cc ++ " -I" ++ h ++ "/ccomp/runtime/ -I" ++ h ++ " -L" ++ h' ++ "/lib", gcFlag gc, "-o", out, "-cleanaft"] <> cFlags <> toLibs libs')

-- | Build a @.lats@ file.
atsLex :: Rules ()
atsLex =
    "*.dats" %> \out -> do
        lats <- liftIO $ readFile (out -<.> "lats")
        (Stdout contents) <- command [Stdin lats] "atslex" []
        liftIO $ writeFile out contents

cleanATS :: Rules ()
cleanATS =

    "clean" ~> do
        removeFilesAfter "." ["//*.c", "//tags"]
        removeFilesAfter ".atspkg" ["//*"]
        removeFilesAfter "ats-deps" ["//*"]

handleSource :: Version -> Version -> FilePath -> Action ()
handleSource v v' sourceFile = do
        sources <- transitiveDeps [sourceFile]
        need sources
        copySources v v' sources

-- | This provides rules for generating C code from ATS source files in the
-- @ats-src@ directory.
cgen :: Version -- ^ Library version
     -> Version -- ^ Compiler version
     -> FilePath -- ^ Directory containing ATS source code
     -> Rules ()
cgen v v' dir =

    "//*.c" %> \out -> do
        let sourceFile = dir ++ "/" ++ (takeBaseName out -<.> "dats")
        handleSource v v' sourceFile
        atsCommand v v' sourceFile out

fixDir :: FilePath -> String -> String
fixDir p =
      TL.unpack
    . TL.replace (TL.pack "./") (TL.pack $ p ++ "/")
    . TL.replace (TL.pack "../") (TL.pack $ joinPath (init $ splitPath p) ++ "/")
    . TL.pack

trim :: String -> String
trim = init . drop 1

transitiveDeps :: [FilePath] -> Action [FilePath]
transitiveDeps [] = pure []
transitiveDeps ps = fmap join $ forM ps $ \p -> do
    contents <- liftIO $ readFile p
    let ats = fromRight mempty . parseATS . lexATS $ contents
    let dir = takeDirectory p
    deps <- filterM doesFileExist $ fixDir dir . trim <$> getDependencies ats
    deps' <- transitiveDeps deps
    pure $ (p:deps) ++ deps'

-- | This uses @pats-filter@ to prettify the errors.
cgenPretty :: Version -- ^ Library version
           -> Version -- ^ Compiler version
           -> FilePath
           -> Rules ()
cgenPretty v v' dir =

    "//*.c" %> \out -> do

        let sourceFile = dir ++ "/" ++ (takeBaseName out -<.> "dats")
        handleSource v v' sourceFile
        (Exit c, Stderr err) :: (Exit, Stderr String) <- atsCommand v v' sourceFile out
        cmd_ [Stdin err] Shell "pats-filter"
        if c /= ExitSuccess
            then error "patscc failure"
            else pure ()