{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE CPP #-}
module Language.Haskell.Generate.Monad
  ( Generate(..), ExpG
  , runGenerate, newName
  , returnE
  , useValue, useCon, useVar
  , caseE
  , applyE, applyE2, applyE3, applyE4, applyE5, applyE6
  , (<>$)
  , GenExp(..)
  , ModuleM(..)
  , ModuleG
  , FunRef(..)
  , Name(..)
  , exportFun
  , addDecl
  , runModuleM
  , generateModule
  , generateExp
  )
  where

import           Control.Applicative
import           Control.Monad
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.State
import           Control.Monad.Trans.Writer
import qualified Data.Set as S
import           Language.Haskell.Exts.Pretty
import           Language.Haskell.Exts.SrcLoc
import           Language.Haskell.Exts.Syntax
import           Language.Haskell.Generate.Expression

--------------------------------------------------------------------------------
-- Generate expressions

-- | This monad keeps track of a counter for generating unique names and the set of modules
-- that are needed for the expression.
newtype Generate a = Generate { unGenerate :: StateT Integer (Writer (S.Set ModuleName)) a } deriving (Functor, Applicative, Monad)

-- | Extract the set of modules and the value from a Generate action.
runGenerate :: Generate a -> (a, S.Set ModuleName)
runGenerate (Generate a) = runWriter $ evalStateT a 0 

-- | This is a type alias for a Generate action that returns an expression of type 't'.
type ExpG t = Generate (Expression t)

-- | Use a haskell-src-exts Exp as the result of a Generate action.
returnE :: Exp -> ExpG t
returnE = return . Expression

-- | Pretty print the expression generated by a given action.
generateExp :: ExpG t -> String
generateExp = prettyPrint . runExpression . fst . runGenerate

-- | Generate a case expression.
caseE :: ExpG x -> [(Pat, ExpG t)] -> ExpG t
caseE v alt = do
  v' <- v
#if MIN_VERSION_haskell_src_exts(1,16,0)   
  alt' <- mapM (\(p,a) -> fmap (\a' -> Alt noLoc p (UnGuardedRhs $ runExpression a') (BDecls [])) a) alt
#else  
  alt' <- mapM (\(p,a) -> fmap (\a' -> Alt noLoc p (UnGuardedAlt $ runExpression a') (BDecls [])) a) alt
#endif
  return $ Expression $ Case (runExpression v') alt'

-- | Import a function from a module. This function is polymorphic in the type of the resulting expression, 
-- you should probably only use this function to define type-restricted specializations. 
--
-- Example:
--
-- > addInt :: ExpG (Int -> Int -> Int) -- Here we restricted the type to something sensible
-- > addInt = useValue "Prelude" $ Symbol "+"
--
useValue :: String -> Name -> ExpG a
useValue md name = Generate $ do
  lift $ tell $ S.singleton $ ModuleName md
  return $ Expression $ Var $ Qual (ModuleName md) name

-- | Import a value constructor from a module. Returns the qualified name of the constructor.
useCon :: String -> Name -> Generate QName
useCon md name = Generate $ do
  lift $ tell $ S.singleton $ ModuleName md
  return $ Qual (ModuleName md) name

-- | Use the value of a variable with the given name.
useVar :: Name -> ExpG t
useVar name = return $ Expression $ Var $ UnQual name

-- | Generate a new unique variable name with the given prefix. Note that this new variable name
-- is only unique relative to other variable names generated by this function. 
newName :: String -> Generate Name
newName pref = Generate $ do
  i <- get <* modify succ
  return $ Ident $ pref ++ show i

-- | Generate a expression from a haskell value. This can for example be used to create lambdas:
-- 
-- >>> putStrLn $ generateExp $ expr (\x f -> f <>$ x)
-- \ pvar_0 -> \ pvar_1 -> pvar_1 pvar_0
--
-- Or string literals:
--
-- >>>  putStrLn $ generateExp $ expr "I'm a string!"
-- ['I', '\'', 'm', ' ', 'a', ' ', 's', 't', 'r', 'i', 'n', 'g', '!']
--
class GenExp t where
  type GenExpType t :: *

  -- | This function generates the haskell expression from the given haskell value.
  expr :: t -> ExpG (GenExpType t)

instance GenExp (ExpG a) where
  type GenExpType (ExpG a) = a
  expr = id

instance GenExp (Expression t) where
  type GenExpType (Expression t) = t
  expr = return

instance GenExp Char where
  type GenExpType Char = Char
  expr = return . Expression . Lit . Char

instance GenExp Integer where
  type GenExpType Integer = Integer
  expr = return . Expression . Lit . Int

instance GenExp Rational where
  type GenExpType Rational = Rational
  expr = return . Expression . Lit . Frac

instance GenExp a => GenExp [a] where
  type GenExpType [a] = [GenExpType a]
  expr = Generate . fmap (Expression . List . map runExpression) . mapM (unGenerate . expr)

instance GenExp x => GenExp (ExpG a -> x) where
  type GenExpType (ExpG a -> x) = a -> GenExpType x
  expr f = do 
    pvarName <- newName "pvar_"
    body <- expr $ f $ return $ Expression $ Var $ UnQual pvarName
    return $ Expression $ Lambda noLoc [PVar pvarName] $ runExpression body

--------------------------------------------------------------------------------
-- Apply functions

-- | Apply a function in a haskell expression to a value.
applyE :: ExpG (a -> b) -> ExpG a -> ExpG b
applyE a b = wrap $ liftM (foldl1 App) $ sequence [unwrap a, unwrap b]
  where wrap = fmap Expression
        unwrap = fmap runExpression

-- | Operator for 'applyE'. 
(<>$) :: ExpG (a -> b) -> ExpG a -> ExpG b
(<>$) = applyE

infixl 1 <>$

-- | ApplyE for 2 arguments
applyE2 :: ExpG (a -> b -> c) -> ExpG a -> ExpG b -> ExpG c
applyE2 a b c = wrap $ liftM (foldl1 App) $ sequence [unwrap a,unwrap b,unwrap c]
  where wrap = fmap Expression
        unwrap = fmap runExpression

-- | Apply a function to 3 arguments
applyE3 :: ExpG (a -> b -> c -> d) -> ExpG a -> ExpG b -> ExpG c -> ExpG d
applyE3 a b c d = wrap $ liftM (foldl1 App) $ sequence [unwrap a,unwrap b,unwrap c,unwrap d]
  where wrap = fmap Expression
        unwrap = fmap runExpression

-- | Apply a function to 4 arguments
applyE4 :: ExpG (a -> b -> c -> d -> e) -> ExpG a -> ExpG b -> ExpG c -> ExpG d -> ExpG e
applyE4 a b c d e = wrap $ liftM (foldl1 App) $ sequence [unwrap a,unwrap b,unwrap c,unwrap d,unwrap e]
  where wrap = fmap Expression
        unwrap = fmap runExpression

-- | Apply a function to 5 arguments
applyE5 :: ExpG (a -> b -> c -> d -> e -> f) -> ExpG a -> ExpG b -> ExpG c -> ExpG d -> ExpG e -> ExpG f
applyE5 a b c d e f = wrap $ liftM (foldl1 App) $ sequence [unwrap a,unwrap b,unwrap c,unwrap d,unwrap e,unwrap f]
  where wrap = fmap Expression
        unwrap = fmap runExpression

-- | Apply a function to 6 arguments
applyE6 :: ExpG (a -> b -> c -> d -> e -> f -> g) -> ExpG a -> ExpG b -> ExpG c -> ExpG d -> ExpG e -> ExpG f -> ExpG g
applyE6 a b c d e f g = wrap $ liftM (foldl1 App) $ sequence [unwrap a,unwrap b,unwrap c,unwrap d,unwrap e,unwrap f,unwrap g]
  where wrap = fmap Expression
        unwrap = fmap runExpression

--------------------------------------------------------------------------------
-- Generate modules

-- | A module keeps track of the needed imports, but also has a list of declarations in it.
newtype ModuleM a = ModuleM (Writer (S.Set ModuleName, [Decl]) a) deriving (Functor, Applicative, Monad)

-- | This is the resulting type of a function generating a module. It is a ModuleM action returning the export list.
type ModuleG = ModuleM (Maybe [ExportSpec])

-- | A reference to a function. With a reference to a function, you can apply it (by lifting it into ExprT using 'expr') to some value
-- or export it using 'exportFun'.
data FunRef t = FunRef Name

instance GenExp (FunRef t) where
  type GenExpType (FunRef t) = t
  expr (FunRef n) = return $ Expression $ Var $ UnQual n

-- | Generate a ExportSpec for a given function item.
exportFun :: FunRef t -> ExportSpec 
#if MIN_VERSION_haskell_src_exts(1,16,0) 
exportFun (FunRef name) = EVar NoNamespace (UnQual name)  
#else
exportFun (FunRef name) = EVar (UnQual name)
#endif

-- | Add a declaration to the module. Return a reference to it that can be used to either apply the function to some values or export it.
addDecl :: Name -> ExpG t -> ModuleM (FunRef t)
addDecl name e = ModuleM $ do
  let (body, mods) = runGenerate e
  tell (mods, [FunBind [Match noLoc name [] Nothing (UnGuardedRhs $ runExpression body) $ BDecls []]])
  return $ FunRef name

-- | Extract the Module from a module generator.
runModuleM :: ModuleG -> String -> Module
runModuleM (ModuleM act) name = 
#if MIN_VERSION_haskell_src_exts(1,16,0) 
  Module noLoc (ModuleName name) [] Nothing export (map (\md -> ImportDecl noLoc md True False False Nothing Nothing Nothing) $ S.toList imps) decls
#else
  Module noLoc (ModuleName name) [] Nothing export (map (\md -> ImportDecl noLoc md True False Nothing Nothing Nothing) $ S.toList imps) decls
#endif
  where (export, (imps, decls)) = runWriter act

-- | Generate the source code for a module.
generateModule :: ModuleG -> String -> String
generateModule = fmap prettyPrint . runModuleM