open Error
open Expr
open Lexer

type parse_result = (expr, parser_error list) result
type state = { tokens : token list; errors_rev : parser_error list }

module State = struct
  let is_at_end state = (List.hd state.tokens).token_type == Eof

  let append_error msg pos state =
    let e = { pos; msg } in
    { state with errors_rev = e :: state.errors_rev }

  let advance state = { state with tokens = List.tl state.tokens }
  let peek state = List.hd state.tokens

  let next state =
    assert (not ((List.hd state.tokens).token_type == Eof));
    (List.hd state.tokens, advance state)

  let advance_if state tt =
    if (peek state).token_type == tt then (true, advance state) else (false, state)

  let matches state tts =
    let f = ( == ) (peek state).token_type in
    List.fold_left (fun acc tt -> acc || f tt) false tts

  let collect_chain (state : state) (tts : token_type list) higher_prec :
      (expr * token) array * state =
    (* ([], state) *)
    let state_ref = ref state in
    let out_list_rev = ref [] in
    while (not (is_at_end !state_ref)) && matches !state_ref tts do
      let token, state = next !state_ref in
      let expr, state = higher_prec state in
      state_ref := state;
      out_list_rev := (expr, token) :: !out_list_rev
    done;
    (Array.of_list (List.rev !out_list_rev), !state_ref)

  let mul_or_div (state : state) : expr * state =
    let token, state = next state in
    (make_string @@ show_token_type token.token_type, state)

  let sum_or_diff (state : state) : expr * state =
    let expr, state = mul_or_div state in
    (* if (not (is_at_end state)) && matches state [ Plus; Minus ] then
         let token, state = next state in
         let expr2, state = sum_or_diff state in
         let (op : binary_op) =
           match token.token_type with
           | Plus -> Plus
           | Minus -> Minus
           | _ -> assert false (* should only be here if tt is + - *)
         in
         let expr = make_binary op expr expr2 in
         (expr, state)
       else (expr, state) *)
    (* turn expr and state to refs for the loop *)
    (* Printf.printf "expr: %s\n\n" (show_expr expr);
       let expr_ref, state_ref = (ref expr, ref state) in
       while (not (is_at_end !state_ref)) && matches !state_ref [ Plus; Minus ] do
         let token, state = next !state_ref in
         let expr2, state = mul_or_div state in
         let (op : binary_op) =
           match token.token_type with
           | Plus -> Plus
           | Minus -> Minus
           | _ -> assert false (* should only be here if tt is + - *)
         in
         let expr = make_binary op !expr_ref expr2 in
         (* (expr_ref, state_ref) := (expr, state) *)
         Printf.printf "expr: %s\n\n" (show_expr expr);
         expr_ref := expr;
         state_ref := state
       done;
       (!expr_ref, !state_ref) *)
    let exprs_tokens, state = collect_chain state [ Plus; Minus ] mul_or_div in
    let f acc (expr, token) =
      let op : binary_op =
        match token.token_type with Plus -> Plus | Minus -> Minus | _ -> assert false
      in
      make_binary op acc expr
    in
    let expr = Array.fold_left f expr exprs_tokens in
    (expr, state)

  let rec inequality (state : state) : expr * state =
    let expr, state = sum_or_diff state in
    if (not (is_at_end state)) && matches state [ Greater; GreaterEqual; Less; LessEqual ] then
      let token, state = next state in
      let expr2, state = inequality state in
      (* TODO: maybe rework to only have Less and Greater as ops; performance? *)
      let (op : binary_op) =
        match token.token_type with
        | Greater -> Greater
        | Less -> Less
        | GreaterEqual -> GreaterEqual
        | LessEqual -> LessEqual
        | _ -> assert false (* should only be here if tt is > < >= <= *)
      in
      let expr = make_binary op expr expr2 in
      (expr, state)
    else (expr, state)

  let rec equality (state : state) : expr * state =
    let expr, state = inequality state in
    if matches state [ EqualEqual; BangEqual ] then
      let token, state = next state in
      let expr2, state = equality state in
      let expr =
        match token.token_type with
        | EqualEqual -> make_binary Equal expr expr2
        | BangEqual ->
            let expr = make_binary Equal expr expr2 in
            make_unary Not expr
        | _ -> assert false (* should only be here if tt is == != *)
      in
      (expr, state)
    else (expr, state)

  let expression (state : state) : expr * state = equality state
end

let parse (tokens : token list) : parse_result =
  let state = { tokens; errors_rev = [] } in
  let expr, state = State.expression state in
  let state =
    if not (State.is_at_end state) then
      let tt = (State.peek state).token_type in
      let msg = Printf.sprintf "Unexpected %s at end" (show_token_type tt) in
      State.append_error msg (State.peek state).pos state
    else state
  in
  (* if List.length state.errors_rev != 0 then Ok expr else Error (List.rev state.errors_rev) *)
  match state.errors_rev with [] -> Ok expr | es -> Error (List.rev es)