module AST where

import Data.List

import Defs
import Pretty


data Program = Program [DVar] [DFunc]
  deriving (Show, Eq)

data DVar = DVar Type Name Expression
  deriving (Show, Eq)

data DFunc = DFunc (Maybe Type) Name [(Type, Name)] Block
  deriving (Show, Eq)

data Type
    = TInt | TChar | TArr Type (Maybe Size)
  deriving (Show, Eq)

data Block = Block [Statement]
  deriving (Show, Eq)

data Statement
    = SDecl Type Name Expression
    | SAs AsExpression Expression
    | SIf Expression Block Block
    | SWhile Expression Block
    | SReturn (Maybe Expression)
    | SExpr Expression
  deriving (Show, Eq)

data AsExpression
    = AEVar Name (Maybe Type)
    | AESubscript AsExpression Expression (Maybe Type)
  deriving (Show, Eq)

data Expression
    = EBin BinaryOp Expression Expression (Maybe Type)
    | EUn UnaryOp Expression (Maybe Type)
    | ELit Literal (Maybe Type)
    | ESubscript Expression Expression (Maybe Type)
    | ECast Type Expression
    | ENew Type Expression
  deriving (Show, Eq)

data BinaryOp
    = BOAdd | BOSub | BOMul | BODiv | BOMod | BOPow
    | BOAnd | BOOr
    | BOEq | BONeq | BOGt | BOLt | BOGeq | BOLeq
  deriving (Show, Eq)

data UnaryOp
    = UONot | UONeg
  deriving (Show, Eq)

data Literal
    = LInt Integer
    | LChar Char
    | LVar Name
    | LCall Name [Expression]
  deriving (Show, Eq)


sizeof :: Type -> Size
sizeof TInt = 8
sizeof TChar = 1
sizeof (TArr _ _) = 8


instance Pretty Program where
    prettyI i (Program vars funcs) =
        concatMap (++ ("\n" ++ indent i)) $
            map (prettyI i) vars ++ map (prettyI i) funcs
      where
        indent n = replicate (2*n) ' '

instance Pretty DVar where
    prettyI i (DVar t n e) =
        prettyI i t ++ " " ++ n ++ " := " ++ prettyI i e ++ ";"

instance Pretty DFunc where
    prettyI i (DFunc mt n al b) =
        "func" ++ maybe "" ((' ' :) . prettyI i) mt ++ " " ++ n ++ "(" ++
            intercalate ","
                (map (\(at,an) -> prettyI i at ++ " " ++ an) al) ++
            ") " ++ prettyI i b
 
instance Pretty Type where
    prettyI _ TInt = "int"
    prettyI _ TChar = "char"
    prettyI _ (TArr t Nothing) = pretty t ++ "[]"
    prettyI _ (TArr t (Just sz)) = pretty t ++ "[" ++ show sz ++ "]"

instance Pretty Block where
    prettyI _ (Block []) = "{}"
    prettyI i (Block l) =
        "{" ++
            concatMap (("\n" ++ indent (i+1)) ++) (map (prettyI (i+1)) l) ++
            "\n" ++ indent i ++ "}"
      where
        indent n = replicate (2*n) ' '

instance Pretty Statement where
    prettyI i (SDecl t n e) =
        prettyI i t ++ " " ++ n ++ " := " ++ prettyI i e ++ ";"
    prettyI i (SAs target e) =
        prettyI i target ++ " = " ++ prettyI i e ++ ";"
    prettyI i (SIf c b1 b2) =
        "if " ++ prettyI i c ++ " " ++ prettyI i b1 ++ " else " ++ prettyI i b2
    prettyI i (SWhile c b) =
        "while " ++ prettyI i c ++ " " ++ prettyI i b
    prettyI _ (SReturn Nothing) =
        "return;"
    prettyI i (SReturn (Just e)) =
        "return " ++ prettyI i e ++ ";"
    prettyI i (SExpr e) = prettyI i e ++ ";"

instance Pretty Expression where
    prettyI i (EBin bo a b (Just t)) =
        "(" ++ prettyI i (EBin bo a b Nothing) ++ " :: " ++ prettyI i t ++ ")"
    prettyI i (EBin bo a b Nothing) =
        "(" ++ prettyI i a ++ ") " ++ prettyI i bo ++
            " (" ++ prettyI i b ++ ")"
    prettyI i (EUn uo e (Just t)) =
        "(" ++ prettyI i (EUn uo e Nothing) ++ " :: " ++ prettyI i t ++ ")"
    prettyI i (EUn uo e Nothing) =
        prettyI i uo ++ "(" ++ prettyI i e ++ ")"
    prettyI i (ELit l (Just t)) =
        "(" ++ prettyI i (ELit l Nothing) ++ " :: " ++ prettyI i t ++ ")"
    prettyI i (ELit l Nothing) = prettyI i l
    prettyI i (ESubscript a b (Just t)) =
        "(" ++ prettyI i (ESubscript a b Nothing) ++ " :: " ++ prettyI i t ++ ")"
    prettyI i (ESubscript a b Nothing) = "(" ++ prettyI i a ++ ")[" ++ prettyI i b ++ "]"
    prettyI i (ECast t e) = prettyI i t ++ "(" ++ prettyI i e ++ ")"
    prettyI i (ENew t e) = "new " ++ prettyI i t ++ "[" ++ prettyI i e ++ "]"

instance Pretty AsExpression where
    prettyI i (AEVar n (Just t)) =
        "(" ++ prettyI i (AEVar n Nothing) ++ " :: " ++ prettyI i t ++ ")"
    prettyI _ (AEVar n Nothing) = n
    prettyI i (AESubscript ae e (Just t)) =
        "(" ++ prettyI i (AESubscript ae e Nothing) ++ " :: " ++ prettyI i t ++ ")"
    prettyI i (AESubscript ae e Nothing) = prettyI i ae ++ "[" ++ prettyI i e ++ "]"

instance Pretty BinaryOp where
    prettyI _ BOAdd = "+"
    prettyI _ BOSub = "-"
    prettyI _ BOMul = "*"
    prettyI _ BODiv = "/"
    prettyI _ BOPow = "**"
    prettyI _ BOMod = "%"
    prettyI _ BOAnd = "&&"
    prettyI _ BOOr = "||"
    prettyI _ BOEq = "=="
    prettyI _ BONeq = "!="
    prettyI _ BOGt = ">"
    prettyI _ BOLt = "<"
    prettyI _ BOGeq = ">="
    prettyI _ BOLeq = "<="


instance Pretty UnaryOp where
    prettyI _ UONot = "!"
    prettyI _ UONeg = "-"

instance Pretty Literal where
    prettyI _ (LInt n) = show n
    prettyI _ (LChar c) = show c
    prettyI _ (LVar n) = n
    prettyI i (LCall n al) =
        n ++ "(" ++ intercalate ", " (map (prettyI i) al) ++ ")"