open Common
open Pp

type program = { decls : decl list; expr : e }
and decl = DType of string * Type.t
and e = { desc : desc; loc : Loc.t }

and desc =
  | Const of Const.t
  | Var of string
  | Lam of (string * Type.t) list * e
  | App of e * e list
  | Prim1 of Prim.op1 * e
  | Prim2 of Prim.op2 * e * e
  | If of e * e * e
  | Let of binding * e
  | Letrec of binding list * e
  | Tuple of e list
  | Get of e * int
  | Vect of e * e
  | VectGet of e * e
  | VectSet of e * e * e
  | VectLength of e
  | Ascribe of e * Type.t
  | For of string * e * e * e
  | Variant of string * e list * Type.t
  | Match of e * pattern list * Type.t
  | Seq of e list

and binding = { name : string; ty : Type.t; e : e; binding_loc : Loc.t }

and pattern = {
  casename : string;
  args : string list;
  body : e;
  pattern_loc : Loc.t;
}

let pp_arg fmt (x, ty) = Format.fprintf fmt "(%s : %a)" x Type.pp ty

let rec pp_program fmt p =
  Format.fprintf fmt "@[<v>%a@;<1 0>%a@]"
    (pp_list ~sep:"@ " pp_decl)
    p.decls pp_e p.expr

and pp_decl fmt = function
  | DType (name, ty) -> Format.fprintf fmt "type %s = %a" name Type.pp ty

and pp_e fmt e =
  let open Format in
  match e.desc with
  | Const const -> fprintf fmt "%s" (Const.to_string const)
  | Var x -> fprintf fmt "%s" x
  | Lam (args, e) ->
      fprintf fmt "@[(fun %a ->@;<1 2>%a)@]" (pp_list ~sep:" " pp_arg) args pp_e
        e
  | App (e, exprs) -> pp_app pp_e fmt (e :: exprs)
  | Prim1 (op1, e) -> fprintf fmt "%a %a" Prim.pp_op1 op1 pp_e e
  | Prim2 (op2, e1, e2) ->
      fprintf fmt "@[(%a %a@;<1 2>%a)@]" pp_e e1 Prim.pp_op2 op2 pp_e e2
  | If (e1, e2, e3) ->
      fprintf fmt "@[<hv>(if %a then@;<1 2>%a@;<1 1>else@;<1 2>%a)@]" pp_e e1
        pp_e e2 pp_e e3
  | Let (binding, e) ->
      fprintf fmt "@[<v>@[<b>let@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]" pp_binding
        binding pp_e e
  | Letrec (bindings, e) ->
      fprintf fmt "@[<v>@[<b>let rec@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]"
        (pp_list ~sep:"@;<1 0>and " pp_binding)
        bindings pp_e e
  | Tuple exprs -> fprintf fmt "@[(%a)@]" (pp_list ~sep:",@;<1 2>" pp_e) exprs
  | Get (e, i) -> fprintf fmt "%a.(%d)" pp_e e i
  | Vect (e1, e2) -> pp_app_name "vector" pp_e fmt [ e1; e2 ]
  | VectGet (e1, e2) -> pp_app_name "vector_get" pp_e fmt [ e1; e2 ]
  | VectSet (e1, e2, e3) -> pp_app_name "vector_set" pp_e fmt [ e1; e2; e3 ]
  | VectLength e -> pp_app_name "vector_length" pp_e fmt [ e ]
  | Ascribe (e', t) -> fprintf fmt "@[<hv>(%a@;<1 1>:: %a)@]" pp_e e' Type.pp t
  | For (x, e1, e2, e3) ->
      fprintf fmt "@[<v>@[for %s = %a to %a do@]@;<1 2>%a@;<1 0>done@]" x pp_e
        e1 pp_e e2 pp_e e3
  | Variant (name, exprs, _) ->
      if List.length exprs = 0 then fprintf fmt "%s" name
      else fprintf fmt "@[%s (@[%a@])@]" name (pp_list ~sep:",@ " pp_e) exprs
  | Match (e, patterns, _) ->
      fprintf fmt "@[<hv>@[match %a@ with@]@;<1 0>| %a@]" pp_e e
        (pp_list ~sep:"@;<1 0>| " pp_pattern)
        patterns
  | Seq exprs ->
      fprintf fmt "@[<v>(%a)@]"
        (pp_list ~sep:";@;<1 1>" (fun fmt e -> fprintf fmt "@[%a@]" pp_e e))
        exprs

and pp_binding fmt binding =
  Format.fprintf fmt "%s : %a =@;<1 2>%a" binding.name Type.pp binding.ty pp_e
    binding.e

and pp_pattern fmt pattern =
  Format.fprintf fmt "@[%s %a ->@ %a@]" pattern.casename
    (pp_list ~sep:" " Format.pp_print_string)
    pattern.args pp_e pattern.body

let of_parse_tree (program : 'a Parse_tree.Program.t) =
  let case_env = Hashtbl.create 16 in
  let lookup_case_ty name =
    match Hashtbl.find_opt case_env name with
    | Some ty -> ty
    | None -> failwith (Format.sprintf "unbound variant case name %s" name)
  in
  let rec go_ty (ty : 'a Parse_tree.Type.t) : Type.t =
    match ty.term with
    | Parse_tree.Type.Any -> Type.Any
    | Parse_tree.Type.Int -> Type.Int
    | Parse_tree.Type.Bool -> Type.Bool
    | Parse_tree.Type.Float -> Type.Float
    | Parse_tree.Type.Arrow (tys, ty) ->
        Type.Arrow (List.map go_ty tys, go_ty ty)
    | Parse_tree.Type.Product tys -> Type.Product (List.map go_ty tys)
    | Parse_tree.Type.Vec ty -> Type.Vec (go_ty ty)
    | Parse_tree.Type.Var id -> Type.Var id
    | Parse_tree.Type.Variant cases -> Type.Variant (List.map go_case cases)
  and go_case (Parse_tree.Type.Case (name, tys)) =
    Type.Case (name, List.map go_ty tys)
  in
  let go_const = function
    | Parse_tree.Const.Int i -> Const.Int i
    | Parse_tree.Const.Bool b -> Const.Bool b
    | Parse_tree.Const.Float f -> Const.Float f
  in
  let go_decl decl =
    let (Parse_tree.Decl.DType (name, ty)) = decl.Parse_tree.term in
    let ty' = Type.Rec (name, go_ty ty) in
    match ty.term with
    | Parse_tree.Type.Variant cases ->
        cases
        |> List.iter (fun (Parse_tree.Type.Case (case_name, _)) ->
               Hashtbl.add case_env case_name ty');
        DType (name, ty')
    | _ -> failwith "expected declared type to be a variant type"
  in
  let open Parse_tree in
  let rec go_e (e : 'a Expr.t) =
    let desc =
      match e.term with
      | Const const -> Const (go_const const)
      | Var x -> Var x
      | Lam (args, e) ->
          Lam (List.map (fun (x, ty) -> (x, go_ty ty)) args, go_e e)
      | App (e, exprs) -> App (go_e e, List.map go_e exprs)
      | Prim1 (op1, e) -> Prim1 (op1, go_e e)
      | Prim2 (op2, e1, e2) -> Prim2 (op2, go_e e1, go_e e2)
      | If (e1, e2, e3) -> If (go_e e1, go_e e2, go_e e3)
      | Let (binding, e) -> Let (go_binding binding, go_e e)
      | Letrec (bindings, e) -> Letrec (List.map go_binding bindings, go_e e)
      | Tuple exprs -> Tuple (List.map go_e exprs)
      | Get (e, i) -> Get (go_e e, i)
      | Vect (e1, e2) -> Vect (go_e e1, go_e e2)
      | VectGet (e1, e2) -> VectGet (go_e e1, go_e e2)
      | VectSet (e1, e2, e3) -> VectSet (go_e e1, go_e e2, go_e e3)
      | VectLength e -> VectLength (go_e e)
      | Ascribe (e, ty) -> Ascribe (go_e e, go_ty ty)
      | For (x, e1, e2, e3) -> For (x, go_e e1, go_e e2, go_e e3)
      | Seq exprs -> Seq (List.map go_e exprs)
      | Variant (name, exprs) ->
          let ty = lookup_case_ty name in
          Variant (name, List.map go_e exprs, ty)
      | Match (e, patterns) ->
          let ty = lookup_case_ty (List.hd patterns).term.case_name in
          Match (go_e e, List.map go_pattern patterns, ty)
    in
    { desc; loc = e.meta }
  and go_binding binding =
    let ty_opt = Option.map go_ty binding.term.ty in
    let ty = Option.value ty_opt ~default:Any in
    {
      name = binding.term.name;
      ty;
      e = go_e binding.term.e;
      binding_loc = binding.meta;
    }
  and go_pattern pattern =
    {
      casename = pattern.term.case_name;
      args = pattern.term.args;
      body = go_e pattern.term.body;
      pattern_loc = pattern.meta;
    }
  in
  let go_program (program : 'a Program.t) =
    let decls = List.map go_decl program.term.decls in
    let expr = go_e program.term.expr in
    { decls; expr }
  in
  go_program program
