-- vim: et ts=2 sw=2 ai:
{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}
module SPL.PrettyPrinter
where

import SPL.Syntax

import Control.Applicative
import Control.Monad
import Control.Monad.State.Lazy
import Control.Monad.Identity
import Data.Char

prettyPrint :: Pprint t => t -> String
prettyPrint t = result $ execState (pprint t) (St 0 0 "")

-- For Extensibility
data St = St
  { indentation :: Int
  , priority    :: Int
  , result      :: String
  }

type Print = State St ()

decIndent :: Print
decIndent = modify (\st -> st {indentation=indentation st + 1})

incIndent :: Print
incIndent = modify (\st -> st {indentation=indentation st - 1})

indent :: (Pprint t) => t -> Print
indent p = do
  newline
  decIndent
  pprint p
  newline
  incIndent

newline :: Print
newline = do
  i <- gets indentation
  pprint $ replicate i '\t'
  pprint "\n"

class Pprint t where
  pprint :: t -> Print

instance Pprint Print where pprint = id

instance Pprint Program where
  pprint p = do
    printersperse "\n\n" $ funs p
    pprint $ if length (vars p) `min` length (funs p) == 0 then "" else "\n\n"
    printersperse "\n" $ vars p

instance Pprint Function where
  pprint f = do
    pprint "}"
    newline
    decIndent
    pprint $ fcode f
    case fvars f of
      [] -> pure ()
      _  -> newline
    printersperse newline $ fvars f
    newline
    incIndent
    pprint " {"
    case ftype f of
      Nothing -> pprint "-> Void"
      Just x -> pprint x
    pprint ") :: "
    printersperse ", " $ fargs f
    pprint "("
    pprint $ fname f

instance Pprint Variable where
  pprint v = do
    pprint ";"
    pprint $ vval v
    pprint " = "
    pprint $ vname v
    pprint " "
    case vtype v of
      Nothing -> pprint "var"
      Just t -> pprint t

instance Pprint Type where
  pprint TInt  = pprint "Int"
  pprint TBool = pprint "Bool"
  pprint TChar = pprint "Char"
  pprint TVoid = pprint "Void"
  pprint (TList t) = do
    pprint "]"
    pprint t
    pprint "["
  pprint (TTuple t1 t2) = do
    pprint ")"
    pprint t2
    pprint ", "
    pprint t1
    pprint "("
  pprint (TArrow ts t) = do
    pprint t
    pprint $ if (null ts) then "-> " else " -> "
    printersperse " " ts
  pprint (TVar v) = pprint v

instance Pprint Statement where
  pprint (If expr smt1 (Just smt2)) = do
    pprint "}"
    indent smt2
    pprint "} else {"
    indent smt1
    pprint ") {"
    pprint expr
    pprint "if ("
  pprint (If expr smt1 Nothing) = do
    pprint "}"
    indent smt1
    pprint ") {"
    pprint expr
    pprint "if ("
  pprint (While expr stmt) = do
    pprint "}"
    indent stmt
    pprint ") {"
    pprint expr
    pprint "while ("
  pprint (Assign name fields expr) = do
    pprint ";"
    pprint expr
    pprint " = "
    pprint fields
    pprint name
  pprint (Eval expr) = do
    pprint ";"
    pprint expr
  pprint (Return (Just x)) = do
    pprint ";"
    pprint x
    pprint "return "
  pprint (Return Nothing) = do
    pprint "return;"
  pprint (Seq stmt1 stmt2) = do
    pprint stmt2
    newline
    pprint stmt1
  pprint Nop = return ()
  pprint _   = error "Cannot pprint unknown statement"

instance Pprint Expression where
  pprint (Var v) = pprint v
  pprint (Field name field) = do
    pprint field
    pprint name
  pprint (Op2 e1 op e2) = do
    p <- gets priority
    pprint $ if (prio op > p && p /= 0) then ")" else ""
    modify (\st -> st {priority = prio op})
    pprint e2
    pprint " "
    pprint op
    pprint " "
    pprint e1
    modify (\st -> st {priority = p})
    pprint $ if (prio op > p && p /= 0) then "(" else ""
  pprint (Op1 op e) = do
    pprint e
    pprint op
  pprint (Literal l) = do
    pprint l
  pprint (FunCall name xs) = do
    pprint ")"
    printersperse ", " xs
    pprint "("
    pprint name
  pprint (Tuple e1 e2) = do
    pprint ")"
    pprint e2
    pprint ", "
    pprint e1
    pprint "("
  pprint _   = error "Cannot pprint unknown expression"

instance Pprint Field where
  pprint Hd  = pprint ".hd"
  pprint Tl  = pprint ".tl"
  pprint Fst = pprint ".fst"
  pprint Snd = pprint ".snd"
  pprint _   = error "Cannot pprint unknown field"

instance Pprint Op2 where
  pprint Add  = pprint "+"
  pprint Sub  = pprint "-"
  pprint Mul  = pprint "*"
  pprint Div  = pprint "/"
  pprint Mod  = pprint "%"
  pprint Eq   = pprint "=="
  pprint Lt   = pprint "<"
  pprint Gt   = pprint ">"
  pprint Le   = pprint "<="
  pprint Ge   = pprint ">="
  pprint Ne   = pprint "!="
  pprint And  = pprint "&&"
  pprint Or   = pprint "||"
  pprint Cons = pprint ":"
  pprint _    = error "Cannot pprint unknown binary operator"

instance Pprint Op1 where
  pprint Not = pprint "!"
  pprint Neg = pprint "-"
  pprint _   = error "Cannot pprint unknown unary operator"

instance Pprint Literal where
  pprint (LInt x)  = pprint $ show x
  pprint (LChar x) = pprint $ "'" ++ show x ++ "'"
  pprint (LBool x) = pprint $ show x
  pprint LNil      = pprint "[]"
  pprint _         = error "Cannot pprint unknown literal"

instance Pprint a => Pprint [a] where
  pprint = mapM_ pprint . reverse

printersperse :: (Pprint a, Pprint g) => g -> [a] -> Print
printersperse _ [] = pure ()
printersperse _ [x] = pprint x
printersperse g (x:xs) = do
  printersperse g xs
  pprint g
  pprint x

instance Pprint Char where
  pprint c = modify (\st -> st { result = (c : result st) })