open Structs;;
open Pieuvre;;
open Tactic;;

(* Lorsqu'une tactique ne peut se lancer. Contient la tactique ainsi qu'un message contenant l'erreur.*)
exception TacticException of tactic*string;;
(* Lorsqu'une tactique est mal écrite*)
exception TacticParseException;;
(* Exception renvoyée lorsque la fonction de trou est appelée sur une liste trop grande. *)
(* Ne devrait jamais être lancée *)
exception TrouException;;

(* Parsage des arguments*)
let filename = ref "" in

let reduce_option = ref false in
let alpha_option = ref false in
let typecheck_option = ref false in
let computetype_option = ref false in

Arg.parse
    [("-reduce", Arg.Set reduce_option, "Show the step-by-step reduction of a lambda-term");
     ("-alpha", Arg.Set alpha_option, "Check is two lambda-terms separated by '&' are alpha-convertible");
     ("-typecheck", Arg.Set typecheck_option, "Check if a lambda term has a given type");
     ("-computetype", Arg.Set computetype_option, "Computes the type of the input λ-term")]
    (fun s -> filename := s)
    "The available options are:";

(* Ouverture éventuelle du fichier *)
let file = match !filename with
  | "" -> None
  | fn -> Some (open_in fn)
in
let is_interactive = match file with
  | None -> true
  | _ -> false
in

if !reduce_option then (
    let lexbuf = Lexing.from_channel (match file with
      | None -> stdin
      | Some file -> file
    )
    in
    let lambda_term =
        try
            Parser.main_lambda Lexer.token lexbuf
        with e -> (
            Printf.printf "Can't read lambda term\n";
            raise e
       )
    in
    reduce lambda_term;
    exit 0
);

if !alpha_option then (
    let lexbuf = Lexing.from_channel (match file with
      | None -> stdin
      | Some file -> file
    )
    in
    let lam1, lam2 =
        try
            Parser.main_two_lambda Lexer.token lexbuf
        with e -> (
            Printf.printf "Can't read lambda terms\n";
            raise e
       )
    in
    if alpha lam1 lam2 then (
        Printf.printf "%s and %s are α-equivalent\n" (string_of_lam lam1) (string_of_lam lam2)
    ) else (
        Printf.printf "%s and %s are not α-equivalent\n" (string_of_lam lam1) (string_of_lam lam2)
    );
    exit 0
);

if !typecheck_option then (
    let lexbuf = Lexing.from_channel (match file with
      | None -> stdin
      | Some file -> file
    )
    in
    let lambda_term, ty =
        try
            Parser.main_typed_lambda Lexer.token lexbuf
        with e -> (
            Printf.printf "Can't read typed lambda term\n";
            raise e
       )
    in
    Printf.printf "The type of %s is " (string_of_lam lambda_term);
    if not (typecheck [] lambda_term ty) then (
        Printf.printf "not "
    );
    Printf.printf "%s.\n" (string_of_ty ty);
    exit 0
);

if !computetype_option then (
    let lexbuf = Lexing.from_channel (match file with
      | None -> stdin
      | Some file -> file
    )
    in
    let lambda_term =
        try
            Parser.main_lambda Lexer.token lexbuf
        with e -> (
            Printf.printf "Can't read lambda term\n";
            raise e
        )
    in
    begin match (computeType [] lambda_term) with
      | Some t ->
            Printf.printf "The type of %s is %s\n" (string_of_lam lambda_term) (string_of_ty t)
      | None -> (
            Printf.printf
                "λ-term %s can't be typed! (it is not closed or not typable)\n"
                (string_of_lam lambda_term)
        )
    end;
    exit 0
);

(* Affiche un message si l'entrée est lue sur stdin *)
let show s = match file with
  | None -> Printf.printf "%s" s
  | _ -> ()
in

(* Ouvre un fichier contenant la preuve écrite *)
let log_file = open_out "log.8pus" in

let readline () =
    let line = match file with
      | None -> (
            Printf.printf ">>> ";
            flush stdout;
            read_line ()
        )
      | Some f -> input_line f
    in
    Printf.fprintf log_file "%s\n" line;
    flush log_file;
    line
in

show "Please type the formula to prove\n";

let ty =
    try
        let lexbuf = Lexing.from_string (readline ()) in
        Parser.main_type Lexer.token lexbuf
    with e -> (
        Printf.printf "Can't parse type\n";
        raise e
    )
in

let subgoals: (ty * (var_lambda * ty) list) list ref = ref [(ty, [])] in

(* On donne en paramètre une liste contenant un lambda-terme par trou (ie subgoal),
 * et elle renvoie le lambda-terme complet (en remplaçant les trous) *)
let fill_holes = ref (fun holes -> List.hd holes) in

while !subgoals <> [] do
    if is_interactive then (
        (* Nettoyage du terminal *)
        let _ = Sys.command("clear -x") in

        (* Affichage des hypothèses *)
        List.iter
            (fun (var, h) -> Printf.printf "%s: %s\n" var (string_of_ty h))
            (snd (List.hd !subgoals));

        (* Affichage des sous-buts *)
        Printf.printf "================\n";
        List.iter (fun (ty, _) ->
            Printf.printf "%s\n" (string_of_ty ty)
        ) !subgoals;

        (* Lecture d'une tactique *)
        Printf.printf "What do you want to do?\n"
    );

    let rec read_tactic () =
        try
            let lexbuf = Lexing.from_string (readline ()) in
            Parser.main_tactic Lexer.token lexbuf
        with e -> raise TacticParseException
    in

    let rec applyTactic tactic = (
        let (ty, hyps) = List.hd !subgoals in

        let find_hyp (var: var_lambda) : ty option =
            let rec explore = function
              | [] -> None
              | (var_hyp, hyp) :: hyps when var_hyp = var -> Some hyp
              | _ :: hyps -> explore hyps
            in
            explore hyps
        in

        let f = !fill_holes in
        match tactic with
          | Intro var -> (
                match ty with
                  | TImpl (ty1, ty2) -> (
                        subgoals := List.tl !subgoals;
                        subgoals := (ty2, (var, ty1) :: hyps) :: !subgoals;
                        fill_holes := fun holes -> match holes with
                          | h :: hs -> f (LFun (var, ty1, h) :: hs)
                          | _ -> raise TrouException
                    )
                  | _ -> raise (TacticException(tactic, "Cannot intro when the goal is not an implication."))
            )
          | Intros [] -> ()
          | Intros (t :: ts) -> (
                applyTactic (Intro t);
                applyTactic (Intros ts)
            )
          | Assumption -> (
                let rec explore = function
                  | (var, hyp) :: _ when hyp = ty -> (
                        subgoals := List.tl !subgoals;
                        fill_holes := fun holes -> f ((LVar var) :: holes)
                    )
                  | [] -> raise (TacticException (tactic, "Cannot find such an assumption."))
                  | _ :: hyps -> explore hyps
                in
                explore hyps
            )
          | Apply var -> (
                match find_hyp var with
                  | Some (TImpl (t1, t2)) when t2 = ty -> (
                        subgoals := List.tl !subgoals;
                        subgoals := (t1, hyps) :: !subgoals;
                        fill_holes := function
                          | hole :: holes -> f ((LApp (LVar var, hole)) :: holes)
                          | [] -> raise TrouException
                    )
                  | None -> raise (TacticException (
                        tactic,
                        "Cannot apply hypotesis " ^ var ^ " as it does not exist."
                    ))
                  | _ -> raise (TacticException (
                        tactic,
                        "Cannot apply hypotesis " ^ var ^ " as it is not an implication."
                    ))
            )
          | Elim var -> (
                match find_hyp var with
                  | Some TFalse -> (
                        subgoals := List.tl !subgoals;
                        fill_holes := fun holes -> f ((LExf (LVar var, ty)) :: holes)
                    )
                  | Some TAnd(tl,tr) -> (
                        subgoals := List.tl !subgoals;
                        subgoals := (TImpl (tl, TImpl (tr, ty)), hyps) :: !subgoals;
                        fill_holes := fun holes -> match holes with
                          | h::r -> f ((LApp (LApp (h, LFst (LVar var)), LSnd (LVar var))) :: r)
                          | _ -> raise TrouException
                    )
                  | Some TOr(tl,tr) -> (
                        subgoals := List.tl !subgoals;
                        subgoals := (TImpl (tl, ty), hyps) :: (TImpl (tr, ty), hyps) :: !subgoals;
                        fill_holes := fun holes -> match holes with
                          | hl::hr::r -> f (LCase (LVar var, hl, hr) :: r)
                          | _ -> raise TrouException
                      )
                  | None -> raise (TacticException (
                        tactic,
                        "Cannot apply hypotesis " ^ var ^ " as it does not exist."
                    ))
                  | _ -> raise (TacticException (
                        tactic,
                        "Cannot apply hypotesis " ^ var ^ " as it is neither an implication, a /\\ nor a \\/"
                    ))
            )
          (* Pour montrer A, on montre B -> A et B *)
          | Cut tint -> (
                subgoals := List.tl !subgoals;
                subgoals := (TImpl (tint, ty), hyps) :: (tint, hyps) :: !subgoals;
                fill_holes := function
                  | pf :: px :: s -> f ((LApp (pf, px)) :: s)
                  | _ -> raise TrouException
            )

          | Split -> (
                match ty with
                  | TAnd(t1,t2) -> (
                        subgoals := List.tl !subgoals;
                        subgoals := (t1, hyps) :: (t2, hyps) :: !subgoals;
                        fill_holes := fun holes -> match holes with
                          | h1 :: h2 :: hs -> f (LCouple(h1,h2) :: hs)
                          | _ -> raise TrouException
                    )
                  | _ -> raise (TacticException (tactic, "Cannot split as the goal is not a /\\"))
            )
          | Left -> (
                match ty with
                  | TOr(tl,tr) -> (
                        subgoals := List.tl !subgoals;
                        subgoals := (tl, hyps) :: !subgoals;
                        fill_holes := fun holes -> match holes with
                          | hl :: hs -> f (LIg(hl,tr) :: hs)
                          | _ -> raise TrouException
                    )
                  | _ -> raise (TacticException (tactic, "Cannot apply left as the goal is not a \\/"))
            )
          | Right -> (
                match ty with
                  | TOr(tl,tr) -> (
                        subgoals := List.tl !subgoals;
                        subgoals := (tr, hyps) :: !subgoals;
                        fill_holes := fun holes -> match holes with
                          | hr :: hs -> f (LId(hr,tl) :: hs)
                          | _ -> raise TrouException
                    )
                  | _ -> raise (TacticException (tactic, "Cannot apply right as the goal is not a \\/"))
            )
          | Exact l -> (
                if not (typecheck hyps l ty) then
                    let ty_str = match computeType [] l with
                      | None -> "None"
                      | Some t -> string_of_ty t
                    in
                    raise (TacticException (
                        tactic,
                        "λ-term " ^ (string_of_lam l) ^
                        " can't be typed with " ^ (string_of_ty ty) ^
                        " as its type is " ^ ty_str
                    ))
                else (
                    subgoals := List.tl !subgoals;
                    fill_holes := fun holes -> f (l::holes)
                )
            )
    )
    in

    let rec applyUntilWorking () =
        try (
            let readTactic = read_tactic () in
            applyTactic readTactic
        )
        with
          | TacticException(t,s) ->
                Printf.printf "\027[31mCannot apply the tactic: %s\027[0m\n" s;
                if is_interactive then
                    applyUntilWorking ()
                else
                    raise (TacticException (t,s))
          | TacticParseException ->
                Printf.printf "\027[31mCannot parse the tactic, please refer to pieuvre documentation.\027[0m\n";
                if is_interactive then
                    applyUntilWorking ()
                else
                    raise TacticParseException
          | e ->
                Printf.printf "\027[31mPieuvre Failed Unexpectedly !\027[0m\n";
                raise e
    in
    applyUntilWorking ()
done;

let finalLam = !fill_holes [] in
if (typecheck [] finalLam ty) then (
    Printf.printf "Final proof :\n";
    reduce finalLam;
) else (
    Printf.printf "Invalid proof constructed!\n";
    Printf.printf "%s can't be typed with %s.\n" (string_of_lam finalLam) (string_of_ty ty);
    Printf.printf "The whole development team of pieuvre is sorry for the damage eventually done by this error.\n"
);

close_out log_file;;