{-# LANGUAGE CPP, LambdaCase, MultiWayIf #-}

module Language.HERMIT.PrettyPrinter.Clean
  ( -- * HERMIT's Clean Pretty-Printer for GHC Core
    corePrettyH
  )
where

import Control.Arrow hiding ((<+>))
import Control.Applicative ((<$>))

import Data.List (intersperse)
import Data.Char (isSpace)
import Data.Set (notMember)

import GhcPlugins (TyCon(..), Coercion(..), Var(..), Expr(..))
import qualified GhcPlugins as GHC

import Language.HERMIT.Context
import Language.HERMIT.Core
import Language.HERMIT.GHC
import Language.HERMIT.Kure
import Language.HERMIT.Monad
import Language.HERMIT.Syntax

import Language.HERMIT.PrettyPrinter.Common

import Language.HERMIT.Primitive.Common

import Pair

import Text.PrettyPrint.MarkedHughesPJ as PP

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

data RetExpr
        = RetLam PathH [DocH] PathH DocH
        | RetLet PathH [DocH] PathH DocH
        | RetApp DocH [RetExpr]
        | RetForAll PathH [DocH] PathH DocH
        | RetArrowType [DocH] -- TODO: store PathH here to mark arrows
        | RetExpr DocH
        | RetAtom DocH         -- parens not needed
        | RetEmpty

retApp :: RetExpr -> RetExpr -> RetExpr
retApp f             RetEmpty = f
retApp RetEmpty      e        = e
retApp (RetApp f es) e        = RetApp f (es ++ [e])
retApp f             e        = RetApp (parenExpr f) [e]

retApps :: RetExpr -> [RetExpr] -> RetExpr
retApps = foldl retApp

retLam :: PathH -> RetExpr -> RetExpr -> RetExpr
retLam _ RetEmpty e            = e
retLam p v (RetLam _ vs pb e)  = RetLam p (parenExpr v : vs) pb e
retLam p v e                   = RetLam p [parenExpr v] (p ++ [Lam_Body]) (normalExpr e)

retLet :: PathH -> RetExpr -> RetExpr -> RetExpr
retLet _ RetEmpty body                    = body
retLet p bnd      (RetLet _ bnds pb body) = RetLet p (normalExpr bnd : bnds) pb body
retLet p bnd      body                    = RetLet p [normalExpr bnd] (p ++ [Let_Body]) (normalExpr body)

retForAll :: PathH -> Crumb -> RetExpr -> RetExpr -> RetExpr
retForAll _ _ RetEmpty ty                      = ty
retForAll p _  v       (RetForAll _ vs pb ty)  = RetForAll p (parenExpr v : vs) pb ty
retForAll p cr v       ty                      = RetForAll p [parenExpr v] (p ++ [cr]) (normalExpr ty)

retArrowType :: RetExpr -> RetExpr -> RetExpr
retArrowType ty1 (RetArrowType tys) = RetArrowType (parenExprExceptApp ty1 : tys)
retArrowType ty1 ty2                = RetArrowType [parenExprExceptApp ty1, parenExprExceptApp ty2]

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

isEmptyR :: RetExpr -> Bool
isEmptyR RetEmpty = True
isEmptyR _        = False

isAtom :: RetExpr -> Bool
isAtom (RetAtom _) = True
isAtom _           = False

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

normalExpr :: RetExpr -> DocH
normalExpr  RetEmpty           = empty
normalExpr (RetAtom e)         = e
normalExpr (RetExpr e)         = e
normalExpr (RetLam p vs pb e)  = hang (attrP p (specialSymbol LambdaSymbol) <+> hsep vs <+> attrP pb (specialSymbol RightArrowSymbol)) 2 e
normalExpr (RetLet p vs pb e)  = sep [ attrP p (keyword "let") <+> vcat vs, attrP pb (keyword "in") <+> e ]
normalExpr (RetApp f es)       = let (atoms,exprs) = span isAtom es
                                  in sep [ hsep (f : map normalExpr atoms)
                                         , nest 2 (sep $ map parenExpr exprs) ]
normalExpr (RetForAll p vs pb ty) = attrP p (specialSymbol ForallSymbol) <+> hsep vs <+> attrP pb (symbol '.') <+> ty
normalExpr (RetArrowType tys)     = hsep (intersperse typeArrow tys)

parenExpr :: RetExpr -> DocH
parenExpr RetEmpty      = empty
parenExpr (RetAtom e)   = e
parenExpr (RetApp d []) = d
parenExpr other         = ppParens (normalExpr other)

parenExprExceptApp :: RetExpr -> DocH
parenExprExceptApp e@(RetApp _ _) = normalExpr e
parenExprExceptApp e              = parenExpr e

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

attrPAtomExpr :: PathH -> RetExpr -> RetExpr
attrPAtomExpr p (RetAtom d) = RetAtom (attrP p d)
attrPAtomExpr p (RetExpr d) = RetExpr (attrP p d)
attrPAtomExpr _ e           = e

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

ppParens :: DocH -> DocH
ppParens p = symbol '(' <> p <> symbol ')'

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

specialSymbol :: SpecialSymbol -> DocH
specialSymbol = markColor SyntaxColor . specialFont . char . renderSpecial

symbol :: Char -> DocH
symbol = markColor SyntaxColor . char

keyword :: String -> DocH
keyword = markColor KeywordColor . text


coChar :: Char -> DocH
coChar = coercionColor . char

coSymbol :: SpecialSymbol -> DocH
coSymbol = coercionColor . specialFont . char . renderSpecial

castSymbol :: DocH
castSymbol = coSymbol CastSymbol

coercionSymbol :: DocH
coercionSymbol = coSymbol CoercionSymbol

coercionBindSymbol :: DocH
coercionBindSymbol = coSymbol CoercionBindSymbol

coKeyword :: String -> DocH
coKeyword = coText -- An alternative would be keyword.


tySymbol :: SpecialSymbol -> DocH
tySymbol = typeColor . specialFont . char . renderSpecial

typeSymbol :: DocH
typeSymbol = tySymbol TypeSymbol

typeBindSymbol :: DocH
typeBindSymbol = tySymbol TypeBindSymbol

typeArrow :: DocH
typeArrow = tySymbol RightArrowSymbol

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

hasType :: DocH -> DocH -> DocH
e `hasType` ty = e <+> tySymbol TypeOfSymbol <+> ty

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

-- TODO: PrettyOptions should be in the context.

-- | Pretty print a fragment of GHC Core using HERMIT's \"Clean\" pretty printer.
corePrettyH :: PrettyOptions -> PrettyH CoreTC
corePrettyH opts =
    let
        -- Use for any GHC structure
        ppSDoc :: GHC.Outputable a => PrettyH a
        ppSDoc = do dynFlags <- constT GHC.getDynFlags
                    arr (toDoc . GHC.showSDoc dynFlags . GHC.ppr)
            where toDoc s | any isSpace s = parens (text s)
                          | otherwise     = text s

        ppVar :: Translate PrettyC HermitM Var RetExpr
        ppVar = readerT $ \ v -> GHC.varName ^>> ppName (varColor v)

        varColor :: Var -> SyntaxForColor
        varColor var | GHC.isTyVar var = TypeColor
                     | GHC.isCoVar var = CoercionColor
                     | otherwise       = IdColor

        ppName :: SyntaxForColor -> Translate PrettyC HermitM GHC.Name RetExpr
        ppName color = (GHC.nameOccName >>> GHC.occNameString) ^>> arr (\ name -> let doc  = markColor color (text name)
                                                                                           -- TODO: is "isScriptInfixId" the right predicate to use here?
                                                                                   in RetAtom $ if all isScriptInfixIdChar name
                                                                                                 then ppParens doc
                                                                                                 else doc
                                                                       )

        ppLitTy :: Translate PrettyC HermitM TyLit RetExpr
        ppLitTy = arr $ \case
                          NumTyLit i  -> RetAtom $ tyText (show i)
                          StrTyLit fs -> RetAtom $ tyText (show fs)

        ppTyCon :: Translate PrettyC HermitM TyCon RetExpr
        ppTyCon = GHC.getName ^>> ppName TypeColor

        ppTyConCo :: Translate PrettyC HermitM TyCon RetExpr
        ppTyConCo = GHC.getName ^>> ppName CoercionColor

        -- binders are vars that is bound by lambda or case, etc.
        -- depending on the mode, they might not be displayed
        ppBinderMode :: Translate PrettyC HermitM Var RetExpr
        ppBinderMode = do v <- idR
                          if
                             | GHC.isTyVar v -> case po_exprTypes opts of
                                                                   Omit     -> return (RetEmpty)
                                                                   Abstract -> return (RetAtom typeBindSymbol)
                                                                   _        -> ppVar
                             | GHC.isCoVar v -> case po_coercions opts of
                                                                   Omit     -> return (RetEmpty)
                                                                   Abstract -> return (RetAtom coercionBindSymbol)
                                                                   Show     -> ppVar
                                                                   Kind     -> do pCoKind <- ppCoKind <<^ CoVarCo
                                                                                  return (RetExpr (coercionBindSymbol `hasType` pCoKind))
                                                                  -- TODO: refactor this to be more systematic.  It should be possible to request type sigs for all type bindings.
                             | otherwise       -> ppVar

        ppModGuts :: PrettyH GHC.ModGuts
        ppModGuts = do name <- ppSDoc <<^ GHC.mg_module
                       vtys <- mapT (ppVar &&& (ppType <<^ GHC.idType)) <<< modGutsT progIdsT (\ _ ids -> ids)
                       let defs = [ normalExpr v `hasType` ty | (v,ty) <- vtys ]
                       return $ hang (keyword "module" <+> name <+> keyword "where") 2 (vcat defs)

        -- DocH is not a monoid.
        ppCoreProg :: PrettyH CoreProg
        ppCoreProg = progConsT ppCoreBind ppCoreProg ($+$) <+ progNilT empty

        ppCoreBind :: PrettyH GHC.CoreBind
        ppCoreBind =   (nonRecT idR (ppCoreExprR &&& ppTypeSig) (,) >>> ppDef)
                    <+ (do p <- rootPathT
                           recT (const ppCoreDef) (\ bnds -> attrP p (keyword "rec") <+> vcat bnds)
                       )


        ppCoreAlt :: PrettyH GHC.CoreAlt
        ppCoreAlt = do p <- rootPathT
                       altT (readerT $ \case
                                          GHC.DataAlt dcon -> return (GHC.dataConName dcon) >>> ppName IdColor >>^ parenExpr
                                          GHC.LitAlt lit   -> return lit >>> ppSDoc
                                          GHC.DEFAULT      -> return (symbol '_')
                            )
                            (\ _ -> ppBinderMode)
                            ppCoreExpr
                            (\ con vs e -> hang (con <+> hsep (map parenExpr vs) <+> attrP p (specialSymbol RightArrowSymbol)) 2 e)

        ppCoreDef :: PrettyH CoreDef
        ppCoreDef = defT idR (ppCoreExprR &&& ppTypeSig) (,) >>> ppDef

        ppDef :: PrettyH (Var,(RetExpr,DocH))
        ppDef = do p          <- rootPathT
                   (v,(e,ty)) <- idR
                   case po_coercions opts of
                     Omit | GHC.isCoVar v  -> return empty
                     Kind | GHC.isCoVar v  -> return $ case po_exprTypes opts of
                                                         Show -> (coercionBindSymbol `hasType` ty) $+$ (coercionBindSymbol <+> symbol '=' <+> coercionSymbol)
                                                         _    -> coercionBindSymbol <+> attrP p (symbol '=') <+> normalExpr e
                     _                     -> do pv  <- normalExpr ^<< ppBinderMode <<< return v
                                                 let pre  = pv <+> symbol '='
                                                     body = case e of -- This is an ugly hack
                                                              RetLam p' vs pb e0 -> hang (pre <+> attrP p' (specialSymbol LambdaSymbol) <+> hsep vs <+> attrP pb (specialSymbol RightArrowSymbol)) 2 e0
                                                              _                 -> hang pre 2 (normalExpr e)
                                                 return $ case po_exprTypes opts of
                                                            Omit | GHC.isTyVar v -> empty
                                                            Show                 -> (pv `hasType` ty) $+$ body
                                                            _                    -> body


        ppCoreExpr :: PrettyH GHC.CoreExpr
        ppCoreExpr = ppCoreExprR >>^ normalExpr

        ppCoreExprR :: Translate PrettyC HermitM GHC.CoreExpr RetExpr
        ppCoreExprR = rootPathT >>= ppCoreExprPR
          where
            ppCoreExprPR :: PathH -> Translate PrettyC HermitM GHC.CoreExpr RetExpr
            ppCoreExprPR p =
                      lamT ppBinderMode ppCoreExprR (retLam p)

                   <+ letT (RetExpr <$> ppCoreBind) ppCoreExprR (retLet p)
                   -- HACKs
    {-
                   <+ (acceptR (\ e -> case e of
                                         GHC.App (Var v) (GHC.Type t) | po_exprTypes opts == Abstract -> True
                                         _ -> False) >>>
                               (appT ppCoreExprR ppCoreExprR (\ (RetAtom e1) (RetAtom e2) ->
                                        RetAtom (e1 <+> e2))))
    -}
                   <+ (acceptR (\ e -> case e of
                                         App (Type _) (Lam {}) | po_exprTypes opts == Omit -> True
                                         App (App (Var _) (Type _)) (Lam {}) | po_exprTypes opts == Omit -> True
                                         _ -> False) >>>
                               (appT ppCoreExprR ppCoreExprR (\ (RetAtom e1) (RetLam p' vs pb e0) ->
                                        RetExpr $ hang (e1 <+>
                                                            symbol '(' <>
                                                            attrP p' (specialSymbol LambdaSymbol) <+>
                                                            hsep vs <+>
                                                            attrP pb (specialSymbol RightArrowSymbol)) 2 (e0 <> symbol ')')))
                      )

                   <+ appT ppCoreExprR ppCoreExprR retApp
                   <+ caseT ppCoreExpr ppVar ppTypeModeR (const ppCoreAlt) (\ s w ty alts -> RetExpr $ attrP p ((keyword "case" <+> s <+> keyword "of" <+> parenExpr w <+> parenExpr ty) $$ nest 2 (vcat alts)))

                   <+ varT (attrPAtomExpr p <$> (do (c,i) <- exposeT
                                                    if (GHC.isLocalId i) && (i `notMember` boundVars c)
                                                      then GHC.varName ^>> ppName WarningColor
                                                      else ppVar
                                                )
                           )

                   <+ litT ((RetAtom . attrP p) <$> ppSDoc)
                   <+ typeT (attrPAtomExpr p <$> ppTypeModeR)
                   <+ coercionT (attrPAtomExpr p <$> ppCoercionModeR)
                   <+ castT ppCoreExprR ppCoercionModeR (\ e co -> if isEmptyR co
                                                                    then e
                                                                    else RetExpr (parenExprExceptApp e <+> attrP p castSymbol <+> parenExprExceptApp co)
                                                        )
                   <+ tickT ppSDoc ppCoreExprR (\ tk e -> RetExpr $ attrP p (text "Tick" $$ nest 2 (tk <+> parenExpr e)))

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

        ppType :: PrettyH Type
        ppType = ppTypeR >>^ normalExpr

        ppTypeModeR :: Translate PrettyC HermitM Type RetExpr
        ppTypeModeR = case po_exprTypes opts of
                        Omit     -> return RetEmpty
                        Abstract -> return (RetAtom typeSymbol)
                        _        -> ppTypeR

        ppTypeR :: Translate PrettyC HermitM Type RetExpr
        ppTypeR = rootPathT >>= ppTypePR
          where
            ppTypePR :: PathH -> Translate PrettyC HermitM Type RetExpr
            ppTypePR p =
                   tyVarT (attrPAtomExpr p <$> ppVar)
                <+ litTyT (attrPAtomExpr p <$> ppLitTy)
                <+ appTyT ppTypeR ppTypeR retApp
                <+ funTyT ppTypeR ppTypeR retArrowType
                <+ forAllTyT ppVar ppTypeR (retForAll p ForAllTy_Body)
                <+ tyConAppT (ppTyCon &&& idR) (const ppTypeR)
                     (\ (pCon,tyCon) tys -> if | GHC.isFunTyCon tyCon && length tys == 2 -> let [ty1,ty2] = tys in retArrowType ty1 ty2
                                               | tyCon == GHC.listTyCon -> RetAtom $ attrP p $ tyText "[" <> (case tys of
                                                                                                                []  -> empty
                                                                                                                t:_ -> normalExpr t)
                                                                                                          <> tyText "]"
                                               | GHC.isTupleTyCon tyCon -> RetAtom $ attrP p $ tyText "(" <> (if null tys
                                                                                                                then empty
                                                                                                                else foldr1 (\ ty r -> ty <> tyText "," <+> r) (map normalExpr tys)
                                                                                                             )
                                                                                                          <> tyText ")"
                                               | otherwise              -> retApps (attrPAtomExpr p pCon) tys
                     )

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

        ppCoercion :: PrettyH Coercion
        ppCoercion = ppCoercionR >>^ normalExpr

        ppCoercionModeR :: Translate PrettyC HermitM Coercion RetExpr
        ppCoercionModeR = case po_coercions opts of
                            Omit     -> return RetEmpty
                            Abstract -> return (RetAtom coercionSymbol)
                            Show     -> ppCoercionR
                            Kind     -> ppCoKind >>^ (\ k -> RetExpr (coercionSymbol `hasType` k))

        ppCoercionR :: Translate PrettyC HermitM Coercion RetExpr
        ppCoercionR = rootPathT >>= ppCoercionPR
          where
            ppCoercionPR :: PathH -> Translate PrettyC HermitM Coercion RetExpr
            ppCoercionPR p =
                           reflT (ppTypeModeR >>^ \ ty -> RetAtom $ attrP p $ if isEmptyR ty then coText "refl" else coText "<" <> normalExpr ty <> coText ">")
                        <+ coVarCoT (attrPAtomExpr p <$> ppVar)
                        <+ symCoT (ppCoercionR >>^ \ co -> RetExpr (attrP p (coKeyword "sym") <+> parenExpr co))
                        <+ forAllCoT ppBinderMode ppCoercionR (retForAll p ForAllCo_Body)
                        <+ transCoT ppCoercionR ppCoercionR (\ co1 co2 -> RetExpr (parenExprExceptApp co1 <+> attrP p (coChar ';') <+> parenExprExceptApp co2))
                        <+ unsafeCoT ppTypeModeR ppTypeModeR (\ ty1 ty2 -> (if isEmptyR ty1 && isEmptyR ty2 then RetAtom else RetExpr)
                                                                           (attrP p (coKeyword "unsafe") <+> parenExpr ty1 <+> parenExpr ty2)
                                                             )
                        <+ nthCoT idR ppCoercionR (\ n co -> RetExpr (attrP p (coKeyword "nth") <+> attrP p (coText $ show n) <+> parenExpr co))
                        <+ instCoT ppCoercionR ppTypeModeR (\ co ty -> if isEmptyR ty
                                                                         then RetExpr (attrP p (coText "inst") <+> parenExpr co)
                                                                         else RetExpr (parenExprExceptApp co <+> attrP p (coChar '@') <+> parenExprExceptApp ty)
                                                           )
                        <+ tyConAppCoT (attrPAtomExpr p <$> ppTyConCo) (const ppCoercionR) retApps
                        <+ appCoT ppCoercionR ppCoercionR retApp
#if __GLASGOW_HASKELL__ > 706
        -- TODO: Figure out how to properly pp new branched Axioms and Left/Right Coercions
                        <+ axiomInstCoT (coAxiomName ^>> ppName CoercionColor) (RetAtom <$> ppSDoc) (const ppCoercionR) (\ ax idx coes -> RetExpr (attrP p (coText "axiomInst") <+> attrP p (parenExpr ax) <+> attrP p (parenExpr idx) <+> sep (map parenExpr coes)))
                        <+ lrCoT (coercionColor <$> ppSDoc) ppCoercionR (\ lr co -> RetExpr (attrP p lr <+> parenExpr co))
#else
                        <+ axiomInstCoT (coAxiomName ^>> ppName CoercionColor) (const ppCoercionR) (\ ax coes -> RetExpr (attrP p (coText "axiomInst") <+> attrP p (parenExpr ax) <+> sep (map parenExpr coes)))
#endif

        ppCoKind :: PrettyH Coercion
        ppCoKind = (GHC.coercionKind >>> unPair) ^>> (ppTypeModeR *** ppTypeModeR) >>^ ( \(ty1,ty2) -> parenExprExceptApp ty1 <+> coText "~#" <+> parenExprExceptApp ty2)

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

        ppTypeSig :: PrettyH GHC.CoreExpr
        ppTypeSig = coercionT ppCoKind <+ (GHC.exprType ^>> ppType)

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

    in  promoteT (ppCoreExpr     :: PrettyH GHC.CoreExpr)
     <+ promoteT (ppCoreProg     :: PrettyH CoreProg)
     <+ promoteT (ppCoreBind     :: PrettyH GHC.CoreBind)
     <+ promoteT (ppCoreDef      :: PrettyH CoreDef)
     <+ promoteT (ppModGuts      :: PrettyH GHC.ModGuts)
     <+ promoteT (ppCoreAlt      :: PrettyH GHC.CoreAlt)
     <+ promoteT (ppType         :: PrettyH GHC.Type)
     <+ promoteT (ppCoercion     :: PrettyH Coercion)

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