module Utils.TH where
-- Common Utility Functions that require TemplateHaskell

-- import Data.Char

import ClassyPrelude.Yesod
import Language.Haskell.TH
import Language.Haskell.TH.Datatype
-- import Control.Monad
-- import Control.Monad.Trans.Class
-- import Control.Monad.Trans.Maybe
-- import Control.Monad.Trans.Except

import Data.List ((!!), foldl)

------------
-- Tuples --
------------

-- Alternatively uses lenses: "^. _3" projects the 3rd component of an n-tuple for any n >=3, requires import Control.Lens
projNI :: Int -> Int -> ExpQ -- generic projection gives I-th element of N-tuple, i.e. snd3 = $(projNI 3 2)  --ghci -fth
-- $(projN n m) :: (t1,..,tn) -> tm (for m<=n)
projNI n i = do
  x <- newName "x"
  let rhs = varE x
  let pat = tupP $ replicate (pred i) wildP ++ varP x : replicate (n-i) wildP
  lamE [pat] rhs


-- | Generic projections N-tuples that are actually left-associative pairs
--   i.e. @$(leftAssociativePairProjection c n m :: (..(t1 `c` t2) `c` .. `c` tn) -> tm@ (for m<=n)
leftAssociativePairProjection :: Name -> Int -> Int -> ExpQ
leftAssociativePairProjection constructor n i = do
    x <- newName "x"
    lamE [pat x n] (varE x)
  where
    pat x 1       = varP x
    pat x w
      | w==i      = conP constructor [wildP, varP x]
      | otherwise = conP constructor [pat x (pred w), wildP]


---------------
-- Functions --
---------------

permuteFun :: [Int] -> ExpQ -- generic permutation of function arguments, i.e. $(permuteFun [2,1]) == flip
permuteFun perm = lamE pat rhs
  where pat = map varP $ fn:xs
        rhs = foldl appE (varE fn) $ map varE ps
--         rhs = appE (varE fn) (varE $ xs!!1)
        ln  = length perm
        xs  = [ mkName $ "x" ++ show j | j <- [1..ln] ]
        ps  = [ xs !! (j-1) | j <- perm ]
        fn  = mkName "fn"

altFun :: [Int] -> ExpQ -- generic permutation/repetition of function arguments, i.e. $(permuteFun [2,1]) == flip
altFun perm = lamE pat rhs
  where pat = map varP $ fn:xs
        rhs = foldl appE (varE fn) $ map varE ps
--         rhs = appE (varE fn) (varE $ xs!!1)
        mx  = maximum $ impureNonNull perm
        xs  = [ mkName $ "x" ++ show j | j <- [1..mx] ]
        ps  = [ xs !! (j-1) | j <- perm ]
        fn  = mkName "fn"

-- |
curryN :: Int -> ExpQ
curryN n = do
  fn <- newName "foo"
  xs <- replicateM n $ newName "x"
  let pat = map VarP (fn:xs)
  let tup = TupE (map VarE xs)
  let rhs = AppE (VarE fn) tup
  return $ LamE pat rhs

uncurryN :: Int -> ExpQ
uncurryN n = do
  fn <- newName "foo"
  xs <- replicateM n $ newName "x"
  let pat = [VarP fn, TupP (map VarP xs)]
  let rhs = foldl AppE (VarE fn) (map VarE xs)
  return $ LamE pat rhs


-- Special Show-Instances for Themes
deriveShowWith :: (String -> String) -> Name -> Q [Dec]
deriveShowWith = deriveSimpleWith ''Show 'show

-- deriveDisplayWith :: (String -> String) -> Name -> Q [Dec]
-- deriveDisplayWith = deriveSimpleWith ''DisplayAble 'display

deriveSimpleWith ::  Name -> Name -> (String -> String) -> Name -> Q [Dec]
deriveSimpleWith cls fun strOp ty = do
    (TyConI tyCon)  <- reify ty
    (tyConName, cs) <- case tyCon of
        DataD [] nm [] _ cs _ -> return (nm, cs)
        _ -> fail "deriveShowTheme: tyCon must be a plain datatype enumeration"
    let instanceT = conT cls `appT` conT tyConName
    return <$> instanceD (return []) instanceT [genDecs cs]
  where
    genDecs :: [Con] -> Q Dec
    genDecs cs = funD fun (map genClause cs)

    genClause :: Con -> Q Clause
    genClause (NormalC name []) =
      let pats = [ConP name []]
          body = NormalB $ LitE $ StringL $ strOp $ nameBase name
      in return $ Clause pats body []
    genClause _ = fail "deriveShowTheme: constructors not allowed to have arguments"

embedRenderMessage :: Name -- ^ Foundation type
                   -> Name -- ^ Type to embed into message type
                   -> (Text -> Text) -- ^ Mangle constructor names
                   -> DecsQ
-- ^ @embedRenderMessage ''Foundation ''MessageType mangle@ declares a
--   `RenderMessage Foundation MessageType` instance expecting the default
--   message-datatype (@FoundationMessage@) to contain one constructor for each
--   constructor of @MessageType@, taking the same arguments:
--
-- > data NewMessage = NewMessageOne | NewMessageTwo
-- > data FoundationMessage = MsgOne | MsgTwo
-- >
-- > -- embedRenderMessage ''Foundation ''NewMessage (drop 2 . splitCamel)
-- > instance RenderMessage Foundation NewMessage where
-- >   renderMessage f ls = renderMessage f ls . \case
-- >     NewMessageOne -> MsgOne
-- >     NewMessageTwo -> MsgTwo
embedRenderMessage f inner mangle = do
  DatatypeInfo{..} <- reifyDatatype inner
  let
    matches :: [MatchQ]
    matches = flip map datatypeCons $ \ConstructorInfo{..} -> do
      vars <- forM constructorFields $ \_ -> newName "x"
      let body = foldl (\e v -> e `appE` varE v) (conE . mkName . unpack $ "Msg" <> mangle (pack $ nameBase constructorName)) vars
      match (conP constructorName $ map varP vars) (normalB body) []

  f' <- newName "f"
  ls <- newName "ls"

  pure <$> instanceD (cxt []) [t|RenderMessage $(conT f) $(conT inner)|]
    [ funD 'renderMessage
      [ clause [varP f', varP ls] (normalB $ [e|renderMessage $(varE f') $(varE ls) . $(lamCaseE matches)|]) []
      ]
    ]

embedRenderMessageVariant :: Name -- ^ Foundation Type
                          -> Name -- ^ Name of newtype
                          -> (Text -> Text) -- ^ Mangle constructor names
                          -> DecsQ
embedRenderMessageVariant f newT mangle = do
  [ConstructorInfo{ constructorName = newtypeName, constructorFields = [ ConT newtypeInner ] }] <- datatypeCons <$> reifyDatatype newT
  DatatypeInfo{..} <- reifyDatatype newtypeInner

  let
    matches :: [MatchQ]
    matches = flip map datatypeCons $ \ConstructorInfo{..} -> do
      vars <- forM constructorFields $ \_ -> newName "x"
      let body = foldl (\e v -> e `appE` varE v) (conE . mkName . unpack $ "Msg" <> mangle (pack $ nameBase constructorName)) vars
      match (conP newtypeName [conP constructorName $ map varP vars]) (normalB body) []

  f' <- newName "f"
  ls <- newName "ls"

  pure <$> instanceD (cxt []) [t|RenderMessage $(conT f) $(conT newT)|]
    [ funD 'renderMessage
      [ clause [varP f', varP ls] (normalB $ [e|renderMessage $(varE f') $(varE ls) . $(lamCaseE matches)|]) []
      ]
    ]


dispatchTH :: Name -- ^ Datatype to pattern match
           -> ExpQ
-- ^ Produces a lambda-case-expression matching all constructors of the named datatype and calling a function (named after the constructor prefixed with @dispatch@) on the fields of each constructor
dispatchTH dType = do
  DatatypeInfo{..} <- reifyDatatype dType
  let
    matches = map mkMatch datatypeCons
    mkMatch ConstructorInfo{..} = do
      pats <- forM constructorFields $ \_ -> newName "x"
      let fName = mkName $ "dispatch" <> nameBase constructorName
      match (conP constructorName $ map varP pats) (normalB $ foldl (\e pat -> e `appE` varE pat) (varE fName) pats) []
  lamCaseE matches