open Common
open Pp

type ('term, 'meta) annotated = { term : 'term; meta : 'meta }

let term annotated = annotated.term
let meta annotated = annotated.meta

let annotate (term : 'term) (meta : 'meta) : ('term, 'meta) annotated =
  { term; meta }

module Type = struct
  type 'a t = ('a term, 'a) annotated

  and 'a term =
    | Any
    | Int
    | Bool
    | Float
    | Arrow of 'a t list * 'a t
    | Product of 'a t list
    | Vec of 'a t
    | Var of string
    | Variant of 'a case list

  and 'a case = Case of string * 'a t list

  let rec map_meta (f : 'a -> 'b) (ty : 'a t) : 'b t =
    let meta = f ty.meta in
    let term =
      match ty.term with
      | Any -> Any
      | Int -> Int
      | Bool -> Bool
      | Float -> Float
      | Arrow (tys, ty) ->
          let tys' = List.map (map_meta f) tys in
          let ty' = map_meta f ty in
          Arrow (tys', ty')
      | Product ts -> Product (List.map (map_meta f) ts)
      | Vec t -> Vec (map_meta f t)
      | Var type_name -> Var type_name
      | Variant cases ->
          Variant
            (List.map
               (fun (Case (name, tys)) ->
                 Case (name, List.map (map_meta f) tys))
               cases)
    in
    annotate term meta

  let rec pp ~pp_meta fmt ty =
    let pp = pp ~pp_meta in
    let pp_case = pp_case ~pp_meta in
    let pp_term fmt ty =
      match ty.term with
      | Int -> Format.fprintf fmt "int"
      | Bool -> Format.fprintf fmt "bool"
      | Float -> Format.fprintf fmt "float"
      | Any -> Format.fprintf fmt "?"
      | Arrow (tys, ty) ->
          Format.fprintf fmt "(%a -> %a)"
            (Common.Pp.pp_list ~sep:" -> " pp)
            tys pp ty
      | Product l ->
          Format.fprintf fmt "(%a)" (Common.Pp.pp_list ~sep:" * " pp) l
      | Vec t -> Format.fprintf fmt "vec[%a]" pp t
      | Var type_name -> Format.fprintf fmt "%s" type_name
      | Variant cases ->
          Format.fprintf fmt "(%a)" (Common.Pp.pp_list ~sep:" | " pp_case) cases
    in
    Format.fprintf fmt "%a%a" pp_term ty pp_meta ty.meta

  and pp_case ~pp_meta fmt = function
    | Case (name, tys) ->
        if tys = [] then Format.fprintf fmt "%s" name
        else
          Format.fprintf fmt "%s of %a" name
            (Common.Pp.pp_list ~sep:" * " (pp ~pp_meta))
            tys
end

module Const = struct
  type t = Int of int | Bool of bool | Float of float

  let pp fmt = function
    | Int n -> Format.fprintf fmt "%d" n
    | Bool b -> Format.fprintf fmt "%b" b
    | Float f -> Format.fprintf fmt "%f" f
end

module Expr = struct
  type 'a t = ('a expr, 'a) annotated

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

  and 'a binding = ('a binding_term, 'a) annotated
  and 'a binding_term = { name : string; ty : 'a Type.t option; e : 'a t }
  and 'a pattern = ('a pattern_term, 'a) annotated
  and 'a pattern_term = { case_name : string; args : string list; body : 'a t }

  let make_binding_term name e ty : 'a binding_term = { name; ty; e }

  let make_binding name e ty meta : 'a binding =
    annotate (make_binding_term name e ty) meta

  let make_pattern_term case_name args body : 'a pattern_term =
    { case_name; args; body }

  let make_pattern case_name args body meta : 'a pattern =
    annotate (make_pattern_term case_name args body) meta

  let rec map_ty (f : 'a Type.t -> 'a Type.t) (e : 'a t) : 'a t =
    let term =
      match e.term with
      | Const c -> Const c
      | Var x -> Var x
      | Lam (args, e) ->
          Lam (List.map (fun (x, ty) -> (x, f ty)) args, map_ty f e)
      | App (exprs, e) -> App (map_ty f exprs, List.map (map_ty f) e)
      | Prim1 (op1, e) -> Prim1 (op1, map_ty f e)
      | Prim2 (op2, e1, e2) -> Prim2 (op2, map_ty f e1, map_ty f e2)
      | If (e1, e2, e3) -> If (map_ty f e1, map_ty f e2, map_ty f e3)
      | Let (binding, e) ->
          Let
            ( annotate
                {
                  binding.term with
                  ty = Option.map f binding.term.ty;
                  e = map_ty f binding.term.e;
                }
                (meta binding),
              map_ty f e )
      | Letrec (bindings, e) ->
          Letrec
            ( List.map
                (fun binding ->
                  annotate
                    {
                      binding.term with
                      ty = Option.map f binding.term.ty;
                      e = map_ty f binding.term.e;
                    }
                    (meta binding))
                bindings,
              map_ty f e )
      | Tuple es -> Tuple (List.map (map_ty f) es)
      | Get (e, i) -> Get (map_ty f e, i)
      | Vect (e1, e2) -> Vect (map_ty f e1, map_ty f e2)
      | VectGet (e1, e2) -> VectGet (map_ty f e1, map_ty f e2)
      | VectSet (e1, e2, e3) -> VectSet (map_ty f e1, map_ty f e2, map_ty f e3)
      | VectLength e -> VectLength (map_ty f e)
      | Ascribe (e, ty) -> Ascribe (map_ty f e, f ty)
      | For (x, e1, e2, e3) -> For (x, map_ty f e1, map_ty f e2, map_ty f e3)
      | Variant (name, es) -> Variant (name, List.map (map_ty f) es)
      | Match (e, patterns) ->
          Match
            ( map_ty f e,
              List.map
                (fun p ->
                  make_pattern p.term.case_name p.term.args
                    (map_ty f p.term.body) p.meta)
                patterns )
      | Seq es -> Seq (List.map (map_ty f) es)
    in
    annotate term (meta e)

  let rec map_meta (f : 'a -> 'b) (e : 'a t) : 'b t =
    let meta = f e.meta in
    let term =
      match e.term with
      | Const c -> Const c
      | Var x -> Var x
      | Lam (args, e) ->
          Lam
            ( List.map (fun (x, ty) -> (x, Type.map_meta f ty)) args,
              map_meta f e )
      | App (exprs, e) -> App (map_meta f exprs, List.map (map_meta f) e)
      | Prim1 (op1, e) -> Prim1 (op1, map_meta f e)
      | Prim2 (op2, e1, e2) -> Prim2 (op2, map_meta f e1, map_meta f e2)
      | If (e1, e2, e3) -> If (map_meta f e1, map_meta f e2, map_meta f e3)
      | Let (binding, e) ->
          let meta = f binding.meta in
          let term =
            {
              binding.term with
              ty = Option.map (Type.map_meta f) binding.term.ty;
              e = map_meta f binding.term.e;
            }
          in
          let binding' = annotate term meta in
          Let (binding', map_meta f e)
      | Letrec (bindings, e) ->
          Letrec
            ( List.map
                (fun binding ->
                  let meta = f binding.meta in
                  let term =
                    {
                      binding.term with
                      ty = Option.map (Type.map_meta f) binding.term.ty;
                      e = map_meta f binding.term.e;
                    }
                  in
                  annotate term meta)
                bindings,
              map_meta f e )
      | Tuple es -> Tuple (List.map (map_meta f) es)
      | Get (e, i) -> Get (map_meta f e, i)
      | Vect (e1, e2) -> Vect (map_meta f e1, map_meta f e2)
      | VectGet (e1, e2) -> VectGet (map_meta f e1, map_meta f e2)
      | VectSet (e1, e2, e3) ->
          VectSet (map_meta f e1, map_meta f e2, map_meta f e3)
      | VectLength e -> VectLength (map_meta f e)
      | Ascribe (e, ty) -> Ascribe (map_meta f e, Type.map_meta f ty)
      | For (x, e1, e2, e3) ->
          For (x, map_meta f e1, map_meta f e2, map_meta f e3)
      | Variant (name, es) -> Variant (name, List.map (map_meta f) es)
      | Match (e, patterns) ->
          Match
            ( map_meta f e,
              List.map
                (fun p ->
                  let meta = f p.meta in
                  make_pattern p.term.case_name p.term.args
                    (map_meta f p.term.body) meta)
                patterns )
      | Seq es -> Seq (List.map (map_meta f) es)
    in
    annotate term meta

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

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

  and pp_binding ~pp_meta fmt binding =
    let ty_str =
      match binding.term.ty with
      | Some ty -> Format.asprintf " : %a" (Type.pp ~pp_meta) ty
      | None -> ""
    in
    Format.fprintf fmt "%s%s =@;<1 2>%a" binding.term.name ty_str (pp ~pp_meta)
      binding.term.e

  and pp_pattern ~pp_meta fmt pattern =
    Format.fprintf fmt "@[%s %a ->@ %a@]" pattern.term.case_name
      (pp_list ~sep:" " Format.pp_print_string)
      pattern.term.args (pp ~pp_meta) pattern.term.body
end

module Decl = struct
  type 'a t = ('a decl, 'a) annotated
  and 'a decl = DType of string * 'a Type.t

  let map_ty (f : 'a Type.t -> 'a Type.t) (decl : 'a t) : 'a t =
    let term =
      match decl.term with DType (type_name, ty) -> DType (type_name, f ty)
    in
    annotate term (meta decl)

  let map_meta (f : 'a -> 'b) (decl : 'a t) : 'b t =
    let term =
      match decl.term with
      | DType (type_name, ty) -> DType (type_name, Type.map_meta f ty)
    in
    annotate term (f decl.meta)

  let pp ~pp_meta fmt decl =
    match decl.term with
    | DType (type_name, ty) ->
        Format.fprintf fmt "type %s = %a%a" type_name (Type.pp ~pp_meta) ty
          pp_meta decl.meta
end

module Program = struct
  type 'a t = ('a program, 'a) annotated
  and 'a program = { decls : 'a Decl.t list; expr : 'a Expr.t }

  let map_ty (f : 'a Type.t -> 'a Type.t) (program : 'a t) : 'a t =
    let term =
      {
        decls = List.map (Decl.map_ty f) program.term.decls;
        expr = Expr.map_ty f program.term.expr;
      }
    in
    annotate term (meta program)

  let map_meta (f : 'a -> 'b) (program : 'a t) : 'b t =
    let meta = f program.meta in
    let term =
      {
        decls = List.map (Decl.map_meta f) program.term.decls;
        expr = Expr.map_meta f program.term.expr;
      }
    in
    annotate term meta

  let pp ~pp_meta fmt program =
    Format.fprintf fmt "@[<v>%a@]@.%a%a"
      (Common.Pp.pp_list ~sep:"@ " (Decl.pp ~pp_meta))
      program.term.decls (Expr.pp ~pp_meta) program.term.expr pp_meta
      program.meta
end
