open Syntax
module Env = Map.Make (String)

exception Error of string

let context = Llvm.global_context ()

let verify_function f =
  try
    (match Llvm_analysis.verify_function f with
    | true -> ()
    | false ->
        Format.printf "invalid function generated\n%s\n%!"
          (Llvm.string_of_llvalue f);
        Llvm_analysis.assert_valid_function f);
    f
  with e ->
    Llvm.delete_function f;
    raise e

let declare_function module_ name f_type args =
  match Llvm.lookup_function name module_ with
  | None -> Llvm.declare_function name f_type module_
  | Some f ->
      if Array.length (Llvm.basic_blocks f) == 0 then () else ();
      if Array.length (Llvm.params f) == List.length args then ()
      else
        raise
          (Error
             (Format.asprintf
                "function signature mismatch, '%s' expected %d args, got %d \
                 args"
                name (List.length args)
                (Array.length (Llvm.params f))));
      f

let lookup_function module_ fname =
  match Llvm.lookup_function fname module_ with
  | Some f -> f
  | None ->
      Llvm.dump_module module_;
      failwith ("function not defined: " ^ fname)

let codegen_ty_int (ty_int : ty_int) =
  match ty_int with
  | Int1 -> Llvm.i1_type context
  | Int32 -> Llvm.i32_type context
  | Int64 -> Llvm.i64_type context

let rec codegen_ty (ty : ty) =
  match ty with
  | TInt ty_int -> codegen_ty_int ty_int
  | TDouble -> Llvm.double_type context
  | TPtr -> Llvm.pointer_type context
  | TVoid -> Llvm.void_type context
  | TArray (ty, size) -> Llvm.array_type (codegen_ty ty) size

let rec codegen_expr module_ builder (expr : expr) env =
  match expr with
  | Int (ty_int, n) -> Llvm.const_of_int64 (codegen_ty_int ty_int) n false
  | Double f -> Llvm.const_float (Llvm.double_type context) f
  | Var x -> (
      match Env.find_opt x env with
      | Some value -> value
      | None -> failwith (Format.sprintf "undefined variable: '%s'" x))
  | Let (x, e1, e2) ->
      let v1 = codegen_expr module_ builder e1 env in
      let env' = Env.add x v1 env in
      codegen_expr module_ builder e2 env'
  | Prim1 (op1, e1) -> (
      let v1 = codegen_expr module_ builder e1 env in
      match op1 with
      | Not -> Llvm.build_not v1 "not" builder
      | Neg -> Llvm.build_neg v1 "neg" builder
      | FNeg -> Llvm.build_fneg v1 "fneg" builder
      | Float_of_int ->
          Llvm.build_sitofp v1 (Llvm.double_type context) "float_of_int" builder
      | Int_of_float ->
          Llvm.build_fptosi v1 (Llvm.i64_type context) "int_of_float" builder)
  | Prim2 (op2, e1, e2) -> (
      let v1 = codegen_expr module_ builder e1 env in
      let v2 = codegen_expr module_ builder e2 env in
      match op2 with
      | Eq -> Llvm.build_icmp Llvm.Icmp.Eq v1 v2 "eq" builder
      | Lt -> Llvm.build_icmp Llvm.Icmp.Slt v1 v2 "lt" builder
      | Le -> Llvm.build_icmp Llvm.Icmp.Sle v1 v2 "le" builder
      | Gt -> Llvm.build_icmp Llvm.Icmp.Sgt v1 v2 "gt" builder
      | Ge -> Llvm.build_icmp Llvm.Icmp.Sge v1 v2 "ge" builder
      | And -> Llvm.build_and v1 v2 "and" builder
      | Or -> Llvm.build_or v1 v2 "or" builder
      | Plus -> Llvm.build_add v1 v2 "add" builder
      | Minus -> Llvm.build_sub v1 v2 "sub" builder
      | Times -> Llvm.build_mul v1 v2 "mul" builder
      | Div -> Llvm.build_sdiv v1 v2 "sdiv" builder
      | Mod -> Llvm.build_srem v1 v2 "srem" builder
      | Shl -> Llvm.build_shl v1 v2 "shl" builder
      | Lshr -> Llvm.build_lshr v1 v2 "lshr" builder
      | Ashr -> Llvm.build_ashr v1 v2 "ashr" builder
      | FPlus -> Llvm.build_fadd v1 v2 "fadd" builder
      | FMinus -> Llvm.build_fsub v1 v2 "fsub" builder
      | FTimes -> Llvm.build_fmul v1 v2 "fmul" builder
      | FDiv -> Llvm.build_fdiv v1 v2 "fdiv" builder
      | FLt -> Llvm.build_fcmp Llvm.Fcmp.Olt v1 v2 "flt" builder
      | FLe -> Llvm.build_fcmp Llvm.Fcmp.Ole v1 v2 "fle" builder
      | FEq -> Llvm.build_fcmp Llvm.Fcmp.Oeq v1 v2 "feq" builder
      | FGe -> Llvm.build_fcmp Llvm.Fcmp.Oge v1 v2 "fge" builder
      | FGt -> Llvm.build_fcmp Llvm.Fcmp.Ogt v1 v2 "fgt" builder)
  | If (c, e1, e2) ->
      (* cond *)
      let c_val = codegen_expr module_ builder c env in
      let zero = Llvm.const_int (Llvm.i1_type context) 0 in
      let cond_val = Llvm.build_icmp Llvm.Icmp.Ne c_val zero "cond" builder in
      let start_bb = Llvm.insertion_block builder in
      let func = Llvm.block_parent start_bb in
      (* then *)
      let then_bb = Llvm.append_block context "then" func in
      Llvm.position_at_end then_bb builder;
      let then_val = codegen_expr module_ builder e1 env in
      let new_then_bb = Llvm.insertion_block builder in
      (* else *)
      let else_bb = Llvm.append_block context "else" func in
      Llvm.position_at_end else_bb builder;
      let else_val = codegen_expr module_ builder e2 env in
      let new_else_bb = Llvm.insertion_block builder in
      (* merge *)
      let merge_bb = Llvm.append_block context "if_join" func in
      Llvm.position_at_end merge_bb builder;
      let incoming = [ (then_val, new_then_bb); (else_val, new_else_bb) ] in
      let phi = Llvm.build_phi incoming "phi" builder in
      (* return to start block *)
      Llvm.position_at_end start_bb builder;
      let _ = Llvm.build_cond_br cond_val then_bb else_bb builder in
      (* set a unconditional branch at the end of the 'then' block and the 'else' block to the 'merge' block. *)
      Llvm.position_at_end new_then_bb builder;
      let _ = Llvm.build_br merge_bb builder in
      Llvm.position_at_end new_else_bb builder;
      let _ = Llvm.build_br merge_bb builder in
      (* finally, set the builder to the end of the merge block. *)
      Llvm.position_at_end merge_bb builder;
      phi
  | Switch (e, cases, default) ->
      let start_bb = Llvm.insertion_block builder in
      let func = Llvm.block_parent start_bb in
      let case_bbs =
        List.map
          (fun (case_id, _) ->
            Llvm.append_block context (Format.sprintf "case_%Ld" case_id) func)
          cases
      in
      let default_bb = Llvm.append_block context "default" func in
      let merge_bb = Llvm.append_block context "switch_join" func in
      (* cond *)
      let e_val = codegen_expr module_ builder e env in
      let switch_val =
        Llvm.build_switch e_val default_bb (List.length cases) builder
      in
      (* case blocks *)
      let incoming_cases =
        List.map2
          (fun (case_id, case_expr) case_bb ->
            Llvm.position_at_end case_bb builder;
            let case_val = codegen_expr module_ builder case_expr env in
            let case_id_val =
              Llvm.const_of_int64 (Llvm.i64_type context) case_id false
            in
            ignore (Llvm.build_br merge_bb builder);
            Llvm.add_case switch_val case_id_val case_bb;
            (case_val, case_bb))
          cases case_bbs
      in
      (* default block *)
      Llvm.position_at_end default_bb builder;
      let default_val = codegen_expr module_ builder default env in
      ignore (Llvm.build_br merge_bb builder);
      (* merge block *)
      Llvm.position_at_end merge_bb builder;
      let incoming = (default_val, default_bb) :: incoming_cases in
      let phi = Llvm.build_phi incoming "phi" builder in
      phi
  | For (var_name, start, stop, body) ->
      let start_val = codegen_expr module_ builder start env in
      (* loop header *)
      let preheader_bb = Llvm.insertion_block builder in
      let func = Llvm.block_parent preheader_bb in
      let loop_bb = Llvm.append_block context "loop" func in
      (* fall through *)
      ignore (Llvm.build_br loop_bb builder);
      (* start insertion in loop_bb *)
      Llvm.position_at_end loop_bb builder;

      (* start phi node *)
      let var = Llvm.build_phi [ (start_val, preheader_bb) ] var_name builder in
      let env' = Env.add var_name var env in
      ignore (codegen_expr module_ builder body env');

      (* step *)
      let step = Llvm.const_int (Llvm.i64_type context) 1 in
      let next_var = Llvm.build_add var step "next_var" builder in

      (* condition *)
      let stop_val = codegen_expr module_ builder stop env in
      let cond =
        Llvm.build_icmp Llvm.Icmp.Slt next_var stop_val "loop_cond" builder
      in

      (* after loop block *)
      let loop_end_bb = Llvm.insertion_block builder in
      let after_bb = Llvm.append_block context "after_loop" func in

      (* branch *)
      ignore (Llvm.build_cond_br cond loop_bb after_bb builder);

      (* set insertion point to after_bb *)
      Llvm.position_at_end after_bb builder;

      (* add incoming to phi node *)
      Llvm.add_incoming (next_var, loop_end_bb) var;

      (* for expr always returns zero *)
      Llvm.const_int (Llvm.i64_type context) 0
  | Apply (e, args, ret_ty) ->
      let arg_vals =
        args
        |> List.map fst
        |> List.map (fun expr -> codegen_expr module_ builder expr env)
        |> Array.of_list
      in
      let arg_tys =
        args
        |> List.map snd
        |> List.map (fun ty -> codegen_ty ty)
        |> Array.of_list
      in
      let ret_ty = codegen_ty ret_ty in
      let func_type = Llvm.function_type ret_ty arg_tys in
      let func =
        match e with
        | Func fname -> declare_function module_ fname func_type args
        | Var _ -> codegen_expr module_ builder e env
        | _ ->
            failwith
              (Format.asprintf
                 "expected the name of a top level function or a variable \
                  whose value is a function pointer, but got %a"
                 pp_expr e)
      in
      Llvm.build_call func_type func arg_vals "call" builder
  | ApplyExt (name, tys, ret_ty, args) ->
      let block = Llvm.insertion_block builder in
      Llvm.position_at_end block builder;
      let func_type =
        Llvm.function_type (codegen_ty ret_ty)
          (Array.of_list (List.map codegen_ty tys))
      in
      let func = declare_function module_ name func_type args in
      Llvm.set_linkage Llvm.Linkage.External func;
      let arg_vals =
        args
        |> List.map (fun expr -> codegen_expr module_ builder expr env)
        |> Array.of_list
      in
      let tmp_name = match ret_ty with TVoid -> "" | _ -> "call_ext" in
      Llvm.build_call func_type func arg_vals tmp_name builder
  | Load (ty, e) ->
      let t = codegen_ty ty in
      let v = codegen_expr module_ builder e env in
      Llvm.build_load t v "" builder
  | Store (e_value, e_ptr) ->
      let v = codegen_expr module_ builder e_value env in
      let ptr = codegen_expr module_ builder e_ptr env in
      Llvm.build_store v ptr builder
  | Gep (ty, e, idx) ->
      let ty = codegen_ty ty in
      let v = codegen_expr module_ builder e env in
      let idx_v = Llvm.const_int (Llvm.i64_type context) idx in
      Llvm.build_gep ty v [| idx_v |] "" builder
  | Malloc ty ->
      let lltype = codegen_ty ty in
      Llvm.build_malloc lltype "" builder
  | Alloca ty ->
      let lltype = codegen_ty ty in
      Llvm.build_alloca lltype "" builder
  | Seq exprs -> (
      match exprs with
      | [] ->
          failwith "empty Seq is not valid during lowering of Cmm to LLVM IR"
      | [ expr ] -> codegen_expr module_ builder expr env
      | expr :: rest ->
          let _ = codegen_expr module_ builder expr env in
          codegen_expr module_ builder (Seq rest) env)
  | Func fname -> lookup_function module_ fname
  | Extractelement (e_vec, e_idx) ->
      let vec = codegen_expr module_ builder e_vec env in
      let idx = codegen_expr module_ builder e_idx env in
      Llvm.build_extractelement vec idx "extractelement" builder
  | Insertelement (e_vec, e_val, e_idx) ->
      let vec = codegen_expr module_ builder e_vec env in
      let value = codegen_expr module_ builder e_val env in
      let idx = codegen_expr module_ builder e_idx env in
      Llvm.build_insertelement vec value idx "insertelement" builder
  | Inttoptr (e, ty) ->
      Llvm.build_inttoptr
        (codegen_expr module_ builder e env)
        (codegen_ty ty) "inttoptr" builder
  | Ptrtoint (e, ty) ->
      Llvm.build_ptrtoint
        (codegen_expr module_ builder e env)
        (codegen_ty ty) "ptrtoint" builder
  | Intcast (e, ty) ->
      Llvm.build_intcast
        (codegen_expr module_ builder e env)
        (codegen_ty ty) "intcast" builder
  | Bitcast (e, ty) ->
      Llvm.build_bitcast
        (codegen_expr module_ builder e env)
        (codegen_ty ty) "bitcast" builder

let codegen_fundef module_ builder env fundef =
  let args, arg_tys = List.split fundef.args in
  let func_type =
    Llvm.function_type (codegen_ty fundef.return)
      (Array.of_list (List.map codegen_ty arg_tys))
  in
  let func = declare_function module_ fundef.name func_type args in
  (* create basic block for the function *)
  let bb = Llvm.append_block context "entry" func in
  Llvm.position_at_end bb builder;
  let params = Llvm.params func in
  (* set param names *)
  let param_env =
    args
    |> List.mapi (fun i arg ->
           let param = params.(i) in
           Llvm.set_value_name arg param;
           (arg, param))
    |> List.to_seq
  in
  (* add params to env *)
  let env' = Env.add_seq param_env env in
  let ret_val = codegen_expr module_ builder fundef.body env' in
  let _ = Llvm.build_ret ret_val builder in
  func

let codegen : compunit -> Llvm.llmodule =
 fun compunit ->
  let _ = Llvm.enable_pretty_stacktrace () in
  let module_ = Llvm.create_module context "main" in
  let builder = Llvm.builder context in
  let rec go compunit env =
    match compunit with
    | [] -> module_
    | fundef :: rest ->
        let func = codegen_fundef module_ builder env fundef in
        (match Llvm_analysis.verify_module module_ with
        | None -> ()
        | Some reason -> Format.printf "invalid module generated\n%s\n%!" reason);
        go rest (Env.add fundef.name func env)
  in
  let compiled_module = go compunit Env.empty in
  compiled_module
