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

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

import           Control.Monad
import           Control.Monad.IO.Class
import           Data.Binary                (Binary (..))
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          hiding (doesDirectoryExist)
import           Development.Shake.Cabal
import           Development.Shake.FilePath
import           GHC.Generics               (Generic)
import           Language.ATS
import           Language.ATS.Generate
import           System.Directory           (copyFile, createDirectoryIfMissing, doesDirectoryExist, listDirectory)
import           System.Exit                (ExitCode (ExitSuccess))

data ForeignCabal = ForeignCabal { cabalFile  :: TL.Text -- ^ @.cabal@ file associated with the library
                                 , objectFile :: TL.Text -- ^ Object file to be generated
                                 }

getSubdirs :: FilePath -> IO [FilePath]
getSubdirs p = do
    ds <- listDirectory p
    case ds of
        [] -> pure []
        xs -> do
            ds' <- filterM doesDirectoryExist (((p <> "/") <>) <$> xs)
            ss <- mapM getSubdirs ds'
            pure $ ds' <> join ss

newtype Version = Version [Integer]
    deriving (Eq, Generic, Binary)

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

genATS :: FilePath -> FilePath -> Rules ()
genATS src target =
    target %> \out -> liftIO $ do
        createDirectoryIfMissing True (takeDirectory out)
        genATSTypes src target

-- dumb heuristic:
-- go to dist-newstyle/build
-- find 'build' subdir
-- file is there

cabalExport :: ForeignCabal -> Rules ()
cabalExport (ForeignCabal cf' obf') = do

    let cf = TL.unpack cf'
        obf = TL.unpack obf'
        obfDir = takeDirectory (obf -<.> "hs")

    trDeps <- liftIO $ getCabalDeps cf
    obf %> \out -> do

        need (cf : fmap ((obfDir <> "/") <>) trDeps)
        command_ [Cwd obfDir] "cabal" ["new-build"]

        let subdir = takeDirectory cf ++ "/"
            endsBuild = (== "build") . last . splitPath
        dir <- filter endsBuild <$> liftIO (getSubdirs $ subdir ++ "dist-newstyle/build")
        let obj = head dir ++ "/" ++ (takeFileName obf)
        liftIO $ copyFile obj out

        let hdr = dropExtension obj ++ "_stub.h"
        liftIO $ copyFile hdr (takeDirectory out ++ "/" ++ hdr)

-- cabal deps
ghcExport :: String -> Rules ()
ghcExport m =
    [ m ++ ".o" ] &%> \[out] -> do
        let fn = out -<.> "hs"
        let dir = takeDirectory out
        need [ fn ]
        command [Cwd dir] "ghc" ["-c", "-O", takeFileName fn ]

-- | 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
       -> [ForeignCabal] -- ^ Cabal-based Haskell libraries
       -> [(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_ cabalExport hs

    out %> \_ -> do
        h <- patsHome v'
        h' <- pkgHome
        let home = h ++ "lib/ats2-postiats-" ++ show v
        sources <- transitiveDeps (snd <$> atg) [sourceFile]
        need (sources ++ (TL.unpack . objectFile <$> 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] -> [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 ()