module AST(
    Name,
    Program(..), Declaration(..), Block(..), Type(..), Literal(..),
    BinaryOperator(..), UnaryOperator(..), Expression(..), Statement(..)) where

import Data.List

import PShow


type Name = String

data Program = Program [Declaration]
  deriving (Show)

data Declaration
    = DecFunction {typeOf :: Type
                  ,nameOf :: Name
                  ,argumentsOf :: [(Type, Name)]
                  ,bodyOf :: Block}
    | DecVariable {typeOf :: Type
                  ,nameOf :: Name
                  ,valueOf :: Maybe Expression}
    | DecTypedef {typeOf :: Type
                 ,nameOf :: Name}
  deriving (Show)

data Block = Block [Statement]
  deriving (Show)

data Type = TypeInt Int
          | TypeUInt Int
          | TypeFloat
          | TypeDouble
          | TypePtr Type
          | TypeName Name
  deriving (Show, Eq)

data Literal = LitInt Integer
             | LitString String
             | LitVar Name
             | LitCall Name [Expression]
  deriving (Show)

data BinaryOperator
    = Plus | Minus | Times | Divide | Modulo
    | Equal | Unequal | Greater | Less | GEqual | LEqual
    | BoolAnd | BoolOr
  deriving (Show, Eq)

data UnaryOperator
    = Negate | Not | Invert | Dereference | Address
  deriving (Show, Eq)

data Expression
    = ExLit Literal
    | ExBinOp BinaryOperator Expression Expression
    | ExUnOp UnaryOperator Expression
  deriving (Show)

data Statement
    = StEmpty
    | StBlock Block
    | StExpr Expression
    | StVarDeclaration Type Name (Maybe Expression)
    | StAssignment Name Expression
    | StIf Expression Statement Statement
    | StWhile Expression Statement
    | StReturn Expression
  deriving (Show)


indent :: Int -> String -> String
indent sz str = intercalate "\n" $ map (prefix++) $ lines str
    where prefix = replicate sz ' '


instance PShow Program where
    pshow (Program decls) = intercalate "\n" (map pshow decls)

instance PShow Declaration where
    pshow (DecFunction t n a b) =
        concat [pshow t, " ", n, "(", intercalate ", " (map pshowArg a), ") ", pshow b]
      where pshowArg (argt, argn) = concat [pshow argt, " ", argn]
    pshow (DecVariable t n Nothing) =
        concat [pshow t, " ", n, ";"]
    pshow (DecVariable t n (Just e)) =
        concat [pshow t, " ", n, " = ", pshow e, ";"]
    pshow (DecTypedef t n) =
        concat ["type ", n, " = ", pshow t, ";"]

instance PShow Block where
    pshow (Block []) = "{}"
    pshow (Block stmts) = concat ["{\n", indent 4 $ intercalate "\n" (map pshow stmts), "\n}"]

instance PShow Type where
    pshow (TypeInt sz) = 'i' : pshow sz
    pshow (TypeUInt sz) = 'u' : pshow sz
    pshow TypeFloat = "float"
    pshow TypeDouble = "double"
    pshow (TypePtr t) = concat ["ptr(", pshow t, ")"]
    pshow (TypeName n) = n

instance PShow Literal where
    pshow (LitInt i) = pshow i
    pshow (LitString s) = pshow s
    pshow (LitVar n) = n
    pshow (LitCall n a) = concat [n, "(", intercalate ", " (map pshow a), ")"]

instance PShow BinaryOperator where
    pshow Plus = "+"
    pshow Minus = "-"
    pshow Times = "*"
    pshow Divide = "/"
    pshow Modulo = "%"
    pshow Equal = "=="
    pshow Unequal = "!="
    pshow Greater = ">"
    pshow Less = "<"
    pshow GEqual = ">="
    pshow LEqual = "<="
    pshow BoolAnd = "&&"
    pshow BoolOr = "||"

instance PShow UnaryOperator where
    pshow Negate = "-"
    pshow Not = "!"
    pshow Invert = "~"
    pshow Dereference = "*"
    pshow Address = "&"

instance PShow Expression where
    pshow (ExLit lit) = pshow lit
    pshow (ExBinOp op a b) = concat [pshow a, " ", pshow op, " ", pshow b]
    pshow (ExUnOp op a) = concat [pshow op, pshow a]

instance PShow Statement where
    pshow StEmpty = ";"
    pshow (StBlock bl) = pshow bl
    pshow (StExpr e) = pshow e ++ ";"
    pshow (StVarDeclaration t n Nothing) = concat [pshow t, " ", n, ";"]
    pshow (StVarDeclaration t n (Just e)) = concat [pshow t, " ", n, " = ", pshow e, ";"]
    pshow (StAssignment n e) = concat [n, " = ", pshow e, ";"]
    pshow (StIf c t StEmpty) = concat ["if (", pshow c, ") ", pshow t]
    pshow (StIf c t@(StBlock _) e) = concat ["if (", pshow c, ") ", pshow t, " else ", pshow e]
    pshow (StIf c t e) = concat ["if (", pshow c, ") ", pshow t, "\nelse ", pshow e]
    pshow (StWhile c s) = concat ["while (", pshow c, ") ", pshow s]
    pshow (StReturn e) = concat ["return ", pshow e, ";"]