let ( let* ) = Result.bind

open Error
open Expr
open Lexer

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

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 := { !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));
  let token = List.hd !state.tokens in
  advance state;
  token

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

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

let cur_pos state = (peek state).pos

let collect_chain (state : state ref) (tts : token_type array)
    (higher_prec : state ref -> expr_result) : ((expr * token) array, parser_error) result =
  (* ([], state) *)
  (* let out_list_rev = ref [] in
     while (not (is_at_end state)) && matches state tts do
       let token = next state in
       let expr = higher_prec state in
       out_list_rev := (expr, token) :: !out_list_rev
     done;
     Ok (Array.of_list (List.rev !out_list_rev)) *)
  let rec collect_chain_rec (acc : (expr * token) list) =
    if (not (is_at_end state)) && matches state tts then
      let token = next state in
      let* expr = higher_prec state in
      let acc = (expr, token) :: acc in
      collect_chain_rec acc
    else Ok acc
  in
  (* match collect_chain_rec [] with Ok l -> Ok (Array.of_list (List.rev l)) | Error e -> Error e *)
  collect_chain_rec [] |> Result.map (fun l -> Array.of_list (List.rev l))

let primary (state : state ref) : expr_result =
  match (peek state).token_type with
  | Number x -> Ok (make_number x)
  | String s -> Ok (make_string s)
  | True -> Ok (make_bool true)
  | False -> Ok (make_bool false)
  | Nil -> Ok (make_nil ())
  | tt ->
      let msg = Printf.sprintf "Unexpected %s, expected valid expression" (show_token_type tt) in
      let pos = (peek state).pos in
      Error { msg; pos }

let rec grouping (state : state ref) : expr_result =
  if matches state [| LeftParen |] then (
    advance state;
    let* expr = expression state in
    if advance_if state RightParen then Ok expr
    else
      let pos = cur_pos state in
      let tt = (peek state).token_type in
      let msg = Printf.sprintf "Expected RightParen, got %s instead" (show_token_type tt) in
      Error { pos; msg })
  else primary state

and neg_not (state : state ref) : expr_result =
  if matches state [| Bang; Minus |] then
    let token = next state in
    let* expr = neg_not state in
    let op =
      match token.token_type with
      | Bang -> Not
      | Minus -> Neg
      | _ -> assert false (* should only be here if tt is - ! *)
    in
    let expr = make_unary op expr in
    Ok expr
  else grouping state

and mul_or_div (state : state ref) : expr_result =
  let* expr = neg_not state in
  let* exprs_tokens = collect_chain state [| Star; Slash |] neg_not in
  let f acc (expr, token) =
    let op : binary_op =
      match token.token_type with
      | Star -> Mul
      | Slash -> Div
      | _ -> assert false (* should only be here if tt is * / *)
    in
    make_binary op acc expr
  in
  let expr = Array.fold_left f expr exprs_tokens in
  Ok expr

and sum_or_diff (state : state ref) : expr_result =
  let* expr = mul_or_div state in
  let* exprs_tokens = 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 (* should only be here if tt is + - *)
    in
    make_binary op acc expr
  in
  let expr = Array.fold_left f expr exprs_tokens in
  Ok expr

and inequality (state : state ref) : expr_result =
  (* TODO: maybe rework to only have Less and Greater as ops; performance? *)
  let* expr = sum_or_diff state in
  let* exprs_tokens =
    collect_chain state [| Greater; GreaterEqual; Less; LessEqual |] sum_or_diff
  in
  let f acc (expr, token) =
    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
    make_binary op acc expr
  in
  let expr = Array.fold_left f expr exprs_tokens in
  Ok expr

and equality (state : state ref) : expr_result =
  let* expr = inequality state in
  let* exprs_tokens = collect_chain state [| EqualEqual; BangEqual |] inequality in
  let f acc (expr, token) =
    match token.token_type with
    | EqualEqual -> make_binary Equal acc expr
    | BangEqual ->
        let expr = make_binary Equal acc expr in
        make_unary Not expr
    | _ -> assert false (* should only be here if tt is == != *)
  in
  let expr = Array.fold_left f expr exprs_tokens in
  Ok expr

and expression (state : state ref) : expr_result = equality state

let rec synchronise (state : state ref) =
  match (peek state).token_type with
  | Semicolon -> advance state
  | Class | Fun | Var | For | If | While | Print | Return | Eof -> ()
  | _ ->
      advance state;
      synchronise state

let parse (tokens : token list) : parse_result =
  let state = ref { tokens; errors_rev = [] } in
  expression state |> Result.map_error (fun e -> [ e ])
(* let expr = 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) *)