module AST where

import Data.List
import qualified GCStore as GCS (Id)


type Name = String

data Program = Program Block
  deriving (Show)

data Block = Block [Statement]
  deriving (Show)

data Statement
    = Declaration Name Expression
    | Assignment Name Expression
    | Condition Expression Block Block
    | Dive Name [Expression] Block
    | Expr Expression
  deriving (Show)

data Expression
    = EBin BinaryOp Expression Expression
    | EUn UnaryOp Expression
    | ELit Literal
  deriving (Show)

data Literal
    = LNum Double
    | LStr String
    | LVar Name
    | LBlock BlockType ArgList Block
    | LGCSId GCS.Id
    | LNil
  deriving (Show)

type ArgList = [Name]

data BlockType = BT0 | BT1 | BT2  -- the number of ?'s
  deriving (Show, Enum, Eq)

data BinaryOp
    = BOPlus | BOMinus | BOMul | BODiv | BOMod | BOPow
    | BOLess | BOGreater | BOEqual | BOLEq | BOGEq
    | BOBoolAnd | BOBoolOr
  deriving (Show, Eq)

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


class ASTPretty a where
    astPrettyI :: Int -> a -> String

    astPretty :: a -> String
    astPretty = astPrettyI 0


indent :: Int -> String
indent n = replicate (4*n) ' '

instance ASTPretty Program where
    astPrettyI i (Program (Block sts)) =
        let pr = map (astPrettyI i) sts
        in intercalate "\n" $ map (uncurry (++)) $ zip ("" : cycle [indent i]) pr

instance ASTPretty Block where
    astPrettyI _ (Block []) = "{}"
    astPrettyI i (Block sts) =
        let lns = map (('\n' : indent (i+1)) ++) $ map (astPrettyI (i+1)) sts
        in "{" ++ concat lns ++ "\n" ++ indent i ++ "}"

instance ASTPretty Statement where
    astPrettyI i (Declaration n e) = n ++ " := " ++ astPrettyI i e ++ ";"
    astPrettyI i (Assignment n e) = n ++ " = " ++ astPrettyI i e ++ ";"
    astPrettyI i (Condition c b1 b2) =
        "if " ++ astPrettyI i c ++ " " ++ astPrettyI i b1 ++ " else " ++ astPrettyI i b2
    astPrettyI i (Dive n [] b) = n ++ " " ++ astPrettyI i b
    astPrettyI i (Dive n al b) =
        n ++ "(" ++ intercalate ", " (map (astPrettyI i) al) ++ ") " ++ astPrettyI i b
    astPrettyI i (Expr e) = astPrettyI i e ++ ";"

instance ASTPretty Expression where
    astPrettyI i (EBin bo e1 e2) =
        "(" ++ astPrettyI i e1 ++ ") " ++ astPrettyI i bo ++ " (" ++ astPrettyI i e2 ++ ")"
    astPrettyI i (EUn uo e) =
        astPrettyI i uo ++ "(" ++ astPrettyI i e ++ ")"
    astPrettyI i (ELit l) = astPrettyI i l

instance ASTPretty Literal where
    astPrettyI _ (LNum m) = show m
    astPrettyI _ (LStr s) = show s
    astPrettyI _ (LVar n) = n
    astPrettyI i (LBlock bt [] b) = replicate (fromEnum bt) '?' ++ astPrettyI i b
    astPrettyI i (LBlock bt al b) =
        replicate (fromEnum bt) '?' ++ "(" ++ intercalate ", " al ++ ")" ++ astPrettyI i b
    astPrettyI _ (LGCSId d) = "<[" ++ show d ++ "]>"
    astPrettyI _ LNil = "nil"

instance ASTPretty BinaryOp where
    astPrettyI _ bo = case bo of
        BOPlus -> "+"
        BOMinus -> "-"
        BOMul -> "*"
        BODiv -> "/"
        BOMod -> "%"
        BOPow -> "**"
        BOLess -> "<"
        BOGreater -> ">"
        BOEqual -> "=="
        BOLEq -> "<="
        BOGEq -> ">="
        BOBoolAnd -> "&&"
        BOBoolOr -> "||"

instance ASTPretty UnaryOp where
    astPrettyI _ uo = case uo of
        UONeg -> "-"
        UONot -> "!"