module AST 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}
    | DecExtern {typeOf :: Type
                ,nameOf :: Name}
  deriving (Show)

data Block = Block [Statement]
  deriving (Show)

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

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

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

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

data Expression    -- (Maybe Type)'s are type annotations by the type checker
    = ExLit Literal (Maybe Type)
    | ExCast Type Expression  -- No type annotation needed
    | ExBinOp BinaryOperator Expression Expression (Maybe Type)
    | ExUnOp UnaryOperator Expression (Maybe Type)
  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 (Maybe Expression)
  deriving (Show)


exLit_ :: Literal -> Expression
exLit_ l = ExLit l Nothing

exBinOp_ :: BinaryOperator -> Expression -> Expression -> Expression
exBinOp_ bo a b = ExBinOp bo a b Nothing

exUnOp_ :: UnaryOperator -> Expression -> Expression
exUnOp_ uo e = ExUnOp uo e Nothing

exTypeOf :: Expression -> Maybe Type
exTypeOf (ExLit _ mt) = mt
exTypeOf (ExCast t _) = Just t
exTypeOf (ExBinOp _ _ _ mt) = mt
exTypeOf (ExUnOp _ _ mt) = mt

-- exSetType :: Type -> Expression -> Expression
-- exSetType t (ExLit l _) = ExLit l t
-- exSetType t (ExBinOp bo e1 e2 _) = ExBinOp bo e1 e2 t
-- exSetType t (ExUnOp bo e _) = ExUnOp bo e t


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, ";"]
    pshow (DecExtern t n) =
        concat ["extern ", pshow t, " ", n, ";"]

instance PShow Block where
    pshow (Block []) = "{}"
    pshow (Block stmts) = concat ["{\n", indent 4 $ intercalate "\n" (map pshow stmts), "\n}"]
      where
        indent :: Int -> String -> String
        indent sz str = intercalate "\n" $ map (prefix++) $ lines str
            where prefix = replicate sz ' '

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
    pshow TypeVoid = "void"
    pshow (TypeFunc ret args) = concat ["func ", pshow ret, "("] ++ intercalate "," (map pshow args) ++ ")"

instance PShow Literal where
    pshow (LitInt i) = pshow i
    pshow (LitUInt i) = pshow i ++ "U"
    pshow (LitFloat x) = pshow x
    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 = "||"
    pshow Index = "[~]"

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

instance PShow Expression where
    pshow (ExLit lit Nothing) = pshow lit
    pshow (ExLit lit (Just t)) = concat ["(", pshow lit, " :: ", pshow t, ")"]
    pshow (ExCast t ex) = concat ["cast(", pshow t, ")(", pshow ex, ")"]
    pshow (ExBinOp Index a b Nothing) = concat ["(", pshow a, ")[", pshow b, "]"]
    pshow (ExBinOp Index a b (Just t)) = concat ["((", pshow a, ")[", pshow b, "] :: ", pshow t, ")"]
    pshow (ExBinOp op a b Nothing) = concat ["(", pshow a, " ", pshow op, " ", pshow b, ")"]
    pshow (ExBinOp op a b (Just t)) = concat ["(", pshow a, " ", pshow op, " ", pshow b, " :: ", pshow t, ")"]
    pshow (ExUnOp op a Nothing) = concat [pshow op, pshow a]
    pshow (ExUnOp op a (Just t)) = concat ["(", pshow op, pshow a, " :: ", pshow t, ")"]

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 Nothing) = "return;"
    pshow (StReturn (Just e)) = concat ["return ", pshow e, ";"]