-- | Pretty-printing
{-# LANGUAGE
      PatternGuards,
      QuasiQuotes,
      TypeSynonymInstances
    #-}
module Ppr (
  pprTyApp,
  -- * Re-exports
  module PprClass,
  module Prec
) where

import Meta.Quasi
import PprClass
import Prec
import Syntax
import Util

import qualified Syntax.Decl
import qualified Syntax.Expr
import qualified Syntax.Notable
import qualified Syntax.Patt
import qualified Loc

import Data.List (intersperse)

instance IsInfix (Type i) where
  isInfix [$ty| ($_, $_) $lid:n |] = isOperator n
  isInfix [$ty| $_ -[$_]> $_ |]    = True
  isInfix _                        = False

-- | To pretty print the application of a type constructor to
--   generic parameters
pprTyApp :: (Ppr a) => Int -> QLid i -> [a] -> Doc
pprTyApp _ ql       []   = ppr ql
pprTyApp p (J [] l) [t1]
  | isOperator l, precOp (unLid l) == Right precBang
    = parensIf (p > precBang) $
        text (unLid l) <> pprPrec (precBang + 1) t1
pprTyApp p (J [] l) [t1, t2]
  -- print @ without space around it:
  | isOperator l, '@':_ <- unLid l, Right prec <- precOp (unLid l)
    = parensIf (p > prec) $
        pprPrec (prec + 1) t1 <> text (unLid l) <> pprPrec prec t2
  | isOperator l, Left prec <- precOp (unLid l)
    = parensIf (p > prec) $
        sep [ pprPrec prec t1,
              text (unLid l) <+> pprPrec (prec + 1) t2 ]
  | isOperator l, Right prec <- precOp (unLid l)
    = parensIf (p > prec) $
        sep [ pprPrec (prec + 1) t1,
              text (unLid l) <+> pprPrec prec t2]
pprTyApp p ql ts = parensIf (p > precApp) $
                     sep [ pprPrec precApp ts,
                           ppr ql ]

instance Ppr (Type i) where
  -- Print sugar for infix type constructors:
  pprPrec p [$ty| $t1 ; $t2 |]
                  = parensIf (p > precSemi) $
                      sep [ pprPrec (precSemi + 1) t1 <> text ";",
                            pprPrec precSemi t2 ]
  -- pprPrec p (TyFun q t1 t2)
  pprPrec p [$ty| $t1 -[$q]> $t2 |]
                  = parensIf (p > precArr) $
                    sep [ pprPrec (precArr + 1) t1,
                          pprArr (view q) <+> pprRight precArr t2 ]
    where pprArr (QeLit Qu) = text "->"
          pprArr (QeLit Qa) = text "-o"
          pprArr _          = text "-[" <> pprPrec precStart q <> text "]>"
  pprPrec p [$ty| ($list:ts) $qlid:n |]
                          = pprTyApp p n ts
    -- debugging: <> text (show (ttId (unsafeCoerce tag :: TyTag)))
  pprPrec p [$ty| '$x |]  = pprPrec p x
  pprPrec p [$ty| $quant:qu '$x. $t |]
                          = parensIf (p > precDot) $
                              ppr qu <+>
                              fsep (map (pprPrec (precDot + 1))
                                        tvs) <>
                              char '.'
                                >+> pprPrec precDot body
      where (tvs, body) = unfoldTyQu qu [$ty| $quant:qu '$x. $t |]
  pprPrec p [$ty| mu '$x. $t |]
                          = parensIf (p > precDot) $
                              text "mu" <+>
                              pprPrec (precDot + 1) x <>
                              char '.'
                                >+> pprPrec precDot t
  pprPrec p [$ty| $anti:a |] = pprPrec p a

instance Ppr (TyPat i) where
  pprPrec p tp0 = case tp0 of
    N _ (TpVar tv var) -> pprParamV var tv
    [$tpQ| ($list:tps) $qlid:ql |]
                       -> pprTyApp p ql tps
    [$tpQ| $antiP:a |] -> ppr a

instance Ppr (QExp i) where
  pprPrec p [$qeQ| $qlit:qu |] = pprPrec p qu
  pprPrec p [$qeQ| $qvar:v |]  = pprPrec p (tvname v)
  pprPrec p [$qeQ| $qdisj:qes |] = case qes of
    []    -> pprPrec p Qu
    [qe]  -> pprPrec p qe
    _     -> parensIf (p > precPlus) $
               fsep $
                 intersperse (text "\\/") $
                   map (pprPrec (precPlus + 1)) qes
  pprPrec p [$qeQ| $qconj:qes |] = case qes of
    []    -> pprPrec p Qa
    [qe]  -> pprPrec p qe
    _     -> parensIf (p > precStar) $
               hcat $
                 intersperse (text "/\\") $
                   map (pprPrec (precStar + 1)) qes
  pprPrec p [$qeQ| $anti:a |] = pprPrec p a

instance Ppr (Prog i) where
  ppr [$prQ| $list:ms |]       = vcat (map ppr ms)
  ppr [$prQ| $expr:e |]        = ppr e
  ppr [$prQ| $list:ms in $e |] = vcat (map ppr ms) $+$
                                 (text "in" >+> ppr e)

instance Ppr (Decl i) where
  ppr [$dc| let $x = $e |] = sep
    [ text "let" <+> ppr x,
      nest 2 $ equals <+> ppr e ]
  ppr [$dc| let $x : $t = $e |] = sep
    [ text "let" <+> ppr x,
      nest 2 $ colon <+> ppr t,
      nest 4 $ equals <+> ppr e ]
  ppr [$dc| type $list:tds |] = pprTyDecs tds
  ppr [$dc| abstype $list:ats0 with $list:ds end |] =
    case ats0 of
      []     ->
        vcat [
          text "abstype with",
          nest 2 $ vcat (map ppr ds),
          text "end"
        ]
      at:ats ->
        vcat [
          vcat (text "abstype" <+> pprAbsTy at :
                [ nest 4 $ text "and" <+> pprAbsTy ati | ati <- ats ])
            <+> text "with",
          nest 2 $ vcat (map ppr ds),
          text "end"
        ]
  ppr [$dc| open $b |] = pprModExp (text "open" <+>) b
  ppr [$dc| module $uid:n : $s = $b |] = pprModExp add1 b where
    add1 body = pprSigExp add2 s <+> equals <+> body
    add2 body = text "module" <+> ppr n <+> colon <+> body
  ppr [$dc| module $uid:n = $b |] = pprModExp add b where
    add body = text "module" <+> ppr n <+> equals <+> body
  ppr [$dc| module type $uid:n = $s |] = pprSigExp add s where
    add body = text "module type" <+> ppr n <+> equals <+> body
  ppr [$dc| local $list:d0 with $list:d1 end |] =
    vcat [
      text "local",
      nest 2 (vcat (map ppr d0)),
      text "with",
      nest 2 (vcat (map ppr d1)),
      text "end"
    ]
  ppr [$dc| exception $uid:n of $opt:mt |] =
    pprExcDec n mt
  ppr [$dc| $anti:a |] = ppr a

pprTyDecs :: [TyDec i] -> Doc
pprTyDecs tds =
  vcat $
    mapHead (text "type" <+>) $
      mapTail ((nest 1) . (text "and" <+>)) $
        map ppr tds

pprExcDec :: Uid i -> Maybe (Type i) -> Doc
pprExcDec u Nothing  =
  text "exception" <+> ppr u
pprExcDec u (Just t) =
  text "exception" <+> ppr u <+> text "of" <+> ppr t

instance Ppr (TyDec i) where
  ppr td = case view td of
    TdAbs n ps vs qs  -> pprProtoV n vs ps >?> pprQuals qs
    TdSyn n [(ps,t)]  -> pprProto n ps >+> equals <+> ppr t
    TdSyn n cs        -> vcat [ char '|' <+> each ci | ci <- cs ]
      where
        each (ps, rhs) = pprProto n ps
                           >+> char '=' <+> ppr rhs
    TdDat n ps alts   -> pprProtoV n (repeat Invariant) ps
                           >?> pprAlternatives alts
    TdAnti a          -> ppr a

pprAbsTy :: AbsTy i -> Doc
pprAbsTy at = case view at of
  AbsTy variances qual (N _ (TdDat name params alts)) ->
    pprProtoV name variances params
      >?> pprQuals qual
      >?> pprAlternatives alts
  AbsTy _ _ td -> ppr td -- shouldn't happen (yet)
  AbsTyAnti a -> ppr a

pprProto  :: Lid i -> [TyPat i] -> Doc
pprProto n ps = ppr (tpApp (J [] n) ps)

pprProtoV :: Lid i -> [Variance] -> [TyVar i] -> Doc
pprProtoV n vs tvs = pprProto n (zipWith tpVar tvs vs)

pprParamV :: Variance -> TyVar i -> Doc
pprParamV Invariant tv = ppr tv
pprParamV v         tv = ppr v <> ppr tv

pprQuals :: QExp i -> Doc
pprQuals [$qeQ| U |] = empty
pprQuals qs          = text ":" <+> pprPrec precApp qs

pprAlternatives :: [(Uid i, Maybe (Type i))] -> Doc
pprAlternatives [] = equals
pprAlternatives (a:as) = sep $
  equals <+> alt a : [ char '|' <+> alt a' | a' <- as ]
  where
    alt (u, Nothing) = ppr u
    alt (u, Just t)  = ppr u <+> text "of" <+> pprPrec precDot t

pprModExp :: (Doc -> Doc) -> ModExp i -> Doc
pprModExp add modexp = case modexp of
  [$me| $quid:n |] -> add (ppr n)
  [$me| $quid:n $list:qls |] ->
    add (ppr n) <+>
    brackets (fsep (punctuate comma (map ppr qls)))
  [$me| struct $list:ds end |] ->
    add (text "struct")
    $$ nest 2 (vcat (map ppr ds))
    $$ text "end"
  [$me| $me1 : $se2 |] ->
    pprSigExp (pprModExp add me1 <+> colon <+>) se2
  [$me| $anti:a |] -> add (ppr a)

pprSigExp :: (Doc -> Doc) -> SigExp i -> Doc
pprSigExp add se0 = body >+> withs where
  (wts, se1) = unfoldSeWith se0
  body       = case se1 of
    [$seQ| $quid:n |] -> add (ppr n)
    [$seQ| $quid:n $list:qls |] ->
      add (ppr n) <+>
      brackets (fsep (punctuate comma (map ppr qls)))
    [$seQ| sig $list:sgs end |] ->
      add (text "sig")
      $$ nest 2 (vcat (map ppr sgs))
      $$ text "end"
    [$seQ| $_ with type $list:_ $qlid:_ = $_ |] ->
      error "BUG! can't happen in pprSigExp"
    [$seQ| $anti:a |] -> add (ppr a)
  withs      =
    sep $
      mapHead (text "with type" <+>) $
        mapTail ((nest 6) . (text "and" <+>)) $
          [ pprTyApp 0 tc tvs <+> equals <+> ppr t
          | (tc, tvs, t) <- wts ]

instance Ppr (SigItem i) where
  ppr sg0 = case sg0 of
    [$sgQ| val $lid:n : $t |] ->
      text "val" <+> ppr n >+> colon <+> ppr t
    [$sgQ| type $list:tds |] ->
      pprTyDecs tds
    [$sgQ| module $uid:u : $s |] ->
      pprSigExp add s where
        add body = text "module" <+> ppr u <+> colon <+> body
    [$sgQ| module type $uid:u = $s |] ->
      pprSigExp add s where
        add body = text "module type" <+> ppr u <+> equals <+> body
    [$sgQ| include $s |] ->
      pprSigExp (text "include" <+>) s
    [$sgQ| exception $uid:u of $opt:mt |] ->
      pprExcDec u mt
    [$sgQ| $anti:a |] ->
      ppr a

instance Ppr (Expr i) where
  pprPrec p e0 = case e0 of
    [$ex| $id:x |]   -> pprPrec p x
    [$ex| $lit:lt |] -> pprPrec p lt
    [$ex| if $ec then $et else $ef |] ->
      parensIf (p > precDot) $
        sep [ text "if" <+> ppr ec,
              nest 2 $ text "then" <+> ppr et,
              nest 2 $ text "else" <+> pprPrec precDot ef ]
    [$ex| $e1; $e2 |] ->
      parensIf (p > precSemi) $
        sep [ pprPrec (precSemi + 1) e1 <> semi,
              ppr e2 ]
    [$ex| let $x = $e1 in $e2 |] ->
      pprLet p (ppr x) e1 e2
    [$ex| match $e1 with $list:clauses |] ->
      parensIf (p > precDot) $
        vcat (sep [ text "match",
                    nest 2 $ ppr e1,
                    text "with" ] : map alt clauses)
      where
        alt (N _ (CaClause xi ei)) =
          hang (char '|' <+> pprPrec precDot xi <+> text "->")
                4
                (pprPrec precDot ei)
        alt (N _ (CaAnti a))      = char '|' <+> ppr a
    [$ex| let rec $list:bs in $e2 |] ->
      text "let" <+>
      vcat (zipWith each ("rec" : repeat "and") bs) $$
      text "in" <+> pprPrec precDot e2
        where
          each kw (N _ (BnBind x t e)) =
            -- This could be better by pulling some args out.
            hang (hang (text kw <+> ppr x)
                       6
                       (colon <+> ppr t <+> equals))
                 2
                 (ppr e)
          each kw (N _ (BnAnti a)) = text kw <+> ppr a
    [$ex| let $decl:d in $e2 |] ->
      text "let" <+> ppr d $$
      (text "in" >+> pprPrec precDot e2)
    [$ex| ($e1, $e2) |] ->
      parensIf (p > precCom) $
        sep [ pprPrec precCom e1 <> comma,
              pprPrec (precCom + 1) e2 ]
    [$ex| fun $_ : $_ -> $_ |] -> pprAbs p e0
    [$ex| $name:x $e2 |]
      | Right p' <- precOp x,
        p' == 10
          -> parensIf (p > p') $
               text x <+> pprPrec p' e2
    [$ex| ($name:x $e12) $e2 |] 
      | (dl, dr, p') <- either ((,,) 0 1) ((,,) 1 0) (precOp x),
        p' < 9
          -> parensIf (p > p') $
               sep [ pprPrec (p' + dl) e12,
                     text x,
                     pprPrec (p' + dr) e2 ]
    [$ex| $e1 $e2 |]
          -> parensIf (p > precApp) $
               sep [ pprPrec precApp e1,
                     pprPrec (precApp + 1) e2 ]
    [$ex| fun '$_ -> $_ |] -> pprAbs p e0
    [$ex| $_ [$_] |] ->
      parensIf (p > precTApp) $
        cat [ pprPrec precTApp op,
              brackets . fsep . punctuate comma $
                map (pprPrec precCom) args ]
      where 
        (args, op) = unfoldExTApp e0
    [$ex| Pack[$opt:t1]($t2, $e) |] ->
      parensIf (p > precApp) $
        text "Pack" <> maybe empty (brackets . ppr) t1 <+>
        parens (sep [ pprPrec (precCom + 1) t2 <> comma,
                      pprPrec precCom e ])
    [$ex| ( $e : $t1 :> $t2 ) |] ->
      parensIf (p > precCast) $
         sep [ pprPrec (precCast + 2) e,
               colon     <+> pprPrec (precCast + 2) t1,
               text ":>" <+> pprPrec (precCast + 2) t2 ]
    [$ex| ( $e : $t1 ) |] ->
      parensIf (p > precCast) $
         sep [ pprPrec (precCast + 2) e,
               colon     <+> pprPrec (precCast + 2) t1 ]
    [$ex| ( $e :> $t1 ) |] ->
      parensIf (p > precCast) $
         sep [ pprPrec (precCast + 2) e,
               text ":>" <+> pprPrec (precCast + 2) t1 ]
    [$ex| $anti:a |] -> pprPrec p a

pprLet :: Int -> Doc -> Expr i -> Expr i -> Doc
pprLet p pat e1 e2 = parensIf (p > precDot) $
  hang (text "let" <+> pat <+> pprArgList args <+> equals
          >+> ppr body <+> text "in")
       (if isLet (view e2)
          then 0
          else 2)
       (pprPrec precDot e2)
  where
    (args, body) = unfoldExAbs e1
    isLet (ExCase _ [_]) = True
    isLet _              = False

pprAbs :: Int -> Expr i -> Doc
pprAbs p e = parensIf (p > precDot) $
    text "fun" <+> argsDoc <+> text "->"
      >+> pprPrec precDot body
  where (args, body)   = unfoldExAbs e
        argsDoc = case args of
          [Left ([$pa| _ |], [$ty|@! unit |])]
                        -> parens empty
          [Left (x, t)] -> ppr x <+> char ':' <+> pprPrec (precArr + 1) t
          _             -> pprArgList args

pprArgList :: [Either (Patt i, Type i) (TyVar i)] -> Doc
pprArgList = fsep . map eachArg . combine where
  eachArg (Left ([$pa| _ |], [$ty|@! unit |]))
                          = parens empty
  eachArg (Left (x, t))   = parens $
                              ppr x
                                >+> colon <+> ppr t
  eachArg (Right tvs)     = brackets .
                              sep .
                                punctuate comma $
                                  map ppr tvs

  combine :: [Either a b] -> [Either a [b]]
  combine  = foldr each [] where
    each (Right b) (Right bs : es) = Right (b : bs) : es
    each (Right b) es              = Right [b] : es
    each (Left a)  es              = Left a : es

instance Ppr (Patt i) where
  pprPrec _ [$pa| _ |]             = text "_"
  pprPrec _ [$pa| $lid:l |]        = ppr l
  pprPrec _ [$pa| $quid:qu |]      = ppr qu
  pprPrec p [$pa| $quid:qu $x |]   = parensIf (p > precApp) $
                                       pprPrec precApp qu <+>
                                       pprPrec (precApp + 1) x
  pprPrec p [$pa| ($x, $y) |]      = parensIf (p > precCom) $
                                       pprPrec precCom x <> comma <+>
                                       pprPrec (precCom + 1) y
  pprPrec p [$pa| $lit:lt |]       = pprPrec p lt
  pprPrec p [$pa| $x as $lid:l |]  = parensIf (p > precDot) $
                                       pprPrec (precDot + 1) x <+>
                                       text "as" <+> ppr l
  pprPrec p [$pa| Pack('$tv,$x) |] = parensIf (p > precApp) $
                                       text "Pack" <+> parens (sep pair)
    where pair = [ pprPrec (precCom + 1) tv <> comma,
                   pprPrec precCom x ]
  pprPrec p [$pa| $anti:a |]       = pprPrec p a

instance Ppr Lit where
  ppr (LtInt i)   = integer i
  ppr (LtFloat f) = double f
  ppr (LtStr s)   = text (show s)
  ppr (LtAnti a)  = ppr a

instance Show (Prog i)   where showsPrec = showFromPpr
instance Show (Decl i)   where showsPrec = showFromPpr
instance Show (TyDec i)  where showsPrec = showFromPpr
instance Show (Expr i)   where showsPrec = showFromPpr
instance Show (Patt i)   where showsPrec = showFromPpr
instance Show Lit        where showsPrec = showFromPpr
instance Show (Type i)   where showsPrec = showFromPpr
instance Show (TyPat i)  where showsPrec = showFromPpr
instance Show (QExp i)   where showsPrec = showFromPpr
instance Show (SigItem i)where showsPrec = showFromPpr

instance Ppr QLit      where pprPrec = pprFromShow
instance Ppr Variance  where pprPrec = pprFromShow
instance Ppr Quant     where pprPrec = pprFromShow
instance Ppr (Lid i)   where pprPrec = pprFromShow
instance Ppr (Uid i)   where pprPrec = pprFromShow
instance Ppr (BIdent i)where pprPrec = pprFromShow
instance Ppr (TyVar i) where pprPrec = pprFromShow
instance Ppr Anti      where pprPrec = pprFromShow
instance (Show p, Show k) => Ppr (Path p k) where pprPrec = pprFromShow