open Lam
open Types

exception TacticFailed

type hlam = (* hollow lam *)
  HFun of (id * Types.ty) * hlam
  | HApp of hlam * hlam
  | HVar of id
  | HExf of hlam * Types.ty
  | Ref of hlam ref
  | Hole

type context = (id * hlam * Types.ty) list
type goal = hlam * Types.ty * context
type proof = goal option * goal list

type tactic =
  Exact of lam
  | Assumption
  | Intro
  | Cut of ty
  | Apply of id

let count = ref 0
let get_fresh_hyp () =
  let n = string_of_int !count in
  incr count;
  let hyp_id = "H" ^ n in
  let var_id = "x_" ^ n in
  (hyp_id, var_id)

let proof_is_over ((g, _) : proof) : bool =
  match g with
  None -> true
  | Some _ -> false

let fill (hole : hlam) (e : hlam) : unit =
  match hole with
  Ref h -> h := e
  | _ -> failwith "not fillable"

let rec hlam_of_lam : lam -> hlam = function
  Fun ((x, t), e) ->
    let e = hlam_of_lam e in
    HFun ((x, t), e)
  | App (e1, e2) ->
    let e1 = hlam_of_lam e1 in
    let e2 = hlam_of_lam e2 in
    HApp (e1, e2)
  | Var x -> HVar x
  | Exf (e, t) ->
    let e = hlam_of_lam e in
    HExf (e, t)

let rec lam_of_hlam : hlam -> lam = function
  HFun ((x, t), e) ->
    let e = lam_of_hlam e in
    Fun ((x, t), e)
  | HApp (e1, e2) ->
    let e1 = lam_of_hlam e1 in
    let e2 = lam_of_hlam e2 in
    App (e1, e2)
  | HVar x -> Var x
  | HExf (e, t) ->
    let e = lam_of_hlam e in
    Exf (e, t)
  | Ref e_ref -> lam_of_hlam !e_ref
  | Hole -> failwith "can not translate unclosed terms"

let rec get_term_by_type (ty : Types.ty) : context -> hlam option =
  function
  [] -> None
  | (_, h, ty') :: _ when ty' = ty -> Some h
  | _ :: cs -> get_term_by_type ty cs

let next_goal (gs : goal list) : (goal option * goal list) =
  match gs with
  [] -> None, []
  | g :: gs -> Some g, gs

let tact_exact_term ((g, gs) : proof) (e : lam) : proof =
  match g with
  None -> failwith "no current goal"
  | Some (h, _, _) ->
    fill h (hlam_of_lam e);
    next_goal gs

let tact_exact_proof ((g, gs) : proof) (h : id) : proof =

let tact_assumption ((g, gs) : proof) : proof =
  match g with
  None -> failwith "no current goal"
  | Some (h, goal_ty, cs) ->
    match get_term_by_type goal_ty cs with
    None -> (* failwith "assumption failed" *) (g, gs)
    | Some h' ->
      fill h h';
      next_goal gs

let tact_intro ((g, gs) : proof) : proof =
  match g with
  None -> failwith "no current goal"
  | Some (h, goal_ty, cs) ->
    match goal_ty with
    Arr (t1, t2) ->
      let (hyp_id, var_id) = get_fresh_hyp () in
      let cs = (hyp_id, HVar var_id, t1) :: cs in
      let new_h = Ref (ref Hole) in
      fill h (HFun ((var_id, t1), new_h));
      Some (new_h, t2, cs), gs
    | _ -> (* failwith "expected function" *) (* (g, gs) *)
            raise TacticFailed

let tact_cut ((g, gs) : proof) (new_t : Types.ty) : proof =
  match g with
  None -> failwith "no current goal"
  | Some (h, goal_ty, cs) ->
    (* subgoal 2 : new_t -> goal_ty *)
    let arrow_h = Ref (ref Hole) in
    let arrow_goal = (arrow_h, Arr (new_t, goal_ty), cs) in
    let gs = arrow_goal :: gs in
    (* subgoal 1 (main goal) : new_t *)
    let new_h = Ref (ref Hole) in
    fill h (HApp (arrow_h, new_h));
    Some (new_h, new_t, cs), gs

let tact_apply ((g, gs) : proof) (hyp_id : id) : proof =
  let rec get_hyp = function
    [] -> failwith "no such hypothesis in context"
    | (hyp_id', h', t') :: _ when hyp_id = hyp_id' -> (h', t')
    | _ :: cs -> get_hyp cs
  in
  match g with
  None -> failwith "no current goal"
  | Some (h, goal_ty, cs) ->
    let h', t' = get_hyp cs in
    match t' with
    Arr (t1, t2) when t2 = goal_ty ->
      let new_h = Ref (ref Hole) in
      fill h (HApp (h', new_h));
      Some (new_h, t1, cs), gs
    | Arr _ -> (* failwith "wrong types" *) (g, gs)
    | _ -> (* failwith "expected ->" *) (g, gs)

let tact_intros : proof -> proof =
  let rec push (p : proof) =
    try
      let p = tact_intro p in
      push p
    with TacticFailed -> p
  in push

let tact_try (p : proof) (t : tactic) : proof =
  try
  match t with
  Exact e -> tact_exact p e
  | Intro -> tact_intro p
  | Assumption -> tact_assumption p
  | Cut t -> tact_cut p t
  | Apply h -> tact_apply p e
  with TacticFailed -> p