{-# LANGUAGE TemplateHaskell, NoCPP #-}
module RlangQQ.Internal where

import System.IO.Temp
import System.IO

import Control.Applicative
import Data.Char
import Data.List
import Data.Maybe
import Data.Monoid
import Data.Tree
import Language.Haskell.TH
import Language.Haskell.TH.Quote
import qualified Data.Foldable as F
import System.Directory
import System.Process
import Text.Printf
import Text.Trifecta
import Text.Read (readMaybe)

import qualified Data.ByteString.Lazy.UTF8 as B
import qualified Data.ByteString.Lazy as B
import qualified Data.ByteString.Lazy
import qualified Data.ByteString.Lazy.Char8 as B8

import qualified Data.Digest.Pure.SHA as SHA

import Control.Monad.State
import Control.Monad.Trans.Maybe
import Data.Generics
import qualified Data.Traversable as T
import qualified Data.Map as M
import Data.Foldable (foldMap)
import qualified Data.Set as S

import System.FilePath

import Paths_Rlang_QQ

import RlangQQ.Binary
import RlangQQ.Antiquote
-- import RlangQQ.ParseKnitted

import Control.Concurrent.Chan
import Control.Concurrent.MVar
import Control.Concurrent

import Control.Lens hiding (noneOf)

import System.IO.Unsafe
import Data.IORef

-- * the main quoter

-- | going via binary serialization (classes from "RlangQQ.Binary"). This is used in 'RlangQQ.r'
quoteRExpression2 i returnChan str0 = do
    let str0Hash = take 8 $ SHA.showDigest (SHA.sha1 (B8.pack str0))
        rawFile, mdFile, rmdFile, inputFile, inputFile2, outputFile,
          figPath :: String
        rawFile    = printf "Rtmp/%s/raw%d.R" str0Hash (i :: Int)
        rmdFile    = printf "Rtmp/%s/raw%d.Rmd" str0Hash i
        rmdFile2   = printf "Rtmp/%s/raw%d_.Rmd" str0Hash i
        mdFile     = printf "Rtmp/%s/raw%d.md" str0Hash i
        figPath    = printf "Rtmp/%s/fig%d/" str0Hash i
        inputFile  = printf "Rtmp/%s/%d_hs.rdx" str0Hash i
        inputFile2 = printf "Rtmp/%s/%d_hs2.rdx" str0Hash i
        outputFile = printf "Rtmp/%s/%d_r.rdx" str0Hash i

        (knitrHdr, withRawFile -> (("setwd('../../'); getwd();"++) -> str, addAntiquotes))
            = splitKnitrHdr str0
    let inRmdChunk s = unlines ["```{r " ++ knitrHdr ++ "}", s, "```"]

    variables <- runIO $ do
        t <- do
          t1 <- getDataFileName "Tree.R"
          b <- doesFileExist t1
          -- for the benefit of "cabal test"
          -- which does not set the environment variable
          -- Rlang_QQ_datadir=rsrc, and you don't necessarily
          -- have Tree.R installed in the ~/.cabal/... yet
          return $ if b then t1 else "rsrc/Tree.R"

        createDirectoryIfMissing True ("Rtmp/" ++ str0Hash)
        writeFile rawFile str
        writeFile rmdFile (inRmdChunk str)
        rt <- readProcess "R" ["--no-save","--quiet"]
            $ printf "source('%s'); rlangQQ_toTree(parse('%s'))" t rawFile
        case parseString parseTree mempty rt of
            Success a -> return $ hsClassify a
            Failure a -> do
                print a
                error "parse failure"

    (variables, chanVars) <- return $ M.partitionWithKey (\k _ -> "hs_" `isPrefixOf` k) variables

    let
        writeInputFile' :: FilePath -> M.Map String Intent -> ExpQ
        writeInputFile' inputFile vars = [|
            B.writeFile inputFile $
                $(appE [| toRDA . Record |] $
                  mkHList
                   [ [| $(label x) .=. ($(label' x) .=. $(var x)) |]
                    | (x, intent) <- M.toList vars, notOut intent ])
                |]

        writeChanInput :: ExpQ
        writeChanInput
         | M.size chanVars == 0 = [| return () |]
         | otherwise = do
            (withOutputs, binds) <- fmap unzip $ sequence [ do
                withOutput <- newName "writer"
                return (withOutput, bindS (tupP [varP (mkName (dropHS x)),
                                            varP withOutput]) [| readChan $(var x) |] )
                    | (x, intent) <- M.toList chanVars, notOut intent ]
            let write = doE $ binds ++ [  noBindS (writeInputFile' inputFile2 chanVars),
                        noBindS [|
                           return $ \outputRecord -> $(doE
                            [ noBindS [| $(varE oi)
                                            (outputRecord `asTypeOf` Record $sampOutput2) :: IO () |]
                                    | oi <- withOutputs ] )
                           |] ]
            write

        writeInputFile :: ExpQ
        writeInputFile = writeInputFile' inputFile variables

        -- a record containing variables which were sent
        -- into R (intent InOut), or undefined for intent Out
        -- which will aid in type inference.
        sampOutput' :: (String -> ExpQ) -> ExpQ
        sampOutput' addLabels = mkHList
                [ [| $(addLabels x) $(case intent of
                                    Out -> [| error ("RlangQQ.Internal.sampOutput" ++
                                                     " should not be evaluated") |]
                                    _ -> var x ) |]

                    | (x, intent) <- M.toList variables, notIn intent ]

        sampOutput, sampOutput2 :: ExpQ
        sampOutput  = sampOutput' $ \x -> [| \ y -> $(label x) .=.  ($(label' x) .=. y) |]
        sampOutput2 = sampOutput' $ \x -> [| \ y -> $(label x) .=.  (                y) |]


        readOutputFile :: ExpQ
        readOutputFile = [|
            let -- Record '[Tagged "x" (Tagged "x1" a)] -> Record '[Tagged "x" a]
                dropInnerTagged x = over unlabeled (^. to Record . unlabeled) x
                fixTy x = x `asTypeOf` Record $sampOutput
            in do
              cts <- B.readFile outputFile
              -- lazy IO makes for confusings: possibly better to switch
              -- over to strict BS instead of this
              Data.ByteString.Lazy.length cts `seq` return ()

              -- to stop ghc-7.6 from having a type error:
              return . $(varE 'dropInnerTagged) . fixTy $ fromRDA cts |]

        outVars :: [String]
        outVars = [ s  | (s, intent) <- M.toList variables
                        ++ M.toList chanVars {- necessary? -}, notIn intent ]

        saveRVars :: String
        saveRVars
            | [] <- outVars = ""
            | otherwise = printf "save(%s, file='%s', compress='gzip');"
                                    (intercalate "," $ reverse outVars)
                                    (takeFileName outputFile)
        rmd2Content :: String
        rmd2Content = inRmdChunk str ++ inRmdChunk saveRVars

        runRNoChan :: ExpQ
        runRNoChan  = do
            [| readProcess "R" ["--no-save", "--quiet"] $ concat
                    [ $(stringE (printf "library(knitr);\
                                       \opts_chunk$set(fig.path='%s',tidy=F);\
                                       \load('%s');" figPath inputFile)),
                      $(stringE (printf "knit('%s', output='%s', quiet=TRUE);" rmdFile2 mdFile))]
                    |]


        whenOutVars :: ExpQ -> ExpQ
        whenOutVars e = unlessQ (null outVars) e

        runRChan :: ExpQ
        runRChan = do
          [| do
            chOut <- $(unlessQ (null outVars && not returnChan) [| newChan `asTypeOf`
                        ((undefined :: HList a -> IO (Chan (Record a))) $sampOutput2) |])

            outVar <- $(whenOutVars [| newEmptyMVar |])

            forkIO $ void $ do
                -- binary/text IO instead?
                (i,o,err,pid) <- runInteractiveProcess "R" ["--no-save", "--quiet"] Nothing Nothing
                forkIO $ putStr =<< hGetContents err
                hPutStrLn i $ printf "library(knitr);\
                                     \opts_chunk$set(fig.path='%s');\
                                     \load('%s');" figPath inputFile
                hFlush i
                let uniqueDoneString = "\"done calculating signaled by a very unique\
                                        \ string that will never happen by chance\""
                forkIO $ forever $ do
                    withOutputFn <- $writeChanInput
                    hPutStrLn i $ printf "load('%s'); knit('%s', output='%s', quiet=TRUE); %s"
                                        inputFile2 rmdFile2 mdFile
                                        uniqueDoneString
                    hFlush i
                    $(whenOutVars [|  putMVar outVar withOutputFn |] )

                forkIO $ forever $ do
                    oEOF <- hIsEOF o -- is this ok?
                    when oEOF (killThread =<< myThreadId)
                    o <- hGetLine o
                    when (o == ("[1] "++uniqueDoneString))
                        $(whenOutVars [| do
                                ov <- takeMVar outVar
                                output <- $readOutputFile

                                ov output
                                $(whenQ returnChan
                                    [| $(varE 'writeChan) chOut output |]
                                    -- rewrite to: [| writeChan chOut output |]
                                    -- when ghc-7.6 support is dropped
                                    )
                                |])
            return chOut
           |]

    runIO $ writeFile rmdFile2 rmd2Content

    let positionalVariablesIx :: M.Map Int (S.Set String)
        positionalVariablesIx = M.fromListWith (<>)
                        $ mapMaybe (\ k -> do
                              n <- readMaybe =<< msum (map (`stripPrefix` k) prefixes)
                              return (n, S.singleton k))
                        $ [ v  | (v, intent) <- M.toList variables,
                                  intent `elem` [In,InOut] ]

    addLambda positionalVariablesIx $ addAntiquotes [| do
        {- already done, but Rtmp could have moved
           since compile time, or another RlangQQ could
           have overwritten the file -}
        createDirectoryIfMissing True ("Rtmp" </> str0Hash)
        writeFile rmdFile2 rmd2Content

        $writeInputFile
        $(if M.null chanVars then [| do
                $runRNoChan
                $( whenOutVars readOutputFile) |]
            else runRChan)

        |]

unlessQ, whenQ :: Bool -> ExpQ -> ExpQ
unlessQ b e | b = [| return () |]  | otherwise = e
whenQ b e = unlessQ (not b) e


splitKnitrHdr :: String -> (String,String)
splitKnitrHdr ('{' : xs) = case f [] xs of (acc,ys) -> (reverse acc, ys)
    where
    f acc (stripPrefix "\\}" -> Just xs) = f ('}':acc) xs
    f acc ('}':xs) = (acc, xs)
splitKnitrHdr xs = ([], xs)

withRawFile :: String -> (String, ExpQ -> ExpQ)
withRawFile str =
    case parseString extractAntiquotes mempty str of
        Failure msg -> (str, \xp -> do
                    reportWarning (show msg)
                    xp)
        Success parsed ->
            (\(a,b,c) -> (a,b)) $
            foldr (\ chunk ~( str, ef, i) ->
                case chunk of
                    Left (LitE (IntegerL p)) -> ("hs_" ++ show p ++ " " ++ str,
                                                  ef, i)
                    Left x ->
                        let v = "hs_interp" ++  show i in
                        ( concat [v, " ", str],
                                    \e0 ->  caseE (return x)
                                        [ match
                                                (mkName (dropHS v) `asP` varP (mkName v))
                                                (normalB (ef e0))
                                                [] ],
                          i+1)
                    Right s -> (s ++ str, ef, i))
                ("", id, 1)
                parsed


addLambda :: M.Map Int (S.Set String) -> ExpQ -> ExpQ
addLambda ps xp
    | Just ((pMax,_), _ ) <- M.maxViewWithKey ps = do
        let allPS = M.fromList [ (n,()) | n <- [1 .. pMax] ]
        case M.keys $ ps M.\\ allPS of
            [] -> return ()
            extra -> reportWarning ("Unexpected positional arguments " ++ show extra)

        case M.keys $ allPS M.\\ ps of
            [] -> return ()
            unused -> reportWarning ("Positional argument not used " ++ show unused)

        case  [ p0 |  p0 <- M.elems ps,
                      S.size p0 > 1 ] of
            [] -> return ()
            duplicates -> reportError ("Positional argument types must be the same " ++ show duplicates )

        lamE [ case M.lookup n ps of
                   Just (S.toList -> [e]) -> varP (mkName e)
                   Nothing -> wildP
                | n <- [1 .. pMax] ] xp
    | otherwise = xp



-- ** utility functions
-- | go from the variable name used on the R-side to the one in the haskell side
dropHS x = fromMaybe x (foldMap stripPrefix prefixes x)


prefixes = ["hs_","ch_"]

-- | HList label
label x = [| Label :: Label $(litT (strTyLit (dropHS x))) |]
label' x = [| Label :: Label $(litT (strTyLit x)) |]

var x = varE $ mkName $ case dropHS x of
          x' | all isDigit x' -> x
             | otherwise -> x'

mkHList :: [ExpQ] -> ExpQ
mkHList = foldl (\ b a -> [| $a .*. $b |]) [| HNil |]

notOut Out = False
notOut _ = True

notIn In = False
notIn _ = True






-- * converting R\'s AST

parseTree' =  do
    symbol "Node"
    nl <- stringLiteral'
    children <- brackets $ commaSep parseTree'
    return (Node nl children)

parseTree = do
    manyTill anyChar (try (string "[1]"))
    manyTill (noneOf "\n") space
    between (oneOf "\"") (oneOf "\"") parseTree'


-- | gets variables like @abc@ provided the R file contained @hs_abc@
hsVars :: Tree String -> [String]
hsVars =  mapMaybe (foldMap stripPrefix prefixes) . F.toList



labelTree :: Tree String -> Tree (String, Int)
labelTree t = flip evalState 0 $ T.forM t $ \x -> do
                n <- get
                put (n+1)
                return (x, n)

-- | what is the usage of a @hs_@ variable on the R side?
data Intent = In -- ^ only read
    | Out -- ^ only assigned to
    | InOut -- ^ both
    deriving (Show, Eq)



classifyExp :: Tree String -> Maybe (String, Intent)
classifyExp (Node oper (a: _))
    | oper `elem` ["<-", "="],
      [r] <- filter (\x -> any (`isPrefixOf` x) prefixes) $ leftmosts a = Just (r, Out)

      where
    leftmosts (Node a []) = [a]
    leftmosts (Node n (a:_)) = n : leftmosts a

classifyExp (Node n _) | any (`isPrefixOf` n) prefixes = Just (n, In)
classifyExp _ = Nothing


hsClassify :: Tree String -> M.Map String Intent
hsClassify
    = M.fromListWith (flip merge)
    . mapMaybe classifyExp
    . toList

    where
    merge Out _ = Out
    merge In Out = InOut
    merge InOut _ = InOut
    merge _ InOut = InOut
    merge In In = In

    -- all Node being viewed
    toList :: Tree a -> [Tree a]
    toList a = a : concatMap toList (subForest a)



-- ** tests

parseTreeTest2 = parseTreeTest =<< getDataFileName "parseTreeExample.R"

parseTreeTest3 contents = withSystemTempFile "RlangQQ.tmp" $ \fn h -> do
    System.IO.hPutStrLn h contents
    hFlush h
    parseTreeTest fn

parseTreeTest inputFile = do
    t <- getDataFileName "Tree.R"
    rt <- readProcess "R" ["--no-save","--quiet"] $ "source('"++ t ++ "'); rlangQQ_toTree(parse('"++ inputFile ++ "'))"
    print rt
    let r = parseString parseTree mempty rt :: Result (Tree String)

    print (fmap hsClassify r)
    -- print (fmap (hsAssignedFirstVars True) r)
    -- print (fmap (hsAssignedFirstVars False) r)
    -- print (zipper r & withins root)




-- * global variable quote index

getRlangQQ_n :: Q Int
getRlangQQ_n = runIO $ do
    n <- readIORef rlangQQ_n
    writeIORef rlangQQ_n (n+1)
    return n

{-# NOINLINE rlangQQ_n #-}
rlangQQ_n :: IORef Int
rlangQQ_n = unsafePerformIO (newIORef 1)