open Format
open Common
open Pp

type 'a program = { funcs : 'a func list; expr : 'a expr }

and 'a func =
  | Func of {
      name : string;
      args : string list;
      body : 'a expr;
      free_vars : string list;
      well_known : bool;
      func_meta : 'a;
    }

and 'a expr = { expr_desc : 'a expr_desc; expr_meta : 'a }

and 'a expr_desc =
  | Ret of 'a cexpr
  | Let of 'a binding * 'a expr
  | Letrec of 'a binding list * 'a expr
  | Seq of 'a expr list

and 'a cexpr = { cexpr_desc : 'a cexpr_desc; cexpr_meta : 'a }

and 'a cexpr_desc =
  | Atom of 'a atom
  | Call of 'a atom * 'a atom list * string option
  | Closure of string * string list
  | Prim1 of Prim.op1 * 'a cexpr
  | Prim2 of Prim.op2 * 'a atom * 'a atom
  | If of 'a atom * 'a expr * 'a expr
  | Match of 'a atom * 'a pattern list * Type.t
  | Tuple of 'a atom list
  | Get of 'a cexpr * int
  | Vect of 'a atom * 'a atom
  | VectGet of 'a atom * 'a atom
  | VectSet of 'a atom * 'a atom * 'a atom
  | VectLength of 'a cexpr
  | Ascribe of 'a cexpr * Type.t
  | AscribeDyn of 'a cexpr * 'a atom * Inversion.t
  | For of string * 'a atom * 'a atom * 'a expr
  | Variant of string * 'a atom list * Type.t

and 'a atom = { atom_desc : 'a atom_desc; atom_meta : 'a }
and 'a atom_desc = Const of Const.t | Var of string
and 'a binding = { name : string; cexpr : 'a cexpr }
and 'a pattern = { casename : string; args : string list; body : 'a expr }

let program funcs expr = { funcs; expr }
let expr expr_meta expr_desc = { expr_desc; expr_meta }
let cexpr cexpr_meta cexpr_desc = { cexpr_desc; cexpr_meta }
let atom atom_meta atom_desc = { atom_desc; atom_meta }
let binding name cexpr = { name; cexpr }

module Env = Map.Make (String)

type string_env = string Env.t

let well_known : (string, bool) Hashtbl.t = Hashtbl.create 10

let is_well_known label =
  Hashtbl.find_opt well_known label |> Option.value ~default:false

let is_lambda (cexpr : 'a Anf.cexpr) =
  match cexpr.cexpr_desc with Anf.Lam _ -> true | _ -> false

let rec closure (e : 'a Anf.expr) : 'a program =
  Gensym.init ();
  let env = Env.empty in
  closure_e env e (fun expr -> { funcs = []; expr })

and closure_e (env : string_env) (e : 'a Anf.expr) k : 'a program =
  let ret (expr_desc : 'a expr_desc) = k (expr e.expr_meta expr_desc) in
  match e.expr_desc with
  | Ret cexpr -> closure_c env cexpr (fun cexpr -> ret (Ret cexpr))
  | Let (b, e) ->
      let env' =
        match b.cexpr.cexpr_desc with
        | Atom { atom_desc = Var x; _ }
        | Ascribe ({ cexpr_desc = Atom { atom_desc = Var x; _ }; _ }, _)
        | AscribeDyn ({ cexpr_desc = Atom { atom_desc = Var x; _ }; _ }, _, _)
          ->
            if Env.mem x env then Env.add b.name (Env.find x env) env else env
        | _ -> env
      in
      closure_c env b.cexpr (fun c ->
          closure_e env' e (fun e -> ret (Let (binding b.name c, e))))
  | Letrec (bindings, expr) -> closure_lambdas env e.expr_meta bindings expr k
  | Seq exprs -> closure_e_list env exprs (fun exprs -> ret (Seq exprs))

and closure_lambdas env meta bindings e k =
  (* create a new label for each binding *)
  let assoc =
    bindings
    |> List.map (fun Anf.{ name; _ } ->
           let label = Gensym.fresh (Format.sprintf "lambda_%s" name) in
           Hashtbl.add well_known label true;
           (name, label))
  in
  let labels = List.map snd assoc in
  let env' = Env.add_seq (List.to_seq assoc) env in
  (* process body *)
  let body_program = closure_e env' e (fun e -> program [] e) in
  let funcs_from_body = body_program.funcs in
  (* process each binding *)
  let bindings, funcs_from_bindings =
    List.map2
      (fun (b : 'a Anf.binding) label ->
        (* Helper to unwrap ascriptions and get the lambda *)
        let rec unwrap_to_lam (cexpr : 'a Anf.cexpr) =
          match cexpr.cexpr_desc with
          | Lam (args, body) -> Some (args, body, cexpr.cexpr_meta)
          | Ascribe (c, _) -> unwrap_to_lam c
          | AscribeDyn (c, _, _) -> unwrap_to_lam c
          | _ -> None
        in
        match unwrap_to_lam b.cexpr with
        | Some (args, body, lam_meta) ->
            let { funcs = funcs_from_body; expr = body } =
              closure_e env' body (fun body -> program [] body)
            in
            let free_vars =
              Anf.free_vars_cexpr b.cexpr
              |> List.filter (fun x -> not @@ String.equal x b.name)
            in
            let body' =
              let binding =
                binding b.name
                  (cexpr lam_meta (Atom (atom lam_meta (Var "self"))))
              in
              expr body.expr_meta (Let (binding, body))
            in
            let func =
              Func
                {
                  name = label;
                  args;
                  body = body';
                  free_vars;
                  well_known = is_well_known label;
                  func_meta = lam_meta;
                }
            in
            let closure = Closure (label, free_vars) in
            let binding' = binding b.name (cexpr b.cexpr.cexpr_meta closure) in
            let funcs = funcs_from_body @ [ func ] in
            (binding', funcs)
        | None ->
            failwith
              (Format.asprintf "expected lambda, but got %a"
                 (Anf.pp_binding ~show_meta:true (fun _ _ -> ()))
                 b))
      bindings labels
    |> List.split
  in
  let funcs_from_bindings = List.flatten funcs_from_bindings in
  let program = k (expr meta (Letrec (bindings, body_program.expr))) in
  let funcs = funcs_from_bindings @ funcs_from_body in
  { funcs; expr = program.expr }

and closure_c (env : string_env) (c : 'a Anf.cexpr) k : 'a program =
  let ret cexpr_desc = k (cexpr c.cexpr_meta cexpr_desc) in
  match c.cexpr_desc with
  | Anf.Atom a -> ret (Atom (closure_a env a))
  | Anf.App (a1, atoms) -> (
      match a1.atom_desc with
      | Anf.Var x when Env.mem x env ->
          let label = Env.find x env in
          let a1 = atom a1.atom_meta (Var x) in
          let atoms = List.map (closure_a env) atoms in
          ret (Call (a1, atoms, Some label))
      | _ -> ret (Call (closure_a env a1, List.map (closure_a env) atoms, None))
      )
  | Anf.Lam _ -> failwith "this case should be handled by closure_e"
  | Anf.Prim1 (op, c) -> closure_c env c (fun c -> ret (Prim1 (op, c)))
  | Anf.Prim2 (op, a1, a2) ->
      ret (Prim2 (op, closure_a env a1, closure_a env a2))
  | Anf.If (a, expr1, expr2) ->
      let a = closure_a env a in
      closure_e env expr1 (fun expr1 ->
          closure_e env expr2 (fun expr2 -> ret (If (a, expr1, expr2))))
  | Anf.Match (a, patterns, ty) ->
      let closure_e_patterns patterns k =
        let bodys = patterns |> List.map (fun Anf.{ body; _ } -> body) in
        let pattern Anf.{ casename; args; _ } body = { casename; args; body } in
        closure_e_list env bodys (fun bodys ->
            k (List.map2 pattern patterns bodys))
      in
      let a = closure_a env a in
      closure_e_patterns patterns (fun patterns ->
          ret (Match (a, patterns, ty)))
  | Anf.Tuple atoms ->
      let atoms = List.map (closure_a env) atoms in
      ret (Tuple atoms)
  | Anf.Get (c, i) -> closure_c env c (fun c -> ret (Get (c, i)))
  | Anf.Vect (a1, a2) -> ret (Vect (closure_a env a1, closure_a env a2))
  | Anf.VectGet (a1, a2) -> ret (VectGet (closure_a env a1, closure_a env a2))
  | Anf.VectSet (a1, a2, a3) ->
      ret (VectSet (closure_a env a1, closure_a env a2, closure_a env a3))
  | Anf.VectLength c -> closure_c env c (fun c -> ret (VectLength c))
  | Anf.Ascribe (c, ty) -> (
      match c.cexpr_desc with
      | Anf.Atom { atom_desc = Var x; atom_meta } when Env.mem x env ->
          let a = atom atom_meta (Var x) in
          let c = cexpr c.cexpr_meta (Atom a) in
          ret (Ascribe (c, ty))
      | _ -> closure_c env c (fun c -> ret (Ascribe (c, ty))))
  | Anf.AscribeDyn (c, a, inv) -> (
      match c.cexpr_desc with
      | Anf.Atom { atom_desc = Var x; atom_meta } when Env.mem x env ->
          let c = cexpr a.atom_meta (Atom (atom atom_meta (Var x))) in
          ret (AscribeDyn (c, closure_a env a, inv))
      | _ ->
          closure_c env c (fun c -> ret (AscribeDyn (c, closure_a env a, inv))))
  | Anf.For (x, a1, a2, expr) ->
      let a1 = closure_a env a1 in
      let a2 = closure_a env a2 in
      closure_e env expr (fun expr -> ret (For (x, a1, a2, expr)))
  | Anf.Variant (name, atoms, ty) ->
      let atoms = List.map (closure_a env) atoms in
      ret (Variant (name, atoms, ty))

and closure_a (env : string_env) (a : 'a Anf.atom) : 'a atom =
  let ret = atom a.atom_meta in
  match a.atom_desc with
  | Anf.Const c -> ret (Const c)
  | Anf.Var v ->
      (if Env.mem v env then
         let label = Env.find v env in
         Hashtbl.add well_known label false);
      ret (Var v)

and closure_e_list (env : string_env) (exprs : 'a Anf.expr list)
    (k : 'a expr list -> 'a program) : 'a program =
  let rec go acc exprs =
    match exprs with
    | e :: tl -> closure_e env e (fun e -> go (e :: acc) tl)
    | [] -> k (List.rev acc)
  in
  go [] exprs

let pp_well_known fmt = function
  | true -> fprintf fmt "<well known> "
  | false -> fprintf fmt "<not well known> "

let pp_free_vars fmt free_vars =
  if List.length free_vars = 0 then ()
  else fprintf fmt "[%a] " (pp_list ~sep:", " pp_print_string) free_vars

let pp_with_meta ~show_meta pp_meta pp_desc get_meta fmt x =
  if show_meta then fprintf fmt "@[<hv>%a@ %a@]" pp_meta (get_meta x) pp_desc x
  else pp_desc fmt x

let rec pp_program ?(show_meta = false) pp_meta fmt { funcs; expr } =
  fprintf fmt "@[<v>%a@;<1 0>%a@]"
    (pp_list ~sep:"@;<1 0>" (pp_func ~show_meta pp_meta))
    funcs
    (pp_expr ~show_meta pp_meta)
    expr

and pp_func ?(show_meta = false) pp_meta fmt = function
  | Func f ->
      fprintf fmt "def %s(self, %a) %a%a{@;<1 2>%a@;<0 0>}" f.name
        (pp_list ~sep:", " pp_print_string)
        f.args pp_free_vars f.free_vars pp_well_known f.well_known
        (pp_expr ~show_meta pp_meta)
        f.body

and pp_expr ?(show_meta = false) pp_meta fmt e =
  pp_with_meta ~show_meta pp_meta
    (pp_expr_desc ~show_meta pp_meta)
    (fun e -> e.expr_meta)
    fmt e

and pp_expr_desc ?(show_meta = false) pp_meta fmt e =
  match e.expr_desc with
  | Ret cexpr -> pp_cexpr ~show_meta pp_meta fmt cexpr
  | Let (binding, expr) ->
      fprintf fmt "@[<v>@[<b>let@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]"
        (pp_binding ~show_meta pp_meta)
        binding
        (pp_expr ~show_meta pp_meta)
        expr
  | Letrec (bindings, expr) ->
      fprintf fmt "@[<v>@[<b>let rec@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]"
        (pp_list ~sep:"@;<1 0>and " (pp_binding ~show_meta pp_meta))
        bindings
        (pp_expr ~show_meta pp_meta)
        expr
  | Seq exprs ->
      fprintf fmt "@[<v>(%a)@]"
        (pp_list ~sep:";@;<1 1>" (fun fmt e ->
             fprintf fmt "@[%a@]" (pp_expr ~show_meta pp_meta) e))
        exprs

and pp_cexpr ?(show_meta = false) pp_meta fmt c =
  pp_with_meta ~show_meta pp_meta
    (pp_cexpr_desc ~show_meta pp_meta)
    (fun c -> c.cexpr_meta)
    fmt c

and pp_cexpr_desc ?(show_meta = false) pp_meta fmt c =
  let pp_c = pp_cexpr ~show_meta pp_meta in
  match c.cexpr_desc with
  | Atom atom -> pp_atom fmt atom
  | Call (atom1, atoms, label_opt) ->
      let pp_label_opt fmt = function
        | Some label -> fprintf fmt "<%s> " label
        | None -> ()
      in
      fprintf fmt "@[call %a%a(%a)@]" pp_label_opt label_opt pp_atom atom1
        (pp_list ~sep:"@;<1 2>" pp_atom)
        atoms
  | Closure (name, free_vars) ->
      fprintf fmt "@[(closure %s [%a])@]" name
        (pp_list ~sep:",@ " pp_print_string)
        free_vars
  | Prim1 (op, cexpr) -> fprintf fmt "%a %a" Prim.pp_op1 op pp_c cexpr
  | Prim2 (op, atom1, atom2) ->
      fprintf fmt "@[(%a %a@;<1 2>%a)@]" pp_atom atom1 Prim.pp_op2 op pp_atom
        atom2
  | If (atom, expr1, expr2) ->
      fprintf fmt "@[<hv>(if %a then@;<1 2>%a@;<1 1>else@;<1 2>%a)@]" pp_atom
        atom
        (pp_expr ~show_meta pp_meta)
        expr1
        (pp_expr ~show_meta pp_meta)
        expr2
  | Match (atom, patterns, _) ->
      fprintf fmt "@[<hv>@[match %a@ with@]@;<1 0>| %a@]" pp_atom atom
        (pp_list ~sep:"@;<1 0>| " (pp_pattern ~show_meta pp_meta))
        patterns
  | Tuple atoms ->
      fprintf fmt "@[(%a)@]" (pp_list ~sep:",@;<1 2>" pp_atom) atoms
  | Get (cexpr, i) -> fprintf fmt "%a.(%d)" pp_c cexpr i
  | Vect (a1, a2) -> pp_app_name "vector" pp_atom fmt [ a1; a2 ]
  | VectGet (a1, a2) -> pp_app_name "vector_get" pp_atom fmt [ a1; a2 ]
  | VectSet (a1, a2, a3) -> pp_app_name "vector_set" pp_atom fmt [ a1; a2; a3 ]
  | VectLength cexpr -> pp_app_name "vector_length" pp_c fmt [ cexpr ]
  | Ascribe (cexpr, t) ->
      fprintf fmt "@[<hv>%a@;<0 1>%a@;<1 1>:: %a@]" pp_meta c.cexpr_meta pp_c
        cexpr Type.pp t
  | AscribeDyn (cexpr, atom, inv) ->
      fprintf fmt "@[<hv>(%a@;<1 1>:: %a(%a))@]" pp_c cexpr Inversion.pp inv
        pp_atom atom
  | For (x, a1, a2, expr) ->
      fprintf fmt "@[<v>@[for %s = %a to %a do@]@;<1 2>%a@;<1 0>done@]" x
        pp_atom a1 pp_atom a2
        (pp_expr ~show_meta pp_meta)
        expr
  | Variant (name, atoms, _) ->
      if List.length atoms = 0 then fprintf fmt "%s" name
      else fprintf fmt "@[%s (@[%a@])@]" name (pp_list ~sep:",@ " pp_atom) atoms

and pp_atom fmt a =
  match a.atom_desc with
  | Const const -> fprintf fmt "%s" (Const.to_string const)
  | Var x -> fprintf fmt "%s" x

and pp_binding ?(show_meta = false) pp_meta fmt binding =
  fprintf fmt "%s =@;<1 2>%a" binding.name
    (pp_cexpr ~show_meta pp_meta)
    binding.cexpr

and pp_pattern ?(show_meta = false) pp_meta fmt pattern =
  fprintf fmt "@[%s %a ->@ %a@]" pattern.casename
    (pp_list ~sep:" " pp_print_string)
    pattern.args
    (pp_expr ~show_meta pp_meta)
    pattern.body
