open Ast

let tab = "    "

(* string -> string *)
let indent = Str.global_replace (Str.regexp "^\\(.\\)") (tab ^ "\\1")

(* monop -> string *)
let monop2str = function
    | Neg -> "-"
    | Not -> "!"

(* binop -> string *)
let binop2str = function
    | Add -> " + "
    | Sub -> " - "
    | Mul -> " * "
    | Div -> " / "
    | Mod -> " % "
    | Eq  -> " == "
    | Ne  -> " != "
    | Lt  -> " < "
    | Le  -> " <= "
    | Gt  -> " > "
    | Ge  -> " >= "
    | And -> " && "
    | Or  -> " || "

(* ctype -> string *)
let rec type2str = function
    | Void -> "void"
    | Bool -> "bool"
    | Int -> "int"
    | Float -> "float"
    | ArrayDec (t, dims)
    | ArrayDef (t, dims) -> (type2str t) ^ "[" ^ (concat ", " dims) ^ "]"

and concat sep nodes = String.concat sep (List.map node2str nodes)

(* node -> string *)
and node2str node =
    let str = node2str in
    match node with

    (* Global *)
    | Program (decls, _) ->
        concat "\n\n" decls
    | Param (param_type, name, _) ->
        (type2str param_type) ^ " " ^ name
    | Dim (name, _) -> name
    | FunDec (ret_type, name, params, _) ->
        let params = concat ", " params in
        "extern " ^ type2str ret_type ^ " " ^ name ^ "(" ^ params ^ ");"
    | FunDef (export, ret_type, name, params, body, _) ->
        let export = if export then "export " else "" in
        let params = "(" ^ (concat ", " params) ^ ")" in
        export ^ type2str ret_type ^ " " ^ name ^ params ^ " " ^ str body
    | GlobalDec (var_type, name, _) ->
        "extern " ^ type2str var_type ^ " " ^ name ^ ";"
    | GlobalDef (export, ret_type, name, init, _) ->
        let export = if export then "export " else "" in
        let init = match init with
            | Some value -> " = " ^ str value
            | None -> ""
        in
        export ^ (type2str ret_type) ^ " " ^ name ^ init ^ ";"

    (* Statements *)
    | VarDec (var_type, name, None, _) ->
        (type2str var_type) ^ " " ^ name ^ ";"
    | VarDec (var_type, name, Some init, _) ->
        (type2str var_type) ^ " " ^ name ^ " = " ^ (str init) ^ ";"
    | Assign (name, value, _) ->
        name ^ " = " ^ (str value) ^ ";"
    | Expr expr ->
        str expr ^ ";"
    | Return (value, _) ->
        "return " ^ (str value) ^ ";"
    | If (cond, body, _) ->
        "if (" ^ str cond ^ ") " ^ str body
    | IfElse (cond, true_body, false_body, _) ->
        "if (" ^ str cond ^ ") " ^ str true_body ^ " else " ^ str false_body
    | While (cond, body, _) ->
        "while (" ^ str cond ^ ") " ^ str body
    | DoWhile (cond, body, _) ->
        "do " ^ str body ^ " while (" ^ str cond ^ ");"
    | For (counter, start, stop, step, body, _) ->
        let step = match step with
            | IntConst (1, _) -> ""
            | value -> ", " ^ str value
        in
        let range = str start ^ ", " ^ str stop ^ step in
        "for (int " ^ counter ^ " = " ^ range ^ ") " ^ str body
    | Allocate (name, dims, _) ->
        name ^ " = __allocate(" ^ concat ", " dims ^ ");"
    | Block body -> "{\n" ^ indent (concat "\n" body) ^ "\n}"

    (* Expressions *)
    | BoolConst (b, _) -> string_of_bool b
    | IntConst (i, _) -> string_of_int i
    | FloatConst (f, _) -> string_of_float f
    | ArrayConst (dims, _) -> "[" ^ concat ", " dims ^ "]"
    | ArrayScalar (value, _) -> str value
    | Var (v, _) -> v
    | Deref (name, dims, _) -> name ^ (str (ArrayConst (dims, noloc)))
    | Monop (op, opnd, _) -> monop2str op ^ str opnd
    | Binop (op, left, right, _) -> "(" ^ str left ^ binop2str op ^ str right ^ ")"
    | Cond (cond, t, f, _) -> (str cond) ^ " ? " ^ str t ^ " : " ^ str f
    | TypeCast (ctype, value, _) -> "(" ^ type2str ctype ^ ")" ^ str value
    | FunCall (name, args, _) -> name ^ "(" ^ (concat ", " args) ^ ")"

    | Arg node
    | VarUse (node, _, _)
    | FunUse (node, _, _) -> str node

    | _ -> raise InvalidNode