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

module Development.Shake.ATS ( -- * Shake Rules
                               cgen
                             , cgenPretty
                             , cleanATS
                             , atsBin
                             , atsLex
                             , cabalExport
                             -- * Actions
                             , patsHome
                             -- * Helper functions
                             , getSubdirs
                             -- Types
                             , Version (..)
                             , ForeignCabal (..)
                             , BinaryTarget (..)
                             ) 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.ATS.Rules
import           Development.Shake.ATS.Type
import           Development.Shake.FilePath
import           Language.ATS
import           Language.ATS.Generate
import           System.Directory            (copyFile, createDirectoryIfMissing)
import           System.Exit                 (ExitCode (ExitSuccess))

-- | Given a plain Haskell source file, generate a @.sats@ file containing
-- analogous types.
genATS :: FilePath -- ^ Haskell source
       -> FilePath -- ^ @.sats@ file to generate
       -> Rules ()
genATS src target =
    target %> \out -> liftIO $ do
        createDirectoryIfMissing True (takeDirectory out)
        genATSTypes src target

-- | 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

-- | Run @patsopt@ given information about
atsCommand :: CmdResult r => Version -- ^ Compiler version
                          -> Version -- ^ Standard library version
                          -> String -- ^ Source file
                          -> String -- ^ C code to be generated
                          -> 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 weird.
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)

-- This is the @$PATSHOMELOCS@ variable to be passed to the shell.
patsHomeLocs :: Int -> String
patsHomeLocs n = intercalate ":" $ (<> ".atspkg/contrib") . ("./" <>) <$> g
    where g = [ join $ replicate i "../" | i <- [1..n] ]

-- TODO depend on GHC version?
makeCFlags :: [String] -- ^ Inputs
           -> [ForeignCabal] -- ^ Haskell libraries
           -> Bool -- ^ Whether to use the Garbage collector
           -> [String]
makeCFlags ss fc b = gcFlag' : (hsExtra <> ss) where
    gcFlag' = bool ("-optc" <>) id noHs $ gcFlag b
    hsExtra = bool ["--make", "-odir", ".atspkg", "-no-hs-main", "-package-db", "~/.cabal/store/ghc-8.2.2/package.db/"] mempty noHs
    noHs = null fc

-- TODO ideally cache C files and use `gcc` on them?
atsBin :: BinaryTarget -> Rules ()
atsBin BinaryTarget{..} = do

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

    unless (null hsLibs) $
        mapM_ cabalExport hsLibs

    binTarget %> \_ -> do
        h <- patsHome compilerVer
        h' <- pkgHome
        let home = h ++ "lib/ats2-postiats-" ++ show libVersion
        sources <- transitiveDeps (snd <$> genTargets) [src]
        need (sources ++ (TL.unpack . objectFile <$> hsLibs))
        copySources libVersion compilerVer sources

        cmd_ ["mkdir", "-p", dropDirectory1 binTarget]
        path <- fromMaybe "" <$> getEnv "PATH"
        let toLibs = fmap ("-l" <>)
        let libs' = ("atslib" :) $ bool libs ("gc" : libs) gc
        command
            [EchoStderr False, AddEnv "PATSHOME" home, AddEnv "PATH" (home ++ "/bin:" ++ path), AddEnv "PATSHOMELOCS" $ patsHomeLocs 5]
            (home ++ "/bin/patscc")
            (mconcat
                [ [src, "-atsccomp", cc ++ " -I" ++ h ++ "/ccomp/runtime/ -I" ++ h ++ " -L" ++ h' ++ "/lib" ++ " -L" ++ home ++ "/ccomp/atslib/lib", "-o", binTarget, "-cleanaft"]
                , makeCFlags cFlags hsLibs gc
                , toLibs libs'
                ])

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] -> [FilePath] -> Action [FilePath]
transitiveDeps _ [] = pure []
transitiveDeps gen ps = fmap join $ forM ps $ \p -> if p `elem` gen then pure mempty else do
    contents <- liftIO $ readFile p
    let ats = fromRight mempty . parseATS . lexATS $ contents
    let dir = takeDirectory p
    deps <- filterM (\f -> ((f `elem` gen) ||) <$> doesFileExist f) $ fixDir dir . trim <$> getDependencies ats
    deps' <- transitiveDeps gen 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 ()