module StringSet = Set.Make (String)

type name = string

type type' =
  (* 1 *)
  | Unit
  (* x *)
  | Variable of name
  (* ^x *)
  | Existential of name
  (* forall a. A *)
  | ForAll of name * type'
  (* A -> B *)
  | Function of type' * type'

let rec free_variables = function
  | Unit -> StringSet.empty
  | Variable x | Existential x -> StringSet.singleton x
  | ForAll (alpha, a) -> StringSet.remove alpha (free_variables a)
  | Function (a, b) -> StringSet.inter (free_variables a) (free_variables b)

let occurs_check name type' = free_variables type' |> StringSet.mem name |> not

let rec is_monotype = function
  | Unit | Variable _ | Existential _ -> true
  | ForAll _ -> false
  | Function (a, b) -> is_monotype a && is_monotype b

let rec replace name with_type in_type =
  match in_type with
  | (Variable n | Existential n) when n = name -> with_type
  | Variable _ | Existential _ | Unit -> in_type
  | ForAll (alpha, _) when name = alpha -> in_type
  | ForAll (alpha, a) -> ForAll (alpha, replace name with_type a)
  | Function (a, b) ->
      Function (replace name with_type a, replace name with_type b)

type term =
  (* () *)
  | Unit
  (* x *)
  | Variable of name
  (* \x. e *)
  | Lambda of name * term
  (* f x *)
  | Application of term * term
  (* x : T *)
  | Annotation of term * type'

type judgment =
  (* universally quantified *)
  | ForAll of name
  | Assume of name * type'
  (* existentially quantified *)
  | Exists of name
  | Solved of name * type'
  (* marker *)
  | Marker of name

type context = judgment list

let existentials =
  List.fold_left
    (fun vars j ->
      match j with
      | Exists n | Solved (n, _) ->
          StringSet.singleton n |> StringSet.union vars
      | _ -> vars)
    StringSet.empty

let in_existentials name gamma = existentials gamma |> StringSet.mem name

exception TypeError of string

let lookup_name context name =
  let find_name = function
    | (ForAll n | Assume (n, _) | Exists n | Solved (n, _)) when n = name ->
        true
    | _ -> false
  in
  List.find find_name context

let rec drop_marker marker context =
  match context with
  | [] -> raise (TypeError "tried to drop missing marker")
  | top :: rest when top = marker -> rest
  | _ :: rest -> drop_marker marker rest

let solve name solution context =
  let solve' = function
    | Exists n when n = name -> Solved (n, solution)
    | j -> j
  in
  List.map solve' context

let rec substitute gamma (a : type') =
  match a with
  | Unit | Variable _ -> a
  | Existential alpha -> (
      match lookup_name gamma alpha with
      | Exists _ -> a
      | Solved (_, tau) -> substitute gamma tau
      | _ -> raise (TypeError "unbound existential"))
  | Function (a, b) -> Function (substitute gamma a, substitute gamma b)
  | ForAll (alpha, a) -> ForAll (alpha, substitute gamma a)

let mk_fresh_generator prefix =
  let counter = ref 0 in
  fun base ->
    let output = prefix ^ base ^ string_of_int !counter in
    counter := !counter + 1;
    output

let fresh_variable = mk_fresh_generator "$"
let fresh_type_variable = mk_fresh_generator "'"

let rec instantiate_left alpha a gamma =
  match a with
  | _ when is_monotype a -> solve alpha a gamma
  | Existential beta -> solve beta (Existential alpha) gamma
  | Function (a1, a2) ->
      let alpha_1 = fresh_type_variable alpha in
      let alpha_2 = fresh_type_variable alpha in
      let theta =
        instantiate_right a1 alpha_1
          (Solved (alpha, Function (Variable alpha_1, Variable alpha_2))
          :: Exists alpha_1 :: Exists alpha_2 :: gamma)
      in
      instantiate_left alpha_2 (substitute theta a2) theta
  | ForAll (beta, b) ->
      let beta' = fresh_type_variable beta in
      instantiate_left alpha b (ForAll beta' :: gamma)
      |> drop_marker (ForAll beta')
  | _ -> raise (TypeError "bad left instantiation")

and instantiate_right a alpha gamma =
  match a with
  | _ when is_monotype a -> solve alpha a gamma
  | Existential beta -> solve beta (Existential alpha) gamma
  | Function (a1, a2) ->
      let alpha_1 = fresh_type_variable alpha in
      let alpha_2 = fresh_type_variable alpha in
      let theta =
        instantiate_left alpha_1 a1
          (Solved (alpha, Function (Variable alpha_1, Variable alpha_2))
          :: Exists alpha_1 :: Exists alpha_2 :: gamma)
      in
      instantiate_right (substitute theta a2) alpha_2 theta
  | ForAll (beta, b) ->
      let beta' = fresh_type_variable beta in
      instantiate_right b alpha (ForAll beta' :: gamma)
      |> drop_marker (ForAll beta')
  | _ -> raise (TypeError "bad right instantiation")

let rec subtype (a : type') (b : type') gamma =
  match (a, b) with
  | Variable a, Variable b when a = b -> gamma
  | Unit, Unit -> gamma
  | Existential a, Existential b when a = b && in_existentials a gamma -> gamma
  | Function (a1, a2), Function (b1, b2) ->
      let theta = subtype b1 a1 gamma in
      subtype (substitute theta a2) (substitute theta b2) theta
  | ForAll (alpha, a), _ ->
      let alpha' = fresh_type_variable alpha in
      let gamma = Exists alpha' :: Marker alpha' :: gamma in
      let a' = replace alpha (Existential alpha') a in
      subtype a' b gamma |> drop_marker (Marker alpha')
  | _, ForAll (beta, b) ->
      let beta' = fresh_type_variable beta in
      let b' = replace beta (Variable beta') b in
      subtype a b' (ForAll beta' :: gamma) |> drop_marker (ForAll beta')
  | Existential alpha, _ when occurs_check alpha b ->
      instantiate_left alpha b gamma
  | _, Existential beta when occurs_check beta a ->
      instantiate_right a beta gamma
  | _ -> raise (TypeError "invalid subtype relation")

let rec check (e : term) (a : type') gamma =
  match (e, a) with _ -> raise (TypeError "bad check")

and synth (e : term) gamma = match e with _ -> raise (TypeError "bad synth")

and apply (a : type') (e : term) gamma =
  match (a, e) with _ -> raise (TypeError "bad apply")