open Ast

let tab = "    "

(* int -> string list -> string list *)
let indent (lines : string list) =
    let prepend_tab line = tab ^ line in
    List.map prepend_tab lines

(* 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) -> (type2str t) ^ "[" ^ (String.concat ", " dims) ^ "]"
    | ArrayDef (t, dims) -> (type2str t) ^ "[" ^ (String.concat ", " (List.map node2str dims)) ^ "]"

(* decl -> string list *)
and node2lines node =
    let all_str = List.map node2str in
    let all_lines = List.map node2lines in
    match node with
    (* Decls *)
    | FunDec (ret_type, name, params, _) ->
        let params = String.concat ", " (all_str 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 = String.concat ", " (all_str params) in
        let header = type2str ret_type ^ " " ^ name ^ "(" ^ params ^ ")" in
        let body = indent (List.concat (all_lines body)) in
        [export ^ header ^ " {"] @
            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 -> " = " ^ node2str 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 ^ " = " ^ (node2str init) ^ ";"]
    | Assign (name, value, _) ->
        [name ^ " = " ^ (node2str value) ^ ";"]
    | Expr expr ->
        [node2str expr ^ ";"]
    | Return (value, _) ->
        ["return " ^ (node2str value) ^ ";"]
    | If (cond, body, _) ->
        let body = indent (List.concat (all_lines body)) in
        ["if (" ^ node2str cond ^ ") {"] @
            body @
        ["}"]
    | IfElse (cond, true_body, false_body, _) ->
        let true_body = indent (List.concat (all_lines true_body)) in
        let false_body = indent (List.concat (all_lines false_body)) in
        ["if (" ^ node2str cond ^ ") {"] @
            true_body @
        ["} else {"] @
            false_body @
        ["}"]
    | While (cond, body, _) ->
        let body = indent (List.concat (all_lines body)) in
        ["while (" ^ node2str cond ^ ") {"] @
            body @
        ["}"]
    | DoWhile (cond, body, _) ->
        let body = indent (List.concat (all_lines body)) in
        ["do {"] @
            body @
        ["} while (" ^ node2str cond ^ ");"]
    | For (counter, start, stop, step, body, _) ->
        let step = match step with
            | IntConst (1, _) -> ""
            | value -> ", " ^ node2str value
        in
        let range = node2str start ^ ", " ^ node2str stop ^ step in
        let body = indent (List.concat (all_lines body)) in
        ["for (int " ^ counter ^ " = " ^ range ^ ") {"] @
            body @
        ["}"]
    | Allocate (name, dims, _) ->
        [name ^ " = __allocate(" ^ String.concat ", " (List.map node2str dims) ^ ");"]

    | Statements (stats, _) -> List.concat (List.map node2lines stats)

    (* Catch-all, whould never happen *)
    | _ -> failwith "invalid node"

(* node -> string *)
and node2str node =
    let concat sep nodes = String.concat sep (List.map node2str nodes) in
    match node with
    (* Global *)
    | Program (decls, _) ->
        let decl2str decl = String.concat "\n" (node2lines decl) in
        String.concat "\n\n" (List.map decl2str decls)
    | Param (param_type, name, _) -> (type2str param_type) ^ " " ^ name

    (* 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, _) -> node2str value
    | Var (v, _) -> v
    | Deref (name, dims, _) -> name ^ (node2str (ArrayConst (dims, noloc)))
    | Monop (op, opnd, _) -> monop2str op ^ node2str opnd
    | Binop (op, left, right, _) ->
        "(" ^ node2str left ^ binop2str op ^ node2str right ^ ")"
    | Cond (cond, t, f, _) ->
        (node2str cond) ^ " ? " ^ node2str t ^ " : " ^ node2str f
    | TypeCast (ctype, value, _) ->
        "(" ^ type2str ctype ^ ")" ^ node2str value
    | FunCall (name, args, _) ->
        name ^ "(" ^ (concat ", " args) ^ ")"

    | node -> String.concat "\n" (node2lines node)