open Ast

let rec c_type = function
  | TInt -> "int"
  | TBool -> "bool"
  | TVoid -> "void"
  | TStruct n -> "struct " ^ n
  | TArray t -> c_type t ^ "*"
  | TComptime t -> c_type t   (* comptime types lower to their runtime equivalent *)
  | TCombo (t, _) -> c_type t (* combo types use the runtime (left) type in C *)

let rec expr_to_c = function
  | IntLit n -> string_of_int n
  | BoolLit true -> "true"
  | BoolLit false -> "false"
  | Var n -> n
  | Binop (op, a, b) -> Printf.sprintf "(%s %s %s)" (expr_to_c a) (Ast.string_of_binop op) (expr_to_c b)
  | Unop (op, e) -> Printf.sprintf "(%s%s)" (Ast.string_of_unop op) (expr_to_c e)
  | Assign (lhs, rhs) -> Printf.sprintf "(%s = %s)" (expr_to_c lhs) (expr_to_c rhs)
  | Call (callee, args) ->
      let args_s = args |> List.map expr_to_c |> String.concat ", " in
      Printf.sprintf "%s(%s)" (expr_to_c callee) args_s
  | ArrayGet (arr, idx) -> Printf.sprintf "%s[%s]" (expr_to_c arr) (expr_to_c idx)
  | StructGet (obj, fld) -> Printf.sprintf "%s.%s" (expr_to_c obj) fld

let indent n = String.make (2 * n) ' '

let rec stmt_to_c ?(level = 1) = function
  | VarDecl (_, n, Some t, None) -> Printf.sprintf "%s%s %s;" (indent level) (c_type t) n
  | VarDecl (_, n, Some t, Some e) -> Printf.sprintf "%s%s %s = %s;" (indent level) (c_type t) n (expr_to_c e)
  | VarDecl (_, n, None, _) -> Printf.sprintf "%s/* unresolved type for %s */" (indent level) n
  | Expr e -> Printf.sprintf "%s%s;" (indent level) (expr_to_c e)
  | Return None -> Printf.sprintf "%sreturn;" (indent level)
  | Return (Some e) -> Printf.sprintf "%sreturn %s;" (indent level) (expr_to_c e)
  | Block body ->
      let inner = body |> List.map (stmt_to_c ~level:(level + 1)) |> String.concat "\n" in
      if String.equal inner "" then Printf.sprintf "%s{ }" (indent level)
      else Printf.sprintf "%s{\n%s\n%s}" (indent level) inner (indent level)
  | If (cond, tbranch, ebranch) ->
      let then_body = tbranch |> List.map (stmt_to_c ~level:(level + 1)) |> String.concat "\n" in
      let else_body = ebranch |> List.map (stmt_to_c ~level:(level + 1)) |> String.concat "\n" in
      if ebranch = [] then
        Printf.sprintf "%sif (%s) {\n%s\n%s}" (indent level) (expr_to_c cond) then_body (indent level)
      else
        Printf.sprintf "%sif (%s) {\n%s\n%s} else {\n%s\n%s}" (indent level) (expr_to_c cond)
          then_body (indent level) else_body (indent level)
  | ForEach (it_t, it_name, iterable, body) ->
      let body_inner = body |> List.map (stmt_to_c ~level:(level + 2)) |> String.concat "\n" in
      (match iterable with
      | Var arr_name ->
          Printf.sprintf
            "%sfor (int __i = 0; __i < %s_len; ++__i) {\n%s%s %s = %s[__i];\n%s\n%s}"
            (indent level) arr_name (indent (level + 1)) (c_type it_t) it_name arr_name body_inner (indent level)
      | _ ->
          Printf.sprintf
            "%s/* foreach over complex iterable not supported in strict C lowering */\n%sfor (int __i = 0; ; ++__i) {\n%s%s %s = /* TODO */ 0;\n%s\n%s}"
            (indent level) (indent level) (indent (level + 1)) (c_type it_t) it_name body_inner (indent level))

let struct_to_c s =
  let fields =
    s.fields
    |> List.map (fun (t, n) -> Printf.sprintf "  %s %s;" (c_type t) n)
    |> String.concat "\n"
  in
  Printf.sprintf "struct %s {\n%s\n};" s.sname fields

let func_to_c f =
  let params =
    f.params
    |> List.map (fun (n, t) -> Printf.sprintf "%s %s" (c_type t) n)
    |> String.concat ", "
  in
  let body = f.body |> List.map (stmt_to_c ~level:1) |> String.concat "\n" in
  Printf.sprintf "%s %s(%s) {\n%s\n}" (c_type f.ret) f.name params body

let top_to_c = function
  | TopStruct s -> struct_to_c s
  | TopFunc f -> func_to_c f
    | TopGlobalVar (_, n, Some t, init) ->
      (match init with
      | None -> Printf.sprintf "%s %s;" (c_type t) n
      | Some e -> Printf.sprintf "%s %s = %s;" (c_type t) n (expr_to_c e))
    | TopGlobalVar (_, n, None, _) ->
      Printf.sprintf "/* unresolved type for global %s */" n

let generate (prog : program) =
  let header = "#include <stdbool.h>\n\n" in
  header ^ String.concat "\n\n" (List.map top_to_c prog) ^ "\n"