let ( let* ) = Result.bind

open Error
open Expr
open Lexer
open Stmt

type parse_result = (stmt_node list, parser_error list) result
type state = { tokens : token list ref; is_in_loop : bool; is_in_fun : bool }
type stmt_result = (stmt_node, parser_error) result
type expr_result = (expr_node, parser_error) result

let make_state tokens = { tokens; is_in_loop = false; is_in_fun = false }

let with_is_in_loop (f : state -> 'a) (state : state) : 'a =
  let new_state = { state with is_in_loop = true } in
  let result = f new_state in
  (* state.tokens <- new_state.tokens; *)
  result

let with_is_in_fun (f : state -> 'a) (state : state) : 'a =
  let new_state = { state with is_in_fun = true; is_in_loop = false } in
  let result = f new_state in
  (* state.tokens <- new_state.tokens; *)
  result

let is_at_end state =
  assert (not (List.is_empty !(state.tokens)));
  (List.hd !(state.tokens)).token_type = Eof

(* let advance state = state.tokens <- List.tl state.tokens *)
let advance state = state.tokens := List.tl !(state.tokens)
let peek state = List.hd !(state.tokens)
let peek_tt (state : state) : token_type = (peek state).token_type
let cur_pos state = (peek state).pos

let next state =
  let token = peek state in
  advance state;
  token

let advance_if state tt =
  if peek_tt state = tt then (
    advance state;
    true)
  else false

let consume state tt =
  if advance_if state tt then Ok ()
  else
    let pos = cur_pos state in
    let tt' = peek_tt state in
    let msg = Printf.sprintf "Expected %s, but got %s" (show_token_type tt) (show_token_type tt') in
    ParserError.make pos msg |> Result.error

let consume_identifier state =
  match peek_tt state with
  | Identifier name ->
      advance state;
      Ok name
  | tt ->
      let pos = cur_pos state in
      let msg = Printf.sprintf "Expected identifier, but got %s" (show_token_type tt) in
      ParserError.make pos msg |> Result.error

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

let collect_chain (tts : token_type array) (collector : state -> ('a, parser_error) result)
    (state : state) : (('a * token) list, parser_error) result =
  let rec collect_chain_rec (acc : ('a * token) list) =
    if (not (is_at_end state)) && matches state tts then
      let token = next state in
      let* expr = collector state in
      let acc = (expr, token) :: acc in
      collect_chain_rec acc
    else Ok acc
  in
  collect_chain_rec [] |> Result.map List.rev

let primary (state : state) : expr_result =
  let pos = cur_pos state in
  match peek_tt state with
  | Number x ->
      advance state;
      make_number pos x |> Result.ok
  | String s ->
      advance state;
      make_string pos s |> Result.ok
  | True ->
      advance state;
      make_bool pos true |> Result.ok
  | False ->
      advance state;
      make_bool pos false |> Result.ok
  | Nil ->
      advance state;
      make_nil pos |> Result.ok
  | Identifier name ->
      advance state;
      make_variable pos name |> Result.ok
  | tt ->
      advance state;
      let msg = Printf.sprintf "Expected valid expression, got %s instead" (show_token_type tt) in
      Error { msg; pos }

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

and call (state : state) : expr_result =
  let* expr = grouping state in
  let pos = cur_pos state in
  let rec parse_calls_rec expr =
    if advance_if state LeftParen then
      let* args =
        if peek_tt state = RightParen then Ok []
        else
          let* first_arg = expression state in
          let* exprs_tokens = collect_chain [| Comma |] expression state in
          let other_args = List.map fst exprs_tokens in
          let args = first_arg :: other_args in
          Ok args
      in
      if List.length args >= 255 then
        let msg = "Can't call with more than 255 arguments" in
        ParserError.make pos msg |> Result.error
      else
        let* () = consume state RightParen in
        let expr = make_call pos expr args in
        parse_calls_rec expr
    else Ok expr
  in
  parse_calls_rec expr

and neg_not (state : state) : expr_result =
  if matches state [| Bang; Minus |] then
    let token = next state in
    let pos = token.pos 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 pos op expr in
    Ok expr
  else call state

and mul_or_div (state : state) : expr_result =
  let* expr = neg_not state in
  let* exprs_tokens = collect_chain [| Star; Slash |] neg_not state in
  let f acc (expr, (token : token)) =
    let pos = token.pos in
    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 pos op acc expr
  in
  let expr = List.fold_left f expr exprs_tokens in
  Ok expr

and sum_or_diff (state : state) : expr_result =
  let* expr = mul_or_div state in
  let* exprs_tokens = collect_chain [| Plus; Minus |] mul_or_div state in
  let f acc (expr, (token : token)) =
    let pos = token.pos in
    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 pos op acc expr
  in
  let expr = List.fold_left f expr exprs_tokens in
  Ok expr

and inequality (state : state) : 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 [| Greater; GreaterEqual; Less; LessEqual |] sum_or_diff state
  in
  let f acc (expr, (token : token)) =
    let pos = token.pos in
    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 pos op acc expr
  in
  let expr = List.fold_left f expr exprs_tokens in
  Ok expr

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

and logical_and (state : state) : expr_result =
  let* expr = equality state in
  let* exprs_tokens = collect_chain [| And |] equality state in
  let f acc (expr, (token : token)) =
    let pos = token.pos in
    assert (token.token_type = And);
    make_logical pos And acc expr
  in
  let expr = List.fold_left f expr exprs_tokens in
  Ok expr

and logical_or (state : state) : expr_result =
  let* expr = logical_and state in
  let* exprs_tokens = collect_chain [| Or |] logical_and state in
  let f acc (expr, (token : token)) =
    let pos = token.pos in
    assert (token.token_type = Or);
    make_logical pos Or acc expr
  in
  let expr = List.fold_left f expr exprs_tokens in
  Ok expr

and assignment (state : state) : expr_result =
  let* expr = logical_or state in
  if Equal = peek_tt state then
    let pos = (next state).pos in
    let* rhs = assignment state in
    match expr.expr with
    | Variable name -> make_assignment pos name rhs |> Result.ok
    | _ ->
        let msg = "Invalid assignment target" in
        ParserError.make pos msg |> Result.error
  else Ok expr

and expression (state : state) : expr_result = assignment state

let rec block (state : state) : stmt_result =
  let pos = cur_pos state in
  let* () = consume state LeftBrace in
  let rec collect_stmts state =
    if is_at_end state then
      let msg = "Unterminated block" in
      ParserError.make pos msg |> Result.error
    else
      match peek_tt state with
      | RightBrace -> Ok []
      | _ ->
          let* stmt = declaration state in
          let* tail = collect_stmts state in
          Ok (stmt :: tail)
  in
  let* stmts = collect_stmts state in
  let* () = consume state RightBrace in
  make_block pos stmts |> Result.ok

and if_then_else (state : state) : stmt_result =
  let pos = cur_pos state in
  let* () = consume state If in
  let* () = consume state LeftParen in
  let* cond = expression state in
  let* () = consume state RightParen in
  let* then_ = statement state in
  let* (else_ : stmt_node option) =
    if advance_if state Else then statement state |> Result.map Option.some else Ok None
  in
  make_if pos cond then_ else_ |> Result.ok

and while_loop (state : state) : stmt_result =
  let pos = cur_pos state in
  let* () = consume state While in
  let* () = consume state LeftParen in
  let* cond = expression state in
  let* () = consume state RightParen in
  let* body = with_is_in_loop statement state in
  make_while pos cond body |> Result.ok

and for_loop (state : state) : stmt_result =
  let pos = cur_pos state in
  let* () = consume state For in
  let* () = consume state LeftParen in
  let* init =
    match peek_tt state with
    | Semicolon ->
        advance state;
        Ok None
    | Var -> var_declaration state |> Result.map Option.some
    | _ -> expr_stmt state |> Result.map Option.some
  in
  let* cond =
    match peek_tt state with
    | Semicolon -> Ok None
    | _ -> expression state |> Result.map Option.some
  in
  (* expression has no final semicolon, so we need to consume it *)
  let* () = consume state Semicolon in
  let* update =
    match peek_tt state with
    | RightParen -> Ok None
    | _ -> expression state |> Result.map Option.some
  in
  let* () = consume state RightParen in
  let* body = with_is_in_loop statement state in
  make_for pos init cond update body |> Result.ok

and expr_stmt (state : state) : stmt_result =
  let pos = cur_pos state in
  let* expr = expression state in
  let* () = consume state Semicolon in
  let stmt = make_expr_stmt pos expr in
  Ok stmt

and statement (state : state) : stmt_result =
  let pos = cur_pos state in
  match peek_tt state with
  | Break ->
      if not state.is_in_loop then (
        advance state;
        let msg = "Can use break only in loops" in
        ParserError.make pos msg |> Result.error)
      else (
        advance state;
        let* () = consume state Semicolon in
        make_break pos |> Result.ok)
  | Continue ->
      if not state.is_in_loop then (
        advance state;
        let msg = "Can use continue only in loops" in
        ParserError.make pos msg |> Result.error)
      else (
        advance state;
        let* () = consume state Semicolon in
        make_continue pos |> Result.ok)
  | Return ->
      if not state.is_in_fun then (
        advance state;
        let msg = "Can use return only in functions" in
        ParserError.make pos msg |> Result.error)
      else (
        advance state;
        let* expr =
          if peek_tt state = Semicolon then make_nil pos |> Result.ok else expression state
        in
        let* () = consume state Semicolon in
        make_return pos expr |> Result.ok)
  | Print ->
      advance state;
      let* expr = expression state in
      let* () = consume state Semicolon in
      let stmt = make_print pos expr in
      Ok stmt
  | LeftBrace -> block state
  | If -> if_then_else state
  | While -> while_loop state
  | For -> for_loop state
  | _ -> expr_stmt state

and var_declaration (state : state) : stmt_result =
  let pos = cur_pos state in
  (* consume var token *)
  let* () = consume state Var in
  let* name = consume_identifier state in
  let* init =
    if Equal = peek_tt state then
      (* found =, parsing initialiser *)
      let* () = consume state Equal in
      let* init = expression state in
      Ok (Some init)
    else (* no initialiser, default to nil *)
      Ok None
  in
  let* () = consume state Semicolon in
  make_var_decl pos name init |> Result.ok

and fun_declaration (state : state) : stmt_result =
  let pos = cur_pos state in
  let* () = consume state Fun in
  let* name = consume_identifier state in
  let* () = consume state LeftParen in
  let* arg_names =
    if peek_tt state = RightParen then Ok []
    else
      let* first_arg = consume_identifier state in
      let* exprs_tokens = collect_chain [| Comma |] consume_identifier state in
      let other_args = List.map fst exprs_tokens in
      let args = first_arg :: other_args in
      Ok args
  in
  if List.length arg_names >= 255 then
    let msg = Printf.sprintf "Function %s can't have more than 255 arguments" name in
    ParserError.make pos msg |> Result.error
  else
    let* () = consume state RightParen in
    let* body = with_is_in_fun block state in
    make_fun_decl pos name arg_names body |> Result.ok

and declaration (state : state) : stmt_result =
  match peek_tt state with
  | Var -> var_declaration state
  | Fun -> fun_declaration state
  | _ -> statement state

let rec synchronise (state : state) =
  match peek_tt state with
  | Semicolon -> advance state
  | Break | Class | Continue | Fun | Var | For | If | While | Print | Return | Eof -> ()
  | _ ->
      advance state;
      synchronise state

let rec parse_impl (state : state) : parse_result =
  if peek_tt state = Eof then Ok []
  else
    let result = declaration state in
    match result with
    | Ok stmt ->
        let* stmts = parse_impl state in
        Ok (stmt :: stmts)
    | Error e -> (
        synchronise state;
        if peek_tt state = Eof then Error [ e ]
        else
          let tail_result = parse_impl state in
          match tail_result with Ok _ -> Error [ e ] | Error es -> Error (e :: es))

let parse (tokens : token list) : parse_result =
  (* filter out all the comment tokens *)
  let tokens =
    List.filter (fun tok -> match tok.token_type with Comment _ -> false | _ -> true) tokens
  in
  let state = make_state (ref tokens) in
  parse_impl state