type typ =
  | TInt
  | TBool
  | TVoid
  | TStruct of string
  | TArray of typ

type binop =
  | Add
  | Sub
  | Mul
  | Div
  | Mod
  | And
  | Or
  | Eq
  | Ne
  | Lt
  | Le
  | Gt
  | Ge

type unop = Neg | Not

type expr =
  | IntLit of int
  | BoolLit of bool
  | Var of string
  | Binop of binop * expr * expr
  | Unop of unop * expr
  | Assign of expr * expr
  | Call of expr * expr list
  | ArrayGet of expr * expr
  | StructGet of expr * string

type stmt =
  | VarDecl of typ * string * expr option
  | Expr of expr
  | If of expr * stmt list * stmt list
  | ForEach of typ * string * expr * stmt list
  | Return of expr option
  | Block of stmt list

type func = {
  name : string;
  params : (typ * string) list;
  ret : typ;
  body : stmt list;
}

type struct_def = {
  sname : string;
  fields : (typ * string) list;
}

type top =
  | TopStruct of struct_def
  | TopFunc of func
  | TopGlobalVar of typ * string * expr option

type program = top list

let string_of_typ =
  let rec go = function
    | TInt -> "int"
    | TBool -> "bool"
    | TVoid -> "void"
    | TStruct n -> "struct " ^ n
    | TArray t -> go t ^ "[]"
  in
  go

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

let string_of_unop = function Neg -> "-" | Not -> "!"

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

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

let rec string_of_stmt ?(level = 0) st =
  let i = indent level in
  match st with
  | VarDecl (t, n, None) -> Printf.sprintf "%sVarDecl(%s %s)" i (string_of_typ t) n
  | VarDecl (t, n, Some e) ->
      Printf.sprintf "%sVarDecl(%s %s = %s)" i (string_of_typ t) n (string_of_expr e)
  | Expr e -> Printf.sprintf "%sExpr(%s)" i (string_of_expr e)
  | Return None -> Printf.sprintf "%sReturn" i
  | Return (Some e) -> Printf.sprintf "%sReturn(%s)" i (string_of_expr e)
  | Block body ->
      let body_s = String.concat "\n" (List.map (string_of_stmt ~level:(level + 1)) body) in
      if String.equal body_s "" then Printf.sprintf "%sBlock" i
      else Printf.sprintf "%sBlock\n%s" i body_s
  | If (cond, tbranch, ebranch) ->
      let then_s = String.concat "\n" (List.map (string_of_stmt ~level:(level + 1)) tbranch) in
      let else_s = String.concat "\n" (List.map (string_of_stmt ~level:(level + 1)) ebranch) in
      if ebranch = [] then Printf.sprintf "%sIf(%s)\n%s" i (string_of_expr cond) then_s
      else Printf.sprintf "%sIf(%s)\n%s\n%sElse\n%s" i (string_of_expr cond) then_s i else_s
  | ForEach (it_t, it_name, iterable, body) ->
      let body_s = String.concat "\n" (List.map (string_of_stmt ~level:(level + 1)) body) in
      Printf.sprintf "%sForEach(%s %s in %s)\n%s" i (string_of_typ it_t) it_name (string_of_expr iterable)
        body_s

let string_of_top = function
  | TopStruct s ->
      let fields =
        s.fields |> List.map (fun (t, n) -> Printf.sprintf "  Field(%s %s)" (string_of_typ t) n)
        |> String.concat "\n"
      in
      if String.equal fields "" then Printf.sprintf "Struct %s" s.sname
      else Printf.sprintf "Struct %s\n%s" s.sname fields
  | TopGlobalVar (t, n, init) ->
      (match init with
      | None -> Printf.sprintf "GlobalVar(%s %s)" (string_of_typ t) n
      | Some e -> Printf.sprintf "GlobalVar(%s %s = %s)" (string_of_typ t) n (string_of_expr e))
  | TopFunc f ->
      let params =
        f.params
        |> List.map (fun (t, n) -> Printf.sprintf "%s %s" (string_of_typ t) n)
        |> String.concat ", "
      in
      let body = String.concat "\n" (List.map (string_of_stmt ~level:1) f.body) in
      if String.equal body "" then Printf.sprintf "Function %s(%s) -> %s" f.name params (string_of_typ f.ret)
      else Printf.sprintf "Function %s(%s) -> %s\n%s" f.name params (string_of_typ f.ret) body

let string_of_program (prog : program) = prog |> List.map string_of_top |> String.concat "\n\n"

let rec equal_typ a b =
  match (a, b) with
  | TInt, TInt | TBool, TBool | TVoid, TVoid -> true
  | TStruct x, TStruct y -> String.equal x y
  | TArray x, TArray y -> equal_typ x y
  | _ -> false