open Ast
open Ast.Operator

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

type associativity =
  | Left_to_right
  | Right_to_left

let assoc_of_string = function
  | "left" -> Left_to_right
  | "right" -> Right_to_left
  | _ -> invalid_arg "assoc_of_string"

let assoc_to_string = function
  | Left_to_right -> "left"
  | Right_to_left -> "right"

let oper_assoc = [
  Exp, Right_to_left;
] |> List.to_seq |> Hashtbl.of_seq

let op_is_right_to_left op =
  let a =
    Hashtbl.find_opt oper_assoc op
    |> Option.value ~default: Left_to_right
  in
  a = Right_to_left

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 token tok seq =
  match seq () with
  | Seq.Nil -> expected @@ Token.to_string tok
  | Seq.Cons (x, seq) ->
    if x = tok then x, seq
    else expected @@ Token.to_string tok

let any_ident seq =
  match seq () with
  | Seq.Nil -> expected "ident"
  | Seq.Cons (x, seq) -> begin match x with
      | Token.Ident id -> id, seq
      | _ -> unexpected_token x
    end

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 oneof fs seq =
  let rec aux = function
    | [] -> assert false
    | [f] -> f seq
    | f::fs -> (try f seq with _ -> aux fs)
  in
  aux fs

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

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

let more f seq =
  let rec aux xs seq =
    try
      let x, seq = f seq in
      aux (x::xs) seq
    with
    | _ -> xs, seq
  in
  let xs, seq = aux [] seq in
  List.rev xs, seq

(* expr := level
 *       | assoc
 *       | let
 *       | apply
 *       | value binop_right
*)
let rec expr pre seq =
  seq |> oneof [
    level;
    assoc;
    let_value;
    apply;
    (either unary value) @> binop pre;
  ]

(* level := "level" {"get" | "set"} [op] *)
and level seq =
  let _, seq = ident "level" seq in
  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"

(* assoc := "assoc" {"get" | "set"} [op] *)
and assoc seq =
  let _, seq = ident "assoc" seq in
  let id, seq = idents (S.of_list ["get"; "set"]) seq in
  let op, seq = operator seq in
  if id = "get" then
    Get_binop_aso op, seq
  else if id = "set" then
    let a, seq = idents (S.of_list ["left"; "right"]) seq in
    Set_binop_aso (op, a), seq
  else
    failwith "Parser.assoc"

(* let_value := "let" ident "=" expr *)
and let_value seq =
  let _, seq = ident "let" seq in
  let id, seq = any_ident seq in
  let _, seq = token Token.Equal seq in
  let e, seq = expr min_int seq in
  Let (id, e), seq

(* apply := ident [value]* *)
and apply seq =
  let id, seq = any_ident seq in
  let args, seq = more value seq in
  Apply (Var id, args), seq

(* unary := - value *)
and unary seq =
  let op, seq =
    match seq () with
    | Seq.Nil -> raise End_of_tokens
    | Seq.Cons (x, seq) ->
      if x = Minus
      then Negate, seq
      else expected "minus"
  in
  let v, seq = value seq in
  Ast.unary op v, seq

(* value := int | float | ( expr ) *)
and value seq =
  match seq () with
  | Seq.Nil -> raise End_of_tokens
  | Seq.Cons (x, seq) -> begin match x with
      | Token.Ident id -> Var id, seq
      | Int x -> Value (Int x), seq
      | Float x -> Value (Float x), seq
      | String x -> Value (String x), seq
      | LParen ->
        let e, seq = expr min_int seq in
        let _, seq = token RParen seq in
        e, 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 op_pre = precedence_of op in
        (* op has to be calculated first *)
        if op_pre > pre
        || (op_is_right_to_left op && op_pre = pre)
        then
          let right, seq = expr op_pre seq in
          binop pre (Ast.binop left op right) seq
        else
          left, Seq.cons x seq

      | Token.RParen -> left, Seq.cons x seq
      | _ -> unexpected_token x
    end

(* parse tokens *)
let parse ts =
  let ast, rest = expr min_int ts in
  if rest () <> Seq.Nil then failwith "Parser.parse";
  ast