open Syntax

let log (config : Config.t) fmt =
  if config.verbose then Format.eprintf fmt
  else Format.ifprintf Format.err_formatter fmt

let pp_ev ~pp =
 fun fmt ty ->
  let magenta = "\027[35m" in
  let reset = "\027[0m" in
  Format.fprintf fmt "%s⟨%a⟩%s" magenta pp ty reset

let pp_type_ev = pp_ev ~pp:Type.pp
let pp_rt_type_ev = pp_ev ~pp:Evidence.RuntimeType.pp

module TypeEv = struct
  include Evidence.TypeEvidence

  let pp = pp_type_ev
end

module RTypeEv = struct
  include Evidence.RuntimeTypeEvidence

  let pp = pp_rt_type_ev
end

let parse (conf : Config.t) =
  log conf "Parsing\n%!";
  let program_opt =
    match conf.input with
    | Stdin -> Parsing.Driver.parse_stdin ()
    | InputFile filename -> Parsing.Driver.parse_file filename
    | InputString str -> Parsing.Driver.parse_string str
  in
  let parse_tree = Option.get program_opt in
  let pp_meta _ _ = () in
  log conf "%a\n\n%!" (Parse_tree.Program.pp ~pp_meta) parse_tree;
  log conf "Converting to typed AST\n%!";
  let program = Typed_ast.of_parse_tree parse_tree in
  log conf "%a\n\n%!" Typed_ast.pp_program program;
  program

let normalize (conf : Config.t) program =
  log conf "Normalizing\n%!";
  let program' = Normalize.normalize program in
  log conf "%a\n\n%!" Typed_ast.pp_program program';
  program'

let typecheck (conf : Config.t) program =
  match conf.typechecking with
  | Static ->
      log conf "Typechecking\n%!";
      Typechecking.Static.typecheck program;
      program
  | _ -> program

module type EV = sig
  include Evidence.Evidence

  val pp : Format.formatter -> t -> unit
end

let elaborate (type ev) (module Ev : EV with type t = ev and type g = Type.t)
    conf program =
  log conf "Elaborating\n%!";
  let program' =
    match conf.typechecking with
    | Config.Static | Config.Dynamic | Config.Nocheck ->
        let elaborated = Syntax.Elaborate.Elaborate.elaborate program in
        (* TODO: Find an alternative to using Obj.magic *)
        (* Cast from Ev.g to ev (which will always be Type.t for these backends) *)
        (Obj.magic elaborated : ev Ast.t)
    | _ ->
        let module Elaborate = Evidence.Elaborate (Ev) in
        Elaborate.elaborate program
  in
  log conf "%a\n\n%!" (Ast.pp ~show_meta:conf.show_evidence Ev.pp) program';
  program'

let simplify (type ev) (module Ev : EV with type t = ev and type g = Type.t)
    conf program =
  log conf "Simplifying (removing unnecessary ascriptions)\n%!";
  let module Simplify = Simplify.Make (Ev) in
  let program' =
    match conf.typechecking with
    | Config.Static | Config.Dynamic | Config.Nocheck -> program
    | _ -> Simplify.simplify () program
  in
  log conf "%a\n\n%!" (Ast.pp ~show_meta:conf.show_evidence Ev.pp) program';
  program'

let insert_checks (type ev)
    (module Ev : EV with type t = ev and type g = Type.t) (conf : Config.t) e =
  match conf.typechecking with
  | Dynamic ->
      log conf "Inserting dynamic type checks\n%!";
      let t = Typechecking.Dynamic.insert_checks e in
      log conf "%a\n\n%!" (Ast.pp ~show_meta:conf.show_evidence Ev.pp) t;
      e
  | _ -> e

let uniquify (type ev) (module Ev : EV with type t = ev and type g = Type.t)
    conf e =
  log conf "Uniquifying\n%!";
  let t = Uniquify.uniquify e in
  log conf "%a\n\n%!" (Ast.pp ~show_meta:conf.show_evidence Ev.pp) t;
  t

let anf_transform (type ev)
    (module Ev : EV with type t = ev and type g = Type.t) conf e =
  log conf "Transforming to ANF\n%!";
  let anf = Anf.of_ast e in
  log conf "%a\n\n%!" (Anf.pp_expr ~show_meta:conf.show_evidence Ev.pp) anf;
  anf

let purify_letrec (type ev)
    (module Ev : EV with type t = ev and type g = Type.t) conf e =
  log conf "Purifying letrec\n%!";
  let e = Purify.purify_expr e in
  log conf "%a\n\n%!" (Anf.pp_expr ~show_meta:conf.show_evidence Ev.pp) e;
  e

let elaborate_dyn (type ev)
    (module Ev : EV with type t = ev and type g = Type.t)
    (module Tm : Evidence.TypeMatching with type t = ev) (conf : Config.t) e =
  let module ElaborateDyn = Elaborate_dyn.ElaborateDyn (Ev) (Tm) in
  match conf.typechecking with
  | GradualMonoDFO | GradualGuardedDFO | GradualGuarded | GradualMonoClosures
  | GradualMonoValues ->
      log conf "Elaborating dyn (Inserting ascriptions to runtime evidence)\n%!";
      let e = ElaborateDyn.elaborate e in
      log conf "%a\n\n%!" (Anf.pp_expr ~show_meta:conf.show_evidence Ev.pp) e;
      e
  | _ -> e

let closure_convert (type ev)
    (module Ev : EV with type t = ev and type g = Type.t) conf anf =
  log conf "Converting closures\n%!";
  let program = Closure.closure anf in
  log conf "%a\n\n%!"
    (Closure.pp_program ~show_meta:conf.show_evidence Ev.pp)
    program;
  program

let codegen (type ev) (module Lower : Middle_end.Lower.S with type t = ev) conf
    (program : ev Closure.program) =
  log conf "Converting to LLVM IR using the new pipeline\n%!";
  log conf "Converting to C--\n%!";
  let cmm = Lower.codegen_program program in
  log conf "%a\n\n%!" Cmm.Syntax.pp_compunit cmm;
  log conf "Converting to C-- to LLVM IR\n%!";
  let llmodule = Cmm.Codegen.codegen cmm in
  if Option.is_some conf.target then
    Llvm.set_target_triple (Option.get conf.target) llmodule;
  match conf.output with
  | OutputFile file -> Llvm.print_module file llmodule
  | Stderr -> Llvm.dump_module llmodule

let pipeline : type ev.
    (module EV with type t = ev and type g = Type.t) ->
    (module Evidence.TypeMatching with type t = ev) ->
    (module Middle_end.Lower.S with type t = ev) ->
    Config.t ->
    _ =
 fun (module Ev) (module Tm) (module Lower) config ->
  config
  |> parse
  |> normalize config
  |> typecheck config
  |> elaborate (module Ev) config
  |> simplify (module Ev) config
  |> insert_checks (module Ev) config
  |> uniquify (module Ev) config
  |> anf_transform (module Ev) config
  |> elaborate_dyn (module Ev) (module Tm) config
  |> purify_letrec (module Ev) config
  |> closure_convert (module Ev) config
  |> codegen (module Lower) config

let compile (config : Config.t) =
  match config.typechecking with
  | Config.Static | Config.Dynamic | Config.Nocheck | Config.GradualSuccinct
  | Config.GradualMonoClosures | Config.GradualMonoValues
  | Config.GradualGuarded ->
      let (module Lower : Middle_end.Lower.S with type t = Type.t) =
        match config.typechecking with
        | Static | Nocheck -> (module Middle_end.Untagged.Lower)
        | Dynamic -> (module Middle_end.Tagged.Lower)
        | GradualSuccinct -> (module Middle_end.Gradual_succinct.Lower)
        | GradualMonoClosures ->
            let () =
              Middle_end.Combined.proxy_values := true;
              Middle_end.Combined.proxy_closures := false
            in
            (module Middle_end.Combined)
        | GradualMonoValues ->
            let () =
              Middle_end.Combined.proxy_values := false;
              Middle_end.Combined.proxy_closures := false
            in
            (module Middle_end.Combined)
        | GradualGuarded ->
            let () =
              Middle_end.Combined.proxy_values := true;
              Middle_end.Combined.proxy_closures := true
            in
            (module Middle_end.Combined)
        | GradualMonoDFO | GradualGuardedDFO -> assert false
      in
      pipeline
        (module TypeEv)
        (module Evidence.TypeMatching)
        (module Lower)
        config
  | Config.GradualMonoDFO | Config.GradualGuardedDFO ->
      let (module Lower : Middle_end.Lower.S with type t = RTypeEv.t) =
        match config.typechecking with
        | GradualMonoDFO -> (module Middle_end.Gradual)
        | GradualGuardedDFO -> (module Middle_end.Gradual_proxy)
        | _ -> assert false
      in
      pipeline
        (module RTypeEv)
        (module Evidence.RuntimeTypeMatching)
        (module Lower)
        config
