open Ast
open Ast.Binop

module S = Set.Make(String)

exception Expected of string
exception Unexpected_token of string
exception End_of_tokens

let expected t =
  raise @@ Expected t

let unexpected_token t =
  raise @@ Unexpected_token (Token.to_string t)

(* precedence table.
 * my first thought was using some sort of partially-ordered graph for
 * precedency, but infering precedence relation from the graph is hard
 * and the graph can be made to have loops, I just used plain table. *)
let precedence = [
  Add, 10;
  Sub, 10;
  Mul, 20;
  Div, 20;
  Mod, 30;
  Exp, 30;
] |> List.to_seq |> Hashtbl.of_seq

let precedence_of op =
  Hashtbl.find precedence op

let op_is_right_to_left = function
  | Exp -> true
  | _ -> false

let operators = [
  Token.Plus, Add;
  Minus, Sub;
  Asterisk, Mul;
  Slash, Div;
  Carret, Exp;
  Percent, Mod;
] |> List.to_seq |> Hashtbl.of_seq

let token_to_op tok =
  try Hashtbl.find operators tok
  with _ -> failwith "Parser.token_to_op"

let token_is_operator tok =
  Hashtbl.mem operators tok

(* common parsers *)

let idents set seq =
  match seq () with
  | Seq.Nil ->
    let msg = "ident " ^ (S.elements set |> String.concat " or ") in
    expected msg
  | Seq.Cons (x, seq) -> begin
      match x with
      | Token.Ident id when S.mem id set -> id, seq
      | _ -> unexpected_token x
    end

let ident str seq =
  idents (S.singleton str) seq

let operator seq =
  match seq () with
  | Seq.Nil -> expected "operator"
  | Seq.Cons (x, seq) ->
    try token_to_op x, seq with
    | _ -> expected "operator"

(* parser combinators *)

let either f g seq =
  try f seq with _ -> g seq

let (@>) f g seq =
  let a, seq = f seq in
  g a seq

(* parse tokens *)
let parse ts =
  (* value := int | ( expr ) *)
  let rec value seq =
    match seq () with
    | Seq.Nil -> raise End_of_tokens
    | Seq.Cons (x, seq) -> begin match x with
        | Token.Int n -> Value (Int n), seq
        | Float n -> Value (Float n), seq
        | LParen -> expr seq
        | _ -> unexpected_token x
      end

  (* binop := binop op binop *)
  and binop pre left seq =
    match seq () with
    | Seq.Nil -> left, Seq.empty
    | Seq.Cons (x, seq) -> begin match x with
        | op when token_is_operator op ->
          let op = token_to_op op in
          let o = precedence_of op in 
          (* op has to be calculated first *)
          if o > pre || op_is_right_to_left op && o = pre then
            let v, seq = value seq in
            let right, seq = binop o v seq in
            binop pre (Ast.binop left op right) seq
          else
            left, Seq.cons x seq
        | Token.RParen -> left, seq
        | _ -> unexpected_token x
      end

  (* level_inner := "get" | "set" [op] *)
  and level_inner _ seq =
    let id, seq = idents (S.of_list ["get"; "set"]) seq in
    let op, seq = operator seq in
    if id = "get" then
      Get_binop_pre op, seq
    else if id = "set" then
      let v, seq = value seq in
      Set_binop_pre (op, v), seq
    else
      failwith "Parser.level"

  (* expr := "level" level_inner
   *       | value binop_right
   *)
  and expr seq =
    seq |> either
      (ident "level" @> level_inner)
      (value @> binop ~-1)
  in
  let ast, rest = expr ts in
  if rest () <> Seq.Nil then failwith "Parser.parse";
  ast