module AST where

import Data.List
import Data.Maybe

import Defs
import Pretty
import Utils


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

data DTypedef = DTypedef Name Type
  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) | TStruct [(Type, Name)]
    | TName Name
  deriving (Show, Eq, Ord)

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

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

data AsExpression
    = AEVar Name (Maybe Type)
    | AESubscript AsExpression Expression (Maybe Type)
    | AEGet AsExpression Name (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)
    | EGet Expression Name (Maybe Type)
    | ECast Type Expression
    | ENew Type Expression
  deriving (Show, Eq)

data BinaryOp
    = BOAdd | BOSub | BOMul | BODiv | BOMod | BOPow
    | BOAnd | BOOr
    | BOBitAnd | BOBitOr | BOBitXor
    | 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]
    | LStr String
    | LStruct [(Name, Expression)]
  deriving (Show, Eq)


sizeof :: Type -> Size
sizeof TInt = 8
sizeof TChar = 1
sizeof (TArr _ _) = 8
sizeof (TStruct []) = 0
sizeof st@(TStruct ms) =
    roundUp (sum (layoutStruct st) + sizeof (fst $ last ms)) (foldl1 lcm $ 8 : map (alignmentof . fst) ms)
sizeof t@(TName _) = error $ "sizeof on " ++ show t

alignmentof :: Type -> Offset
alignmentof TInt = 8
alignmentof TChar = 1
alignmentof (TArr _ _) = 8
alignmentof (TStruct []) = 1
alignmentof (TStruct ((t,_):_)) = alignmentof t
alignmentof (TName _) = undefined

layoutStruct :: Type -> [Offset]
layoutStruct (TStruct ms) = go ms 0
  where
    go :: [(Type, Name)] -> Offset -> [Offset]
    go [] _ = []
    go ((t,_) : rest) start =
        let o = roundUp start (alignmentof t)
        in o : go rest (o + sizeof t)
layoutStruct _ = undefined

offsetInStruct :: Type -> Name -> Offset
offsetInStruct st@(TStruct ms) name = layoutStruct st !! (fromJust $ findIndex ((==name) . snd) ms)
offsetInStruct _ _ = undefined

structMemberType :: Type -> Name -> Type
structMemberType (TStruct ms) name = fst $ fromJust $ find ((==name) . snd) ms
structMemberType _ _ = undefined


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

instance Pretty DTypedef where
    prettyI i (DTypedef n t) =
        "type " ++ n ++ " := " ++ prettyI i t ++ ";"

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 ++ "]"
    prettyI _ (TStruct []) = "struct {}"
    prettyI _ (TStruct ms) =
        "struct {" ++ intercalate " " [pretty t ++ " " ++ n ++ ";" | (t,n) <- ms] ++ "}"
    prettyI _ (TName n) = n

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 b (Block [])) =
        "if (" ++ prettyI i c ++ ") " ++ prettyI i b
    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 _ (SBreak 0) =
        "break;"
    prettyI _ (SBreak n) =
        "break " ++ show n ++ ";"
    prettyI _ (SReturn Nothing) =
        "return;"
    prettyI i (SReturn (Just e)) =
        "return " ++ prettyI i e ++ ";"
    prettyI i (SExpr e) = prettyI i e ++ ";"
    prettyI _ SDebugger = "debugger;"

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 (EGet e n Nothing) =
        "(" ++ prettyI i e ++ ")." ++ n
    prettyI i (EGet e n (Just t)) =
        "(" ++ prettyI i (EGet e n Nothing) ++ " :: " ++ prettyI i t ++ ")"
    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 ++ "]"
    prettyI i (AEGet ae n (Just t)) =
        "(" ++ prettyI i (AEGet ae n Nothing) ++ " :: " ++ prettyI i t ++ ")"
    prettyI i (AEGet ae n Nothing) =
        prettyI i ae ++ "." ++ n

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 = "<="
    prettyI _ BOBitAnd = "&"
    prettyI _ BOBitOr = "|"
    prettyI _ BOBitXor = "^"

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) ++ ")"
    prettyI _ (LStr s) = show s
    prettyI i (LStruct ms) =
        "{" ++ intercalate ", " (map (\(n,e) -> n ++ " = " ++ prettyI i e) ms) ++ "}"