(* λ calculus, with Forall and simple types *)

type 'a binding =
  | SymBind of 'a
  | QuantifierBind

type 'a context = (string * 'a binding) list

type typ =
  | UnitTy
  | IntTy
  | Arrow of typ * typ
  | Forall of string * typ
  | ClosureTy of (typ context) * typ
  | SymbolTy of string

type expr =
  | Closure of (expr context) * expr
  | Unit
  | Symbol of string
  | Int of int
  | Lam of string * typ * expr  (* turns into Arrow. Takes expr as argument, returns expr *)
  | App of expr * expr
  | TLam of string * expr       (* expr version of Forall. Takes typ as argument, returns expr *)
  | TApp of expr * typ

let rec lookup context key =
  match context with
  | [] -> Error "No result"
  | (k, QuantifierBind) :: rest -> lookup rest key
  | (k, SymBind v) :: rest -> if key = k
                              then Ok v
                              else lookup rest key

let str = String.concat ""

let rec string_of_typ t = match t with
  | UnitTy -> "UnitTy"
  | IntTy -> "IntTy"
  | Arrow _ -> "Arrow"
  | Forall (s, t) -> str ["Forall "; s; " "; string_of_typ t]
  | ClosureTy (_ctx, t) -> str ["Closure "; string_of_typ t]
  | SymbolTy s -> str ["SymbolTy "; s]

let string_of_expr e = match e with
  | Unit -> "Unit"
  | Closure _ -> "Closure"
  | Symbol s -> str ["Symbol "; s]
  | Int i -> str ["Int "; string_of_int i]
  | Lam _ -> "Lam"
  | App _ -> "App"
  | TLam _ -> "TLam"
  | TApp _ -> "TApp"

let rec convert_context converter_fn original_context exprcontext acc =
  (* カリたべたい *)
  let convert_context = convert_context converter_fn original_context in
  match exprcontext with
  | [] -> acc
  | (x, QuantifierBind) :: rest ->
     convert_context rest ((x, QuantifierBind) :: acc)
  | (x, SymBind e) :: rest ->
     (match converter_fn original_context e with
      | Ok t -> convert_context rest ((x, SymBind t) :: acc)
      | Error e -> failwith (str ["typ_context error: "; e]))

let rec typ_of_expr context e =
  match e with
  | Unit -> Ok UnitTy
  | Int i -> Ok IntTy
  | Closure (ctx, e) ->
     let new_context = convert_context typ_of_expr context ctx context in
     typ_of_expr new_context e
  | Symbol s -> lookup context s
  | Lam (x, in_typ, body) ->
     (let new_context = (x, SymBind in_typ) :: context in
      match typ_of_expr new_context body with
      | Ok return_typ -> Ok (Arrow (in_typ, return_typ))
      | Error e -> Error (str [ "Error finding return type of lambda. Error: "
                              ; e]))
  | App (e0,
         e1) ->
     (match (typ_of_expr context e0,
             typ_of_expr context e1) with
      | (Ok (Arrow (t0_0, t0_1)),
         Ok t1) -> if t1 = t0_0
                   then Ok t0_1
                   else Error (str ["Type mismatch"])
      | (Ok (ClosureTy (ctx, Arrow (t0_0, t0_1))),
         Ok t1) -> if t1 = t0_0
                   then Ok t0_1
                   else Error
                          (str
                             ["Type mismatch. t0_0: "
                             ;string_of_typ t0_0
                             ;", t0_1: "
                             ;string_of_typ t0_1
                             ;", t1: "
                             ;string_of_typ t1
                             ])
      | (Ok t, _) -> Error (str ["Not given a lambda as first thing to App. "
                                ;string_of_typ t])
      | _ -> Error (str ["Error getting typ of expr App"]))
  | TLam (a, body) ->
     (let new_context = (a, QuantifierBind) :: context in
      match typ_of_expr new_context body with
      | Ok body_type -> Ok (ClosureTy (new_context,
                                       Forall (a,
                                               body_type)))
      | Error e -> Error (str [ "Failed to get typ of TLam, and thus can not construct Forall. Error: "
                              ; e]))
  | TApp (TLam (a, body),
          t1) -> let new_context = (a, QuantifierBind) :: context in
                 typ_of_expr new_context body
  | TApp (e0, t1) ->
     (match typ_of_expr context e0 with
      | Ok (ClosureTy (ctx,
                       Forall (a, body))) ->
         let new_context = (a, SymBind t1) :: ctx @ context in
         Ok (ClosureTy (new_context,
                        body))
      | Ok (ClosureTy(ctx,
                      t)) ->
         Error (str ["TApp given ClosureTy with non-forall on left-hand-side. IS: "
                    ;string_of_typ t])
      | Ok (Forall (a, body)) -> Error "TApp given naked Forall as first argument"
      | Ok _ -> Error "TApp given non-Forall first argument"
      | Error e -> Error (str ["Error in TApp - error: "
                              ;e]))

let rec eval context e =
    let context: (expr context) = context in
    match e with
    | Unit -> Ok Unit
    | Closure (ctx, e) -> eval (List.append ctx context) e
    | Int i -> Ok (Int i)
    | Lam (x, t, body) -> Ok (Closure (context,
                                       Lam (x, t, body)))
    | Symbol s -> lookup context s
    | App (Closure(ctx,
                   Lam (x, t, body)),
           e1) -> let new_context = (x, SymBind e1) :: ctx @ context in
                  eval new_context body
    | App (e0, e1) -> (match eval context e0 with
                       | Ok (Closure (ctx,
                                      Lam (x, t, body))) ->
                          eval (ctx @ context) (App (Lam (x,
                                                  t,
                                                  body),
                                             e1))
                       | Ok _ -> Error "Can't apply non-Lam"
                       | Error e -> Error (str ["Apply error: "
                                               ;e]))
    | TLam (a, body) -> Ok (TLam (a, body))
    | TApp (TLam (a, body),
            given_typ) -> let new_context = (a, QuantifierBind) :: context in
                          eval new_context body
    | TApp (e0, t) -> (match eval context e0 with
                       | Ok (TLam (a, body)) -> let _ = failwith "Not gonna happen" in
                                                eval context (TApp (TLam (a, body), t))
                       | Ok (Closure (ctx,
                                      TLam (a,
                                            body))) ->
                          eval (ctx @ context) body
                       | Ok _ -> Error "Can't type-apply non-Λ"
                       | Error e -> Error (str [ "Error applying type lambda. Error: "
                                               ; e]))

let tlam a body = TLam (a, body)
let lam x t body = Lam (x, t, body)
let app e0 e1 = App (e0, e1)
let tapp e0 tyT1 = TApp (e0, tyT1)
;; eval [] (tlam "a" (lam "x"
                        (SymbolTy "a")
                        (Symbol "x")))

let k_combinator = (tlam "a" (tlam "_b"
                                (lam
                                   "x"
                                   (SymbolTy "a")
                                   (lam
                                      "_y"
                                      (SymbolTy "_b")
                                      (Symbol "x")))))

let k_intint = (TApp (TApp (k_combinator,
                            IntTy),
                      IntTy))

(* forall a b . a -> b -> a *)
;; typ_of_expr [] (app k_intint (Int 1))