open Ast

let escape_json s =
  let b = Buffer.create (String.length s + 16) in
  String.iter
    (function
      | '"' -> Buffer.add_string b "\\\""
      | '\\' -> Buffer.add_string b "\\\\"
      | '\n' -> Buffer.add_string b "\\n"
      | '\r' -> Buffer.add_string b "\\r"
      | '\t' -> Buffer.add_string b "\\t"
      | c -> Buffer.add_char b c)
    s;
  Buffer.contents b

let q s = Printf.sprintf "\"%s\"" (escape_json s)
let obj fields = "{" ^ String.concat "," fields ^ "}"
let arr xs = "[" ^ String.concat "," xs ^ "]"
let kv k v = q k ^ ":" ^ v

let rec typ_to_json = function
  | TInt -> obj [ kv "kind" (q "int") ]
  | TBool -> obj [ kv "kind" (q "bool") ]
  | TVoid -> obj [ kv "kind" (q "void") ]
  | TStruct name -> obj [ kv "kind" (q "struct"); kv "name" (q name) ]
  | TArray t -> obj [ kv "kind" (q "array"); kv "elem" (typ_to_json t) ]
  | TComptime t -> obj [ kv "kind" (q "comptime"); kv "inner" (typ_to_json t) ]
  | TCombo (t1, t2) -> obj [ kv "kind" (q "combo"); kv "runtime" (typ_to_json t1); kv "comptime" (typ_to_json t2) ]

let rec expr_to_json = function
  | IntLit n -> obj [ kv "node" (q "IntLit"); kv "value" (string_of_int n) ]
  | BoolLit b -> obj [ kv "node" (q "BoolLit"); kv "value" (if b then "true" else "false") ]
  | Var n -> obj [ kv "node" (q "Var"); kv "name" (q n) ]
  | Binop (op, a, b) ->
      obj
        [ kv "node" (q "Binop"); kv "op" (q (Ast.string_of_binop op)); kv "left" (expr_to_json a);
          kv "right" (expr_to_json b) ]
  | Unop (op, e) ->
      obj [ kv "node" (q "Unop"); kv "op" (q (Ast.string_of_unop op)); kv "expr" (expr_to_json e) ]
  | Assign (lhs, rhs) ->
      obj [ kv "node" (q "Assign"); kv "lhs" (expr_to_json lhs); kv "rhs" (expr_to_json rhs) ]
  | Call (callee, args) ->
      obj [ kv "node" (q "Call"); kv "callee" (expr_to_json callee); kv "args" (arr (List.map expr_to_json args)) ]
  | ArrayGet (arr_e, idx) ->
      obj [ kv "node" (q "ArrayGet"); kv "array" (expr_to_json arr_e); kv "index" (expr_to_json idx) ]
  | StructGet (target, field) ->
      obj [ kv "node" (q "StructGet"); kv "target" (expr_to_json target); kv "field" (q field) ]

let rec stmt_to_json = function
  | VarDecl (is_mut, n, t_opt, init) ->
      obj
        [ kv "node" (q "VarDecl"); kv "mut" (if is_mut then "true" else "false");
          kv "name" (q n);
          kv "type" (match t_opt with None -> "null" | Some t -> typ_to_json t);
          kv "init" (match init with None -> "null" | Some e -> expr_to_json e) ]
  | Expr e -> obj [ kv "node" (q "Expr"); kv "expr" (expr_to_json e) ]
  | If (cond, tbranch, ebranch) ->
      obj
        [ kv "node" (q "If"); kv "cond" (expr_to_json cond); kv "then" (arr (List.map stmt_to_json tbranch));
          kv "else" (arr (List.map stmt_to_json ebranch)) ]
  | ForEach (it_t, it_name, iterable, body) ->
      obj
        [ kv "node" (q "ForEach"); kv "iterType" (typ_to_json it_t); kv "iterName" (q it_name);
          kv "iterable" (expr_to_json iterable); kv "body" (arr (List.map stmt_to_json body)) ]
  | Return eo ->
      obj [ kv "node" (q "Return"); kv "expr" (match eo with None -> "null" | Some e -> expr_to_json e) ]
  | Block body -> obj [ kv "node" (q "Block"); kv "body" (arr (List.map stmt_to_json body)) ]

let top_to_json = function
  | TopStruct s ->
      obj
        [ kv "node" (q "Struct"); kv "name" (q s.sname);
          kv "fields"
            (arr
               (List.map
                  (fun (t, n) -> obj [ kv "type" (typ_to_json t); kv "name" (q n) ])
                  s.fields)) ]
  | TopFunc f ->
      obj
        [ kv "node" (q "Function"); kv "name" (q f.name); kv "ret" (typ_to_json f.ret);
          kv "params"
            (arr
               (List.map
                  (fun (n, t) -> obj [ kv "name" (q n); kv "type" (typ_to_json t) ])
                  f.params));
          kv "body" (arr (List.map stmt_to_json f.body)) ]
  | TopGlobalVar (is_mut, n, t_opt, init) ->
      obj
        [ kv "node" (q "GlobalVar"); kv "mut" (if is_mut then "true" else "false");
          kv "name" (q n);
          kv "type" (match t_opt with None -> "null" | Some t -> typ_to_json t);
          kv "init" (match init with None -> "null" | Some e -> expr_to_json e) ]

let generate (prog : program) =
  obj [ kv "node" (q "Program"); kv "tops" (arr (List.map top_to_json prog)) ]