module Language.Java.Paragon.Pretty (pretty, prettyPrint, Pretty(..)) where

import Text.PrettyPrint
import Data.Char (toLower)

import Language.Java.Paragon.Syntax


prettyPrint :: Pretty a => a -> String
prettyPrint = show . pretty

class Pretty a where
  pretty :: a -> Doc
  pretty = prettyPrec 0
  
  prettyPrec :: Int -> a -> Doc
  prettyPrec _ = pretty

-----------------------------------------------------------------------
-- Packages

instance Pretty CompilationUnit where
  pretty (CompilationUnit mpd ids tds) =
    vcat $ maybePP mpd: map pretty ids ++ map pretty tds

instance Pretty PackageDecl where
  pretty (PackageDecl name) = text "package" <+> pretty name <> semi

instance Pretty ImportDecl where
  pretty (ImportDecl st name wc) = 
    text "import" <+> opt st (text "static")
                  <+> pretty name <> opt wc (text ".*")
                  <> semi

-----------------------------------------------------------------------
-- Declarations

instance Pretty TypeDecl where
  pretty (ClassTypeDecl     cd) = pretty cd
  pretty (InterfaceTypeDecl id) = pretty id

instance Pretty ClassDecl where
  pretty (EnumDecl mods ident impls body) = 
    hsep [hsep (map pretty mods)
          , text "enum"
          , pretty ident 
          , ppImplements impls
         ] $$ pretty body

  pretty (ClassDecl mods ident tParams mSuper impls body) =
    hsep [hsep (map pretty mods)
          , text "class"
          , pretty ident
          , ppTypeParams tParams
          , ppExtends (maybe [] return mSuper)
          , ppImplements impls
         ] $$ pretty body

instance Pretty ClassBody where
  pretty (ClassBody ds) = 
    braceBlock (map pretty ds)
    
instance Pretty EnumBody where
  pretty (EnumBody cs ds) =
    braceBlock $ 
        punctuate comma (map pretty cs) ++ 
        opt (not $ null ds) semi : map pretty ds

instance Pretty EnumConstant where
  pretty (EnumConstant ident args mBody) =
    pretty ident 
        -- needs special treatment since even the parens are optional
        <> opt (not $ null args) (ppArgs args) 
      $$ maybePP mBody

instance Pretty InterfaceDecl where
  pretty (InterfaceDecl mods ident tParams impls body) =
    hsep [hsep (map pretty mods)
          , text "interface"
          , pretty ident
          , ppTypeParams tParams
          , ppImplements impls
         ] $$ pretty body

instance Pretty InterfaceBody where
  pretty (InterfaceBody mds) =
    braceBlock (map pretty mds)

instance Pretty Decl where
  pretty (MemberDecl md) = pretty md
  pretty (InitDecl b bl) = 
    opt b (text "static") <+> pretty bl

instance Pretty MemberDecl where
  pretty (FieldDecl mods t vds) =
    hsep (map pretty mods ++ pretty t:map pretty vds) <> semi

  pretty (MethodDecl mods tParams mt ident fParams throws body) =
    hsep [hsep (map pretty mods)
          , ppTypeParams tParams
          , ppResultType mt
          , pretty ident <> ppArgs fParams
          , ppThrows throws
         ] $$ pretty body

  pretty (ConstructorDecl mods tParams ident fParams throws body) =
    hsep [hsep (map pretty mods)
          , ppTypeParams tParams
          , pretty ident
          , ppArgs fParams
          , ppThrows throws
         ] $$ pretty body

  pretty (MemberClassDecl cd) = pretty cd
  pretty (MemberInterfaceDecl id) = pretty id

-- Paragon
  pretty (LockDecl mods ident arity lockProps) =
    hsep [hsep (map pretty mods)
          , text "lock"
          , pretty ident <> ppArity arity
         ] <> maybe empty (\lp -> char ' ' <> pretty lp) lockProps <> semi

{-  pretty (PolicyDecl mods ident pol) =
    hsep [hsep (map pretty mods)
          , text "policy"
          , pretty ident
          , char '='
          , pretty pol
         ] <> semi

  pretty (ActorDecl mods ident mInit) =
    hsep [hsep (map pretty mods)
          , text "actor"
          , pretty ident
          , maybe empty (\init -> char '=' <+> pretty init) mInit
         ] <> semi -}

ppArity :: [Maybe Ident] -> Doc
ppArity [] = empty
ppArity mis = parens $ hsep $ punctuate comma $ map (\i -> {- char '?' <> -} maybePP i) mis

instance Pretty VarDecl where
  pretty (VarDecl vdId mInit) =
    pretty vdId 
        <+> maybe empty (\init -> char '=' <+> pretty init) mInit

instance Pretty VarDeclId where
  pretty (VarId ident) = pretty ident
  pretty (VarDeclArray vId) = pretty vId

instance Pretty VarInit where
  pretty (InitExp e) = pretty e
  pretty (InitArray arrInit) = pretty arrInit

instance Pretty FormalParam where
  pretty (FormalParam mods t b vId) =
    hsep [hsep (map pretty mods)
          , pretty t <> opt b (text "...")
          , pretty vId
         ]

instance Pretty MethodBody where
  pretty (MethodBody mBlock) = maybe semi pretty mBlock

instance Pretty ConstructorBody where
  pretty (ConstructorBody mECI stmts) =
    braceBlock $ maybePP mECI : map pretty stmts

instance Pretty ExplConstrInv where
  pretty (ThisInvoke rts args) =
    ppTypeParams rts <+> text "this" <> ppArgs args <> semi
  pretty (SuperInvoke rts args) =
    ppTypeParams rts <+> text "super" <> ppArgs args <> semi
  pretty (PrimarySuperInvoke e rts args) =
    pretty e <> char '.' <>
      ppTypeParams rts <+> text "super" <> ppArgs args <> semi

instance Pretty Modifier where
  pretty mod = case mod of
                 Reads pol     -> char '?' <> pretty pol
                 Writes pol    -> char '!' <> pretty pol
                 Expects locks -> char '~' <> prettyLocks locks
                 Opens   locks -> char '+' <> prettyLocks locks
                 Closes  locks -> char '-' <> prettyLocks locks

                 _ -> text . map toLower $ show mod

-----------------------------------------------------------------------
-- Statements


instance Pretty Block where
  pretty (Block stmts) = braceBlock $ map pretty stmts

instance Pretty BlockStmt where
  pretty (BlockStmt stmt) = pretty stmt
  pretty (LocalClass cd) = pretty cd
  pretty (LocalVars mods t vds) =
    hsep (map pretty mods) <+> pretty t <+>
      hsep (punctuate comma $ map pretty vds) <> semi

-- Paragon
  pretty (LocalLock mods ident arity lockProps) =
    hsep [hsep (map pretty mods)
          , text "lock"
          , pretty ident <> ppArity arity
         ] <> maybe empty (\lp -> char ' ' <> pretty lp) lockProps <> semi
  
{-  pretty (LocalPolicy mods ident pol) =
    hsep [hsep (map pretty mods)
          , text "policy"
          , pretty ident
          , char '='
          , pretty pol
         ] <> semi
  
  pretty (LocalActor mods ident mInit) =
    hsep [hsep (map pretty mods)
          , text "actor"
          , pretty ident
          , maybe empty (\init -> char '=' <+> pretty init) mInit
         ] <> semi -}

instance Pretty Stmt where
  pretty (StmtBlock block) = pretty block
  pretty (IfThen c th) =
    hsep [text "if", parens (pretty c) , pretty th
         ]

  pretty (IfThenElse c th el) =
    hsep [text "if", parens (pretty c)
          , pretty th , text "else", pretty el
         ]
      
  pretty (While c stmt) =
    hsep [text "while", parens (pretty c), pretty stmt]
  
  pretty (BasicFor mInit mE mUp stmt) =
    hsep [text "for"
          , parens $ hsep [maybePP mInit, semi
                           , maybePP mE, semi
                           , maybe empty (hsep . punctuate comma . map pretty) mUp
                          ]
          , pretty stmt
         ]

  pretty (EnhancedFor mods t ident e stmt) =
    hsep [text "for"
          , parens $ hsep [
                  hsep (map pretty mods)
                , pretty t
                , pretty ident
                , colon
                , pretty e
               ]
          , pretty stmt
         ]

  pretty Empty = semi
  
  pretty (ExpStmt e) = pretty e <> semi

  pretty (Assert ass mE) =
    text "assert" <+> pretty ass
      <+> maybe empty ((colon <>) . pretty) mE <> semi

  pretty (Switch e sBlocks) =
    text "switch" <+> parens (pretty e) 
      $$ braceBlock (map pretty sBlocks)

  pretty (Do stmt e) =
    hsep [text "do", pretty stmt, text "while"
          , parens (pretty e)] <> semi
  
  pretty (Break mIdent) =
    text "break" <+> maybePP mIdent <> semi
  
  pretty (Continue mIdent) =
    text "continue" <+> maybePP mIdent <> semi
  
  pretty (Return mE) =
    text "return" <+> maybePP mE <> semi
  
  pretty (Synchronized e block) =
    text "synchronized" <+> parens (pretty e) $$ pretty block
  
  pretty (Throw e) =
    text "throw" <+> pretty e <> semi
  
  pretty (Try block catches mFinally) =
    text "try" $$ pretty block $$
      vcat (map pretty catches ++ [ppFinally mFinally])
   where ppFinally Nothing = empty
         ppFinally (Just bl) = text "finally" <+> pretty bl

  pretty (Labeled ident stmt) =
    pretty ident <> colon <+> pretty stmt

-- Paragon
  pretty (Open  lock) = text "open"  <+> pretty lock <> semi
  pretty (Close lock) = text "close" <+> pretty lock <> semi

  pretty (OpenBlock  lock block) = text "open"  <+> pretty lock <+> pretty block
  pretty (CloseBlock lock block) = text "close" <+> pretty lock <+> pretty block


{-  pretty (WhenThen lock th) = 
    hsep [text "when", parens (pretty lock) 
          ,text "then", pretty th
         ]
  pretty (WhenThenElse lock th el) =
    hsep [text "when", parens (pretty lock)
          , text "then", pretty th
          , text "else", pretty el
         ]
-}

instance Pretty Catch where
  pretty (Catch fParam block) =
    hsep [text "catch", parens (pretty fParam)] $$ pretty block

instance Pretty SwitchBlock where
  pretty (SwitchBlock lbl stmts) =
    vcat (pretty lbl : map (nest 2 . pretty) stmts)

instance Pretty SwitchLabel where
  pretty (SwitchCase e) = 
    text "case" <+> pretty e <> colon
  pretty Default = text "default:"

instance Pretty ForInit where
  pretty (ForLocalVars mods t vds) =
    hsep $ map pretty mods ++
            pretty t: punctuate comma (map pretty vds)


-----------------------------------------------------------------------
-- Expressions

instance Pretty Exp where
  pretty (Lit l) = pretty l
  
  pretty (ClassLit mT) = 
    ppResultType mT <> text ".class"

  pretty This = text "this"
  
  pretty (ThisClass name) = 
    pretty name <> text ".this"
    
  pretty (Paren e) = parens (pretty e)
  
  pretty (InstanceCreation tArgs ct args mBody) =
    hsep [text "new" 
          , ppTypeParams tArgs 
          , pretty ct <> ppArgs args
         ] $$ maybePP mBody
  
  pretty (QualInstanceCreation e tArgs ident args mBody) =
    hsep [pretty e <> char '.' <> text "new"
          , ppTypeParams tArgs
          , pretty ident <> ppArgs args
         ] $$ maybePP mBody

  pretty (ArrayCreate t es k) =
    text "new" <+> 
      hcat (pretty t : map (brackets . pretty) es 
                ++ replicate k (text "[]"))
  
  pretty (ArrayCreateInit t k init) =
    text "new" 
      <+> hcat (pretty t : replicate k (text "[]")) 
      <+> pretty init

  pretty (FieldAccess fa) = pretty fa
  
  pretty (MethodInv mi) = pretty mi
  
  pretty (ArrayAccess ain) = pretty ain

  pretty (ExpName name) = pretty name
  
  pretty (PostIncrement e) = pretty e <> text "++"
  
  pretty (PostDecrement e) = pretty e <> text "--"

  pretty (PreIncrement e)  = text "++" <> pretty e
  
  pretty (PreDecrement e)  = text "--" <> pretty e

  pretty (PrePlus e) = char '+' <> pretty e
  
  pretty (PreMinus e) = char '-' <> pretty e
  
  pretty (PreBitCompl e) = char '~' <> pretty e
  
  pretty (PreNot e) = char '!' <> pretty e

  pretty (Cast t e) = parens (pretty t) <+> pretty e
  
  pretty (BinOp e1 op e2) =
    hsep [pretty e1, pretty op, pretty e2]
  
  pretty (InstanceOf e rt) =
    hsep [pretty e, text "instanceof", pretty rt]
    
  pretty (Cond c th el) =
    hsep [pretty c, char '?', pretty th, colon, pretty el]

  pretty (Assign lhs aop e) =
    hsep [pretty lhs, pretty aop, pretty e]

-- Paragon
  pretty (PolicyExp pl) = pretty pl

--  pretty (PolicyOf i) = text "policyof" <> parens (pretty i)

  pretty (LockExp l) = char '?' <> pretty l


instance Pretty Literal where
  pretty (Int i) = text (show i)
  pretty (Word i) = text (show i) <> char 'L'
  pretty (Float f) = text (show f) <> char 'F'
  pretty (Double d) = text (show d)
  pretty (Boolean b) = text . map toLower $ show b
  pretty (Char c) = text (show c)
  pretty (String s) = text (show s)
  pretty (Null) = text "null"

instance Pretty Op where
  pretty op = text $ case op of
    Mult    -> "*"
    Div     -> "/"
    Rem     -> "%"
    Add     -> "+"
    Sub     -> "-"
    LShift  -> "<<"
    RShift  -> ">>"
    RRShift -> ">>>"
    LThan   -> "<"
    GThan   -> ">"
    LThanE  -> "<="
    GThanE  -> ">="
    Equal   -> "=="
    NotEq   -> "!="
    And     -> "&"
    Xor     -> "^"
    Or      -> "|"
    CAnd    -> "&&"
    COr     -> "||"
    
instance Pretty AssignOp where
  pretty aop = text $ case aop of
    EqualA  -> "="
    MultA   -> "*="
    DivA    -> "/="
    RemA    -> "%="
    AddA    -> "+="
    SubA    -> "-="
    LShiftA -> "<<="
    RShiftA -> ">>="
    RRShiftA -> ">>>="
    AndA    -> "&="
    XorA    -> "^="
    OrA     -> "|="

instance Pretty Lhs where
  pretty (NameLhs name) = pretty name
  pretty (FieldLhs fa) = pretty fa
  pretty (ArrayLhs ain) = pretty ain

instance Pretty ArrayIndex where
  pretty (ArrayIndex ref e) = pretty ref <> brackets (pretty e)

instance Pretty FieldAccess where
  pretty (PrimaryFieldAccess e ident) =
    pretty e <> char '.' <> pretty ident
  pretty (SuperFieldAccess ident) =
    text "super." <> pretty ident
  pretty (ClassFieldAccess name ident) =
    pretty name <> text ".super." <> pretty ident

instance Pretty MethodInvocation where
  pretty (MethodCall name args) =
    pretty name <> ppArgs args

  pretty (PrimaryMethodCall e tArgs ident args) =
    hcat [pretty e, char '.', ppTypeParams tArgs, 
           pretty ident, ppArgs args]

  pretty (SuperMethodCall tArgs ident args) =
    hcat [text "super.", ppTypeParams tArgs,
           pretty ident, ppArgs args]

  pretty (ClassMethodCall name tArgs ident args) =
    hcat [pretty name, text ".super.", ppTypeParams tArgs,
           pretty ident, ppArgs args]
  
  pretty (TypeMethodCall name tArgs ident args) =
    hcat [pretty name, char '.', ppTypeParams tArgs,
           pretty ident, ppArgs args]

instance Pretty ArrayInit where
  pretty (ArrayInit vInits) =
    braces $ hsep (punctuate comma (map pretty vInits))


ppArgs :: Pretty a => [a] -> Doc
ppArgs = parens . hsep . punctuate comma . map pretty

-----------------------------------------------------------------------
-- Types

instance Pretty Type where
  pretty (PrimType pt) = pretty pt
  pretty (RefType  rt) = pretty rt

instance Pretty RefType where
  pretty (ClassRefType ct) = pretty ct
  pretty (TypeVariable i)  = pretty i
  pretty (ArrayType t mp) = 
      pretty t <> text "[]" <> maybe empty (ppTypeParams . return) mp

instance Pretty ClassType where
  pretty (ClassType itas) =
    hcat . punctuate (char '.') $
      map (\(i,tas) -> pretty i <> ppTypeParams tas) itas

instance Pretty TypeArgument where
  pretty (ActualArg aa) = pretty aa
  pretty (Wildcard mBound) = char '?' <+> maybePP mBound

instance Pretty NonWildTypeArgument where
  pretty (ActualType   t) = pretty t
  pretty (ActualPolicy p) = text "policy" <+> pretty p
  pretty (ActualActor n) = text "actor" <+> pretty n
  pretty (ActualLockState ls) = text "lock[]" <+> ppArgs ls -- HACK ALERT

instance Pretty WildcardBound where
  pretty (ExtendsBound rt) = text "extends" <+> pretty rt
  pretty (SuperBound   rt) = text "super"   <+> pretty rt

instance Pretty PrimType where
  pretty BooleanT = text "boolean"
  pretty ByteT    = text "byte"
  pretty ShortT   = text "short"
  pretty IntT     = text "int"
  pretty LongT    = text "long"
  pretty CharT    = text "char"
  pretty FloatT   = text "float"
  pretty DoubleT  = text "double"

  -- Paragon
  pretty ActorT   = text "actor"
  pretty PolicyT  = text "policy"

instance Pretty TypeParam where
  pretty (TypeParam ident rts) =
    pretty ident 
      <+> opt (not $ null rts) 
           (hsep $ text "extends": 
                    punctuate (text " &") (map pretty rts))
  pretty (ActorParam ident) = text "actor" <+> pretty ident
  pretty (PolicyParam ident) = text "policy" <+> pretty ident
  pretty (LockStateParam ident) = text "lock[]" <+> pretty ident

ppTypeParams :: Pretty a => [a] -> Doc
ppTypeParams [] = empty
ppTypeParams tps = char '<' 
    <> hsep (punctuate comma (map pretty tps))
    <> char '>'

ppImplements :: [RefType] -> Doc
ppImplements [] = empty
ppImplements rts = text "implements" 
    <+> hsep (punctuate comma (map pretty rts))

ppExtends :: [RefType] -> Doc
ppExtends [] = empty
ppExtends rts = text "extends" 
    <+> hsep (punctuate comma (map pretty rts))

ppThrows :: [ExceptionSpec] -> Doc
ppThrows [] = empty
ppThrows ess = text "throws" 
    <+> hsep (punctuate comma (map pretty ess))

instance Pretty ExceptionSpec where
  pretty (ExceptionSpec mods t) =
      hsep (map pretty mods) <+> pretty t

ppResultType :: Maybe Type -> Doc
ppResultType Nothing = text "void"
ppResultType (Just a) = pretty a

-----------------------------------------------------------------------
-- Paragon Policies

instance Pretty PolicyExp where
  pretty (PolicyLit cs) = braces $ hcat (punctuate (char ';') $ map pretty cs)
  pretty (PolicyOf i) = text "policyof" <> parens (pretty i)
  pretty (PolicyTypeVar i) = pretty i

instance Pretty a => Pretty (Clause a) where
  pretty (Clause h b) = pretty h <> 
    char ':' <+> hcat (punctuate (char ',') $ map pretty b)
  
instance Pretty Actor where
  pretty (Actor str) = pretty str
  pretty (Var str)   = char '\'' <> pretty str

instance Pretty ActorName where
  pretty (ActorName n) = pretty n
  pretty (ActorTypeVar i) = pretty i

instance Pretty Atom where
  pretty (Atom pr vs) = pretty pr <> 
    opt (not $ null vs) (parens (hcat (punctuate (char ',') $ map pretty vs)))

instance Pretty Lock where
  pretty (Lock i as) = pretty i <>
    opt (not $ null as) (parens (hcat (punctuate (char ',') $ map pretty as)))
  pretty (LockVar i) = pretty i

instance Pretty LockProperties where
  pretty (LockProperties cs) = braces $ hcat (punctuate (char ';') $ map pretty cs)

prettyLocks :: [Lock] -> Doc
--prettyLocks [l] = pretty l
prettyLocks ls  = parens . hsep . punctuate (char ',') $ map pretty ls

--instance Pretty LockExp where
--  pretty (LockExp ls) = parens $ hcat (map pretty ls)
--  pretty (LockVar v)  = pretty v

-----------------------------------------------------------------------
-- Names and identifiers

instance Pretty Name where
  pretty (Name is) = 
    hcat (punctuate (char '.') $ map pretty is)

instance Pretty Ident where
  pretty (Ident s) = text s

-----------------------------------------------------------------------
-- Help functionality

maybePP :: Pretty a => Maybe a -> Doc
maybePP = maybe empty pretty

opt :: Bool -> Doc -> Doc
opt x a = if x then a else empty

braceBlock :: [Doc] -> Doc
braceBlock xs = char '{'
    $+$ nest 2 (vcat xs)
    $+$ char '}'