-- The JSDict data type and JSDictM monad

{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}

module JSDict
    ( JSDictM
    , shareFun
    , buildFun2
    , buildStmt2
    , buildStmt3
    , newSM
    , expStmt
    , Def (..)
    , runJSDictM
    ) where

import Param
import JS
import Memo

import Language.HJavaScript.Syntax

import Control.Monad.Identity
import Control.Monad.Trans
import Control.Monad.Trans.Writer
import Control.Monad.Trans.State

---------------------------------------------------------------


data Def where
    SMDef :: String -> {-n :-}Int -> ([String]{-Vec n String-} -> Memo String) -> Def
    JSDef :: String -> Def

type JSDictM = WriterT [Def] (StateT Int Identity)

runJSDictM :: JSDictM a -> ((a, [Def]), Int)
runJSDictM = runIdentity . flip runStateT (1 :: Int) . runWriterT

data JSM x r where
    Ret :: r -> JSM () r
    Bind :: JSM () k -> (k -> JSM x r) -> JSM x r
    VDA :: Exp t -> JSM () (Exp t)
    ES :: Exp t -> JSM () ()

type JSMB r = JSM r ()

instance Monad (JSM ()) where
    return = Ret
    (>>=) = Bind

expStmt :: Exp t -> JSM () ()
expStmt = ES

runJSM :: JSM r () -> JSDictM (Block r)
runJSM = fmap fst . runJSM'

infixl 1 |+|

(|+|) :: forall t. Block () -> Stmt t -> Block t
(|+|) = Sequence


runJSM' :: JSM r a -> JSDictM (Block r, a)
runJSM' (Ret r) = return (EmptyBlock, r)
runJSM' (VDA e) = do
    i <- newId
    let v = 'd': i
    return (EmptyBlock |+| VarDeclAssign v e, JConst v)
runJSM' (ES e) =
    return (EmptyBlock |+| ExpStmt e, ())
runJSM' (Bind m f) = do
    (a, x) <- runJSM' m
    (b, y) <- runJSM' $ f x
    return (a .+ b, y)

(.+) :: Block () -> Block a -> Block a
a .+ EmptyBlock = a
a .+ Sequence b s = Sequence (a .+ b) s


newId :: JSDictM String
newId = do
    i <- lift get
    lift $ put $ i + 1
    return $ show i

newVar :: JSDictM (Var a)
newVar = fmap (JParam . ('v':)) newId

newFunName :: JSDictM String
newFunName = fmap ('f':) newId


tellJS :: (Show a) => a -> JSDictM ()
tellJS f = tell [JSDef $ show f]

tellSM :: String -> Int -> ([String] -> Memo String) -> JSDictM ()
tellSM name i f = tell [SMDef name i f]

getPre :: Exp String -> Exp (String -> ()) -> Stmt ()
getPre = curry $
        ExpStmt .
--        JI .
        JCall (JConst "getPre")

infixr 5 `strPlus`

strPlus :: Exp String -> Exp String -> Exp String
strPlus x y = JBinOp x Plus y

(|=|) :: String -> Exp String -> Exp String
t |=| x = JString (t ++ "=") `strPlus` encodeURIComponent x

encodeURIComponent :: Exp String -> Exp String
encodeURIComponent = JCall $ JConst "encodeURIComponent"

infixr 3 |&|

(|&|) :: Exp String -> Exp String -> Exp String
a |&| b = a `strPlus` JString "&" `strPlus` b


newSM
    :: forall t1 t2 t3 a
    . (FromString t1, FromString t2, ToString t3)
    => (t1 -> t2 -> Memo t3)
    -> JSDictM ( Exp t1
              -> Exp t2
              -> (Exp t3 -> Exp a)
              -> Exp ())
newSM body = do
    name <- newFunName
    (w1 :: Var String) <- newVar
    (w2 :: Var String) <- newVar
    (w3 :: Var (String -> ())) <- newVar
    (v1 :: Var t3) <- newVar
    tellJS $ JFunction (Just name) (w1, w2, w3)
                  $ EmptyBlock
                |+| getPre ( "sm" |=| JString name
                         |&| "v1" |=| val w1
                         |&| "v2" |=| val w2)
                    (val w3)

    tellSM name 2 $ \[x, y] -> fmap ts $ body (fs x) (fs y)
    return $ \w1 w2 w3f
        -> JCall (JConst name)
            (w1, w2, JFunction Nothing v1
                        $ EmptyBlock
                      |+| ExpStmt (w3f $ val v1))



buildFun :: forall t r . ParamType t => (Exp t -> Block r) -> JSDictM (Exp t -> Exp r)
buildFun body = do
    name <- newFunName
    v <- newVar
    tellJS $ JFunction (Just name) v $ body $ val v
    return $ JCall $ JConst name

shareFun :: forall t r . ParamType t => (Exp t -> Exp r) -> JSDictM (Exp t -> Exp r)
shareFun f = buildFun $ Sequence EmptyBlock . Return . f

buildFun2 :: forall t1 t2 r . ParamType (t1,t2)
    => (Exp t1 -> Exp t2 -> JSMB r)
    -> JSDictM (Exp t1 -> Exp t2 -> Exp r)
buildFun2 body = do
    name <- newFunName
    v1 <- newVar
    v2 <- newVar
    b <- runJSM $ body (val v1) (val v2)
    tellJS $ JFunction (Just name) (v1,v2) b

    return $ \v1 v2 -> JCall (JConst name) (v1,v2)

buildStmt2 :: forall t1 t2 r . ParamType (t1,t2)
    => (Exp t1 -> Exp t2 -> JSMB r)
    -> JSDictM (Exp t1 -> Exp t2 -> Stmt ())
buildStmt2 f = fmap (\f a b -> ExpStmt $ f a b) $ buildFun2 f

buildFun3 :: forall t1 t2 t3 r . ParamType (t1,t2,t3)
    => (Exp t1 -> Exp t2 -> Exp t3 -> JSMB r)
    -> JSDictM (Exp t1 -> Exp t2 -> Exp t3 -> Exp r)
buildFun3 body = do
    name <- newFunName
    v1 <- newVar
    v2 <- newVar
    v3 <- newVar
    b <- runJSM $ body (val v1) (val v2) (val v3)
    tellJS $ JFunction (Just name) (v1,v2,v3) b
    return $ \v1 v2 v3 -> JCall (JConst name) (v1,v2,v3)

buildStmt3 :: forall t1 t2 t3 r . ParamType (t1,t2,t3)
    => (Exp t1 -> Exp t2 -> Exp t3 -> JSMB ())
    -> JSDictM (Exp t1 -> Exp t2 -> Exp t3 -> Stmt ())
buildStmt3 f = fmap (\f t1 t2 t3 -> ExpStmt $ f t1 t2 t3) $ buildFun3 f