open Common
open Syntax
open Type
open Typed_ast
module Env = Map.Make (String)

let err msg = Error.(error (Ty (IncompatibleType msg)))

(* let rec type_subst ty1 ty2 ty =
   if ty = ty1 then ty2 else
   match ty with
   | Int
   | Bool
   | Any
   | Var _ -> ty
   | Arrow (ty1', ty2') -> Arrow (type_subst ty1 ty2 ty1', type_subst ty1 ty2 ty2')
   | Product tys -> Product (List.map (type_subst ty1 ty2) tys)
   | Vec ty1' -> Vec (type_subst ty1 ty2 ty1')
   | Variant cases -> Variant (List.map (fun (Case (name, tys)) -> Case (name, List.map (type_subst ty1 ty2) tys)) cases)
   | Rec (name, ty1') -> Rec (name, type_subst ty1 ty2 ty1')
*)
(* let unfold rec_ty ty =
   let name, ty' = match rec_ty with
   | Rec (name, ty') -> name, ty'
   | _ -> failwith "expected recursive type as argument to unfold"
   in
   match ty with
   | Rec (name', ty'') ->
     type_subst (Type.Var name') ty' ty''
   | _ -> err "expected recursive type" *)

(* let fold rec_ty ty =
   let name, ty' = match rec_ty with
   | Rec (name, ty') -> name, ty'
   | _ -> failwith "expected recursive type as argument to fold"
   in
   Rec (name, type_subst ty' (Type.Var name) ty) *)

let rec typecheck program =
  let env =
    program.decls
    |> List.fold_left
         (fun env (DType (name, ty)) -> Env.add name ty env)
         Env.empty
  in
  ignore (typeof env program.expr)

and typeof env e =
  match e.desc with
  | Const c -> Const.type_of c
  | Var x -> (
      match Env.find_opt x env with
      | Some ty -> ty
      | None -> raise (err (Format.asprintf "unbound variable '%s'" x)))
  | Lam (args, e) ->
      let env' = Env.add_seq (List.to_seq args) env in
      let _, tys = List.split args in
      let ty' = typeof env' e in
      Arrow (tys, ty')
  | App (e, exprs) -> (
      let ty = typeof env e in
      let tys = List.map (typeof env) exprs in
      match ty with
      | Arrow (a, b) when a = tys -> b
      | Arrow (tys', _) ->
          raise
            (err
               (Format.asprintf
                  "type mismatch, expected function type with arguments of \
                   types '%a', but got function type with arguments of types \
                   '%a'"
                  (Pp.pp_list ~sep:", " Type.pp)
                  tys
                  (Pp.pp_list ~sep:", " Type.pp)
                  tys'))
      | _ ->
          raise
            (err
               (Format.asprintf "expected function type, but got '%a'" Type.pp
                  ty)))
  | Prim1 (op, e1) -> (
      let a, b = ty_of_op1 op in
      match typeof env e1 with
      | ty when ty = a -> b
      | ty ->
          raise
            (err
               (Format.asprintf "expected argument of type '%a', but got '%a'"
                  Type.pp a Type.pp ty)))
  | Prim2 (op, e1, e2) ->
      let ty1 = typeof env e1 in
      let ty2 = typeof env e2 in
      let a, b, c = ty_of_op2 op in
      if ty1 = a && ty2 = b then c
      else
        raise
          (err
             (Format.asprintf
                "expected arguments of types '%a' and '%a', but got '%a' and \
                 '%a'"
                Type.pp a Type.pp b Type.pp ty1 Type.pp ty2))
  | If (e1, e2, e3) -> (
      let ty1 = typeof env e1 in
      let ty2 = typeof env e2 in
      let ty3 = typeof env e3 in
      match (ty1, ty2, ty3) with
      | Bool, a, b when a = b -> a
      | Bool, _, _ ->
          raise
            (err
               (Format.asprintf
                  "expected both branches of if expression to have the same \
                   type, but got '%a' and '%a'"
                  Type.pp ty2 Type.pp ty3))
      | _, _, _ ->
          raise
            (err
               (Format.asprintf
                  "expected condition of if expression to be of type 'bool', \
                   but got '%a'"
                  Type.pp ty1)))
  | Let (binding, e) -> (
      match typeof env binding.e with
      | ty when ty = binding.ty -> typeof (Env.add binding.name ty env) e
      | ty ->
          raise
            (err
               (Format.asprintf "expected '%a', but got '%a'" Type.pp binding.ty
                  Type.pp ty)))
  | Letrec (bindings, e) ->
      let env' =
        bindings |> List.fold_left (fun env b -> Env.add b.name b.ty env) env
      in
      let _ =
        bindings
        |> List.map (fun b ->
               match typeof env' b.e with
               | ty when ty = b.ty -> ty
               | ty ->
                   raise
                     (err
                        (Format.asprintf "expected '%a', but got '%a'" Type.pp
                           b.ty Type.pp ty)))
      in
      typeof env' e
  | Tuple l -> Type.Product (List.map (typeof env) l)
  | Get (e, i) -> (
      match typeof env e with
      | Product types when i < List.length types -> List.nth types i
      | Product types ->
          let len = List.length types in
          raise
            (err
               (Format.asprintf
                  "tuple access out of bounds, tried to access the element at \
                   position %d in a tuple of length %d"
                  i len))
      | ty ->
          raise
            (err
               (Format.asprintf "expected array type, but got '%a'" Type.pp ty))
      )
  | Vect (e1, e2) -> (
      let ty1 = typeof env e1 in
      let ty2 = typeof env e2 in
      match ty1 with
      | Type.Int -> Type.Vec ty2
      | _ ->
          raise
            (err
               (Format.asprintf
                  "expected size of type '%a', but got expression of type '%a'"
                  Type.pp Type.Int Type.pp ty1)))
  | VectGet (e1, e2) -> (
      let ty1 = typeof env e1 in
      let ty2 = typeof env e2 in
      match (ty1, ty2) with
      | Type.Vec ty, Type.Int -> ty
      | _, Type.Int ->
          raise
            (err
               (Format.asprintf
                  "expected expression of type 'Vec _', but got expression of \
                   type '%a'"
                  Type.pp ty1))
      | _ ->
          raise
            (err
               (Format.asprintf
                  "expected index of type '%a', but got expression of type '%a'"
                  Type.pp Type.Int Type.pp ty2)))
  | VectSet (e1, e2, e3) -> (
      let ty1 = typeof env e1 in
      let ty2 = typeof env e2 in
      let ty3 = typeof env e3 in
      match (ty1, ty2, ty3) with
      | Type.Vec ty, Type.Int, ty' ->
          if ty = ty' then ty1
          else
            raise
              (err
                 (Format.asprintf
                    "expected value of type '%a', but got expression of type \
                     '%a"
                    Type.pp ty Type.pp ty'))
      | Type.Vec _, _, _ ->
          raise
            (err
               (Format.asprintf
                  "expected value of type '%a', but got expression of type '%a'"
                  Type.pp Type.Int Type.pp ty2))
      | _, _, _ ->
          raise
            (err
               (Format.asprintf
                  "expected value of type 'vec _', but got expression of type \
                   '%a'"
                  Type.pp ty1)))
  | VectLength e -> (
      match typeof env e with
      | Type.Vec _ -> Type.Int
      | ty ->
          raise
            (err
               (Format.asprintf
                  "expected expression of type 'vec _', but got expression of \
                   type '%a'"
                  Type.pp ty)))
  | Ascribe (e, ty') -> (
      match typeof env e with
      | ty when ty = ty' -> ty
      | ty ->
          raise
            (err
               (Format.asprintf
                  "expected expression of type '%a', but got expression of \
                   type '%a'"
                  Type.pp ty' Type.pp ty)))
  | For (x, e1, e2, e3) ->
      let ty1 = typeof env e1 in
      let ty2 = typeof env e2 in
      if ty1 = Type.Int && ty2 = Type.Int then
        let env' = Env.add x Type.Int env in
        let _ty3 = typeof env' e3 in
        Type.Int
      else
        raise
          (err
             (Format.asprintf
                "expected bounds of type 'int', but got expressions of type \
                 '%a' and '%a'"
                Type.pp ty1 Type.pp ty2))
  | Variant (case_name, exprs, ty) -> (
      let tys = List.map (typeof env) exprs in
      let cases =
        match ty with
        | Rec (ty_name, Variant cases) ->
            cases
            |> List.map (fun (Case (name, tys)) ->
                   Case (name, List.map (Type.subst_ty ty_name ty) tys))
        | _ ->
            raise
              (err
                 (Format.asprintf "expected variant type, but got '%a'" Type.pp
                    ty))
      in
      match List.find_opt (fun (Case (name, _)) -> name = case_name) cases with
      | Some (Case (name, case_tys)) ->
          if case_tys = tys then ty
          else
            raise
              (err
                 (Format.asprintf
                    "In the variant case %s, expected arguments to have types \
                     '%a', but got arguments whose types are '%a'"
                    name
                    (Pp.pp_list ~sep:", " Type.pp)
                    case_tys
                    (Pp.pp_list ~sep:", " Type.pp)
                    tys))
      | None ->
          raise (err (Format.asprintf "variant case '%s' not found" case_name)))
  | Match (e, patterns, ty) ->
      let ty' = typeof env e in
      if ty' <> ty then
        raise
          (err
             (Format.asprintf
                "expected expression being matched to be of type %a, but got \
                 an expression of type %a"
                Type.pp ty Type.pp ty'));
      (* now, check that all the pattern bodies have the same type *)
      let cases =
        match ty with
        | Rec (ty_name, Variant cases) ->
            cases
            |> List.map (fun (Case (name, tys)) ->
                   Case (name, List.map (Type.subst_ty ty_name ty) tys))
        | _ ->
            raise
              (err
                 (Format.asprintf "expected variant type, but got '%a'" Type.pp
                    ty))
      in
      let typeof_pattern (pattern : pattern) =
        let case : Type.case =
          List.find (fun (Type.Case (name, _)) -> name = pattern.casename) cases
        in
        let (Type.Case (_, tys)) = case in
        if List.length pattern.args <> List.length tys then
          raise
            (err
               (Format.asprintf
                  "expected pattern '%s' to have %d arguments, but got %d"
                  pattern.casename (List.length tys) (List.length pattern.args)));
        let env' =
          List.combine pattern.args tys
          |> List.fold_left (fun env (arg, ty) -> Env.add arg ty env) env
        in
        typeof env' pattern.body
      in
      let tys' = patterns |> List.map typeof_pattern in
      let first_ty =
        match tys' with
        | [] -> raise (err "no patterns provided in match expression")
        | ty :: _ -> ty
      in
      tys'
      |> List.iteri (fun i ty ->
             if ty <> List.hd tys' then
               raise
                 (err
                    (Format.asprintf
                       "expected all pattern bodies to have the same type, but \
                        pattern %d has type '%a', while the first pattern has \
                        type '%a'"
                       i Type.pp ty Type.pp (List.hd tys'))));
      (* finally, return the type of the pattern bodies *)
      first_ty
  | Seq exprs -> typeof env (List.hd (List.rev exprs))
