{-# LANGUAGE OverloadedStrings #-}

module Windows12.Ast where

import Data.Text (Text)
import Data.Text.Prettyprint.Doc

data BinOp
  = Add
  | Sub
  | Mul
  | Div
  | Mod
  | Eq
  | Ne
  | Lt
  | Gt
  | Le
  | Ge
  | And
  | Or
  | BitAnd
  | BitOr
  | BitXor
  | ShiftL
  | ShiftR
  deriving (Show, Eq)

data UnOp
  = Neg
  | Not
  | BitNot
  | Deref
  | AddrOf
  deriving (Show, Eq)

data AssignOp
  = BaseAssign
  | AddAssign
  | SubAssign
  | MulAssign
  | DivAssign
  | ModAssign
  | BitAndAssign
  | BitOrAssign
  | BitXorAssign
  | ShiftLAssign
  | ShiftRAssign
  deriving (Show, Eq)

data Expr
  = Id Text
  | IntLit Int
  | UIntLit Word
  | FloatLit Double
  | StrLit Text
  | BoolLit Bool
  | CharLit Char
  | BinOp BinOp Expr Expr
  | UnOp UnOp Expr
  | Call Expr [Expr]
  | Index Expr Expr
  | Member Expr Expr
  | Cast Type Expr
  | Sizeof Type
  | StructInit Text [(Text, Expr)]
  deriving (Show, Eq)

data Stmt
  = Expr Expr
  | Return Expr
  | If Expr [Stmt] (Maybe [Stmt])
  | While Expr [Stmt]
  | Assign AssignOp Expr Expr
  | Block [Stmt]
  | Var Text (Maybe Type) (Maybe Expr)
  deriving (Show, Eq)

data Type
  = IntType
  | UIntType
  | FloatType
  | StrType
  | BoolType
  | CharType
  | PtrType Type
  | ArrayType Type
  | StructType Text
  | EnumType Text
  | VoidType
  deriving (Show, Eq)

data Bind = Bind {bindName :: Text, bindType :: Type}
  deriving (Show, Eq)

data TLStruct = Struct {structName :: Text, structFields :: [Bind]}
  deriving (Show, Eq)

data TLEnum = Enum {enumName :: Text, enumFields :: [Text]}
  deriving (Show, Eq)

data TLFunc = Func {funcName :: Text, funcArgs :: [Bind], funcRetType :: Type, funcBody :: [Stmt]}
  deriving (Show, Eq)

data TL = TLStruct TLStruct | TLEnum TLEnum | TLFunc TLFunc
  deriving (Show, Eq)

data Program = Program [TLStruct] [TLEnum] [TLFunc]
  deriving (Show, Eq)

-- Pretty printing
instance Pretty BinOp where
  pretty Add = "+"
  pretty Sub = "-"
  pretty Mul = "*"
  pretty Div = "/"
  pretty Mod = "%"
  pretty Eq = "=="
  pretty Ne = "!="
  pretty Lt = "<"
  pretty Gt = ">"
  pretty Le = "<="
  pretty Ge = ">="
  pretty And = "&&"
  pretty Or = "||"
  pretty BitAnd = "&"
  pretty BitOr = "|"
  pretty BitXor = "^"
  pretty ShiftL = "<<"
  pretty ShiftR = ">>"

instance Pretty UnOp where
  pretty Neg = "-"
  pretty Not = "!"
  pretty BitNot = "~"
  pretty Deref = "*"
  pretty AddrOf = "&"

instance Pretty AssignOp where
  pretty BaseAssign = "="
  pretty AddAssign = "+="
  pretty SubAssign = "-="
  pretty MulAssign = "*="
  pretty DivAssign = "/="
  pretty ModAssign = "%="
  pretty BitAndAssign = "&="
  pretty BitOrAssign = "|="
  pretty BitXorAssign = "^="
  pretty ShiftLAssign = "<<="
  pretty ShiftRAssign = ">>="

instance Pretty Expr where
  pretty (Id x) = pretty x
  pretty (IntLit x) = pretty x
  pretty (UIntLit x) = pretty x
  pretty (FloatLit x) = pretty x
  pretty (StrLit x) = dquotes (pretty x)
  pretty (BoolLit x) = pretty x
  pretty (CharLit x) = squotes (pretty x)
  pretty (BinOp op l r) = parens (pretty l <+> pretty op <+> pretty r)
  pretty (UnOp op e) = pretty op <> parens (pretty e)
  pretty (Call f args) = parens (pretty f) <> parens (hsep (punctuate comma (map pretty args)))
  pretty (Index arr idx) = parens (pretty arr) <> brackets (pretty idx)
  pretty (Member e m) = pretty e <> "." <> pretty m
  pretty (Cast t e) = parens (pretty t) <> parens (pretty e)
  pretty (Sizeof t) = "sizeof" <> parens (pretty t)
  pretty (StructInit s fields) = pretty s <+> lbrace <> line <> indent 4 (vsep (punctuate comma (map (\(n, e) -> pretty n <+> "=" <+> pretty e) fields))) <> line <> rbrace

instance Pretty Stmt where
  pretty (Expr e) = pretty e <> semi
  pretty (Return e) = "return" <+> pretty e <> semi
  pretty (If cond t f) = "if" <+> pretty cond <+> prettyBlock t <+> maybe "" (\f' -> "else" <+> prettyBlock f') f
  pretty (While cond body) = "while" <+> pretty cond <+> prettyBlock body
  pretty (Assign op l r) = pretty l <+> pretty op <+> pretty r <> semi
  pretty (Block stmts) = braces (vsep (map pretty stmts))
  pretty (Var n t e) = pretty n <+> maybe "" (\t' -> ":" <+> pretty t') t <+> maybe "" (\e' -> "=" <+> pretty e') e <> semi

instance Pretty Type where
  pretty IntType = "int"
  pretty UIntType = "uint"
  pretty FloatType = "float"
  pretty StrType = "str"
  pretty BoolType = "bool"
  pretty CharType = "char"
  pretty (PtrType t) = pretty t <> "*"
  pretty (ArrayType t) = pretty t <> "[]"
  pretty (StructType s) = pretty s
  pretty (EnumType e) = pretty e
  pretty VoidType = "void"

instance Pretty Bind where
  pretty (Bind n t) = pretty n <+> ":" <+> pretty t

instance Pretty TLStruct where
  pretty (Struct n fields) = "struct" <+> pretty n <+> prettyFields fields

instance Pretty TLEnum where
  pretty (Enum n fields) = "enum" <+> pretty n <+> prettyFields fields

instance Pretty TLFunc where
  pretty (Func n args ret body) =
    "fn " <> pretty n <> parens (hsep (punctuate comma (map pretty args))) <+> "->" <+> pretty ret <+> prettyBlock body

instance Pretty TL where
  pretty (TLStruct s) = pretty s
  pretty (TLEnum e) = pretty e
  pretty (TLFunc f) = pretty f

instance Pretty Program where
  pretty (Program structs enums funcs) = vsep (map pretty structs) <> line <> vsep (map pretty enums) <> line <> vsep (map pretty funcs)

prettyFields :: (Pretty a) => [a] -> Doc ann
prettyFields fields = lbrace <> line <> indent 4 (vsep (punctuate comma (map pretty fields))) <> line <> rbrace

prettyBlock :: (Pretty a) => [a] -> Doc ann
prettyBlock stmts = lbrace <> line <> indent 4 (vsep (map pretty stmts)) <> line <> rbrace