{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
module Language.Haskell.Generate.Base 
  ( ExpM(..), ExpG, ExpType
  , runExpM, newName
  , useValue, useCon, useVar
  , caseE
  , applyE, applyE2, applyE3, applyE4, applyE5, applyE6
  , (<>$)
  , GenExp(..)
  , ModuleM(..)
  , ModuleG
  , FunRef(..)
  , Name(..)
  , exportFun
  , addDecl
  , runModuleM
  , generateModule
  )
  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

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

-- | A ExpM is a monad used to track the imports that are needed for a given expression. Usually, you don't have to use
-- this type directly, but use combinators to combine several ExpM into bigger expressions. The t type parameter tracks
-- the type of the expression, so you don't accidently build expression that don't type check.
newtype ExpM t a = ExpM { unExpM :: StateT Integer (Writer (S.Set ModuleName)) a } deriving (Functor, Applicative, Monad)

-- | The ExpG type is a ExpM computation that returns an expression. Usually, this is the end result of a function generating 
-- a haskell expression
type ExpG t = ExpM t Exp

-- | Evaluate a ExpM action, returning the needed modules and the value.
runExpM :: ExpM t a -> (a, S.Set ModuleName)
runExpM (ExpM expt) = runWriter $ evalStateT expt 0

unsafeCoerceE :: ExpM t a -> ExpM t' a
unsafeCoerceE (ExpM x) = ExpM x

-- | Generate a case expression.
caseE :: ExpG x -> [(Pat, ExpG t)] -> ExpG t
caseE v alt = do
  v' <- unsafeCoerceE v
  alt' <- mapM (\(p,a) -> fmap (\a' -> Alt noLoc p (UnGuardedAlt a') (BDecls [])) a) alt
  return $ Case 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 = ExpM $ do
  lift $ tell $ S.singleton $ ModuleName md
  return $ Var $ Qual (ModuleName md) name

-- | Import a value constructor from a module. Returns the qualified name of the constructor.
useCon :: String -> Name -> ExpM t QName
useCon md name = ExpM $ 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 $ 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 -> ExpM t Name
newName pref = ExpM $ do
  i <- get <* modify succ
  return $ Ident $ pref ++ show i

-- | This type family can be used to get the type associated with some expression.
type family ExpType a :: *
type instance ExpType (ExpM t a) = t

-- | Generate a expression from a haskell value. 
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 Char where
  type GenExpType Char = Char
  expr = return . Lit . Char

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

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

instance GenExp a => GenExp [a] where
  type GenExpType [a] = [GenExpType a]
  expr = ExpM . fmap List . mapM (unExpM . expr)

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

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

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

-- | 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 = unsafeCoerceE $ liftM (foldl1 App) $ sequence [ce a,ce b,ce c]
  where ce = unsafeCoerceE

-- | 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 = unsafeCoerceE $ liftM (foldl1 App) $ sequence [ce a,ce b,ce c,ce d]
  where ce = unsafeCoerceE

-- | 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 = unsafeCoerceE $ liftM (foldl1 App) $ sequence [ce a,ce b,ce c,ce d,ce e]
  where ce = unsafeCoerceE

-- | 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 = unsafeCoerceE $ liftM (foldl1 App) $ sequence [ce a,ce b,ce c,ce d,ce e,ce f]
  where ce = unsafeCoerceE

-- | 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 = unsafeCoerceE $ liftM (foldl1 App) $ sequence [ce a,ce b,ce c,ce d,ce e,ce f,ce g]
  where ce = unsafeCoerceE

--------------------------------------------------------------------------------
-- 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 $ Var $ UnQual n

-- | Generate a ExportSpec for a given function item.
exportFun :: FunRef t -> ExportSpec 
exportFun (FunRef name) = EVar (UnQual name)

-- | 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) = runExpM e
  tell (mods, [FunBind [Match noLoc name [] Nothing (UnGuardedRhs body) $ BDecls []]])
  return $ FunRef name

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

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