open Error
open Expr

type stmt =
  | Expr of expr_node
  | Break
  | Continue
  | Print of expr_node
  | VarDecl of { name : string; init : expr_node option }
  | Block of stmt_node list
  | If of { cond : expr_node; then_ : stmt_node; else_ : stmt_node option }
  | While of { cond : expr_node; body : stmt_node }
  | For of {
      init : stmt_node option;
      cond : expr_node option;
      update : expr_node option;
      body : stmt_node;
    }

and stmt_node = { stmt : stmt; pos : code_pos }

let rec show_stmt ?(indent = 0) stmt =
  let indent_s = String.make indent ' ' in
  let show_expr_ind ?(add = 2) = show_expr ~indent:(indent + add) in
  let show_stmt_ind ?(add = 2) = show_stmt ~indent:(indent + add) in
  match stmt with
  | Expr expr -> indent_s ^ "Expr\n" ^ show_expr_ind expr.expr
  | Break -> indent_s ^ "Break"
  | Continue -> indent_s ^ "Continue"
  | Print expr -> indent_s ^ "Print\n" ^ show_expr_ind expr.expr
  | VarDecl { name; init } ->
      let init_s = match init with Some init -> " = \n" ^ show_expr_ind init.expr | None -> "" in
      indent_s ^ "Var " ^ name ^ init_s
  | Block stmts ->
      let stmts_s =
        List.fold_left (fun acc stmt -> acc ^ show_stmt_ind stmt.stmt ^ "\n") "" stmts
      in
      indent_s ^ "Block\n" ^ stmts_s ^ indent_s ^ "End"
  | If { cond; then_; else_ } ->
      let cond_s = show_expr_ind ~add:4 cond.expr in
      let then_s = show_stmt_ind ~add:4 then_.stmt in
      let else_s = Option.map (fun stmt -> show_stmt_ind ~add:4 stmt.stmt) else_ in
      indent_s ^ "If\n" ^ indent_s ^ "  Cond\n" ^ cond_s ^ "\n" ^ indent_s ^ "  Then\n" ^ then_s
      ^ if Option.is_some else_s then "\n" ^ indent_s ^ "  Else\n" ^ Option.get else_s else ""
  | While { cond; body } ->
      let cond_s = show_expr_ind ~add:4 cond.expr in
      let body_s = show_stmt_ind ~add:4 body.stmt in
      indent_s ^ "While\n" ^ indent_s ^ "  Cond\n" ^ cond_s ^ "\n" ^ indent_s ^ "  Body\n" ^ body_s
  | For { init; cond; update; body } ->
      (* let init_s =
           match init with
           | Some { stmt = init; _ } -> indent_s ^ "  Init\n" ^ show_stmt_ind ~add:4 init
           | None -> ""
         in *)
      let init_s =
        init
        |> Option.map (fun { stmt = init; _ } -> indent_s ^ "  Init\n" ^ show_stmt_ind ~add:4 init)
        |> Option.value ~default:""
      in
      (* let cond_s =
         match cond with
         | Some { expr = cond; _ } -> indent_s ^ "  Cond\n" ^ show_expr_ind ~add:4 cond
         | None -> "" *)
      let cond_s =
        cond
        |> Option.map (fun { expr = cond; _ } -> indent_s ^ "  Cond\n" ^ show_expr_ind ~add:4 cond)
        |> Option.value ~default:""
      in
      (* let update_s =
         match update with
         | Some { expr = update; _ } -> indent_s ^ "  Update\n" ^ show_expr_ind ~add:4 update
         | None -> "" *)
      let update_s =
        update
        |> Option.map (fun { expr = update; _ } ->
               indent_s ^ "  Update\n" ^ show_expr_ind ~add:4 update)
        |> Option.value ~default:""
      in
      let body_s = indent_s ^ "  Body\n" ^ show_stmt_ind ~add:4 body.stmt in
      indent_s ^ "For\n" ^ init_s ^ cond_s ^ update_s ^ body_s

let show_stmt_node stmt_node = show_stmt stmt_node.stmt
let make_stmt_node (pos : code_pos) (stmt : stmt) : stmt_node = { stmt; pos }
let make_expr_stmt (pos : code_pos) (expr : expr_node) : stmt_node = Expr expr |> make_stmt_node pos
let make_break (pos : code_pos) : stmt_node = Break |> make_stmt_node pos
let make_continue (pos : code_pos) : stmt_node = Continue |> make_stmt_node pos
let make_print (pos : code_pos) (expr : expr_node) : stmt_node = Print expr |> make_stmt_node pos

let make_var_decl (pos : code_pos) (name : string) (init : expr_node option) =
  VarDecl { name; init } |> make_stmt_node pos

let make_block (pos : code_pos) (stmts : stmt_node list) : stmt_node =
  Block stmts |> make_stmt_node pos

let make_if (pos : code_pos) (cond : expr_node) (then_ : stmt_node) (else_ : stmt_node option) =
  If { cond; then_; else_ } |> make_stmt_node pos

let make_while (pos : code_pos) (cond : expr_node) (body : stmt_node) =
  While { cond; body } |> make_stmt_node pos

let make_for (pos : code_pos) (init : stmt_node option) (cond : expr_node option)
    (update : expr_node option) (body : stmt_node) =
  For { init; cond; update; body } |> make_stmt_node pos