open Ast

module S = Set.Make(String)

(* fatal exception that parsing need to be stopped *)
exception Fatal of exn
exception Expected of int * string
exception Unexpected_token of int * string
exception End_of_tokens

let expected col t =
  raise @@ Expected (col, t)

let unexpected_token col t =
  raise @@ Unexpected_token (col, 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 = [
  "=", 1;
  "<>", 1;
  ">=", 1;
  "<=", 1;
  ">", 1;
  "<", 1;
  "+", 10;
  "-", 10;
  "*", 20;
  "/", 20;
  "%", 30;
  "^", 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 = [
  "^", 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;
  Minus;
  Asterisk;
  Slash;
  Carret;
  Percent;
  Equal;
  Not_equal;
  Greater_equal;
  Less_equal;
  Greater;
  Less;
]

let token_to_op tok =
  Token.to_string tok

let token_is_operator tok =
  List.mem tok operators

let is_keyword = function
  | "if" | "then" | "else" | "let" | "in" -> true
  | _ -> false

(* parser primitives *)

let any seq =
  match seq () with
  | Seq.Nil -> raise End_of_tokens
  | Seq.Cons ((col, x), seq) -> col, x, seq

let token tok seq =
  let col, x, seq = any seq in
  if x = tok then x, seq
  else expected col @@ Token.to_string tok

let any_ident seq =
  let col, x, seq = any seq in
  match x with
  | Token.Ident id -> id, seq
  | _ -> expected col "ident"

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

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

let operator seq =
  let col, x, seq = any seq in
  try token_to_op x, seq with
  | _ -> expected col "operator"

(* parser combinators *)

let mustbe f seq =
  try f seq with
  | e -> raise @@ Fatal e

let oneof fs seq =
  let rec aux = function
    | [] -> assert false
    | [f] -> f seq
    | f::fs ->
      (try f seq with
       | Fatal _ as e -> raise e
       | _ -> aux fs)
  in
  aux fs

let either f g = fun seq ->
  try f seq with
  | Fatal _ as e -> raise e
  | _ -> 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
    | Fatal _ as e -> raise e
    | _ -> xs, seq
  in
  let xs, seq = aux [] seq in
  List.rev xs, seq

(* decl := let_global
 *       | expr
 *)
let rec decl seq =
  seq |> oneof [
    expr min_int;
    let_global;
    nothing;
  ]

and nothing seq =
  match seq () with
  | Seq.Nil -> Nothing, seq
  | Seq.Cons ((col, x), _) -> unexpected_token col x

(* let_global := "let" ident "=" expr *)
and let_global seq =
  let _, seq = ident "let" seq in
  let id, seq = any_ident seq in
  let args, seq = more any_ident seq in
  let _, seq = token Token.Equal seq in
  let e, seq = expr min_int seq in
  (if args = []
   then Let (id, e)
   else Let (id, Nfunction (args, e))),
  seq

(* expr := level
 *       | let_value
 *       | assoc
 *       | apply
 *       | value binop_right
*)
and expr pre seq =
  seq |> oneof [
    ifexpr;
    let_value;
    lambda;
    extern_value;
    oneof [apply; unary; value] @> binop pre;
    (* TODO: place error routine here *)
  ]

(* let_value := "let" id "=" expr "in" expr *)
and let_value seq =
  let _, seq = ident "let" seq in
  let id, seq = mustbe any_ident seq in
  let _, seq = token Equal seq in
  let e, seq = mustbe (expr min_int) seq in
  let _, seq = ident "in" seq in
  let f, seq = mustbe (expr min_int) seq in
  Letin (id, e, f), seq

(* lambda := "fun" [ident]+ "->" expr *)
and lambda seq =
  let _, seq = ident "fun" seq in
  seq |> mustbe (fun seq ->
      let v0, seq = any_ident seq in
      let vars, seq = more any_ident seq in
      let _, seq = token Right_arrow seq in
      let e, seq = expr min_int seq in
      Nfunction (v0::vars, e), seq)

(* ifexpr := "if" expr "then" expr "else" expr *)
and ifexpr seq =
  let _, seq = ident "if" seq in
  seq |> mustbe (fun seq ->
      let co, seq = expr min_int seq in
      let _, seq = ident "then" seq in
      let th, seq = expr min_int seq in
      let _, seq = ident "else" seq in
      let el, seq = expr min_int seq in
      If (co, th, el), seq)

(* apply := value [value]+ *)
and apply seq =
  let v, seq = value seq in
  let a0, seq = value seq in
  let args, seq = more value seq in
  Apply (v, a0::args), seq

(* extern_value := external ident *)
and extern_value seq =
  let _, seq = ident "external" seq in
  let id, seq = mustbe any_ident seq in
  Nexternal id, seq

(* unary := - value *)
and unary seq =
  let op, seq =
    let col, x, seq = any seq in
    if x = Minus
    then "-", seq
    else expected col "minus (-)"
  in
  let v, seq = mustbe 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 ((col, x), seq) -> begin match x with
      | Ident id when is_keyword id -> expected col "value"
      | Ident "true" -> Nbool true, seq
      | Ident "false" -> Nbool false, seq
      | Ident id -> Var id, seq
      | Int x -> Nint x, seq
      | Float x -> Nfloat x, seq
      | String x -> Nstring x, seq
      | Hash ->
        let _, t, seq = any seq in
        Nsymbol (Token.to_string t), seq
      | LParen ->
        seq |> either
          (fun seq ->
             let _, seq = token RParen seq in
             Nunit, seq)
          (fun seq ->
             let e, seq = mustbe (expr min_int) seq in
             let _, seq = mustbe (token RParen) seq in
             e, seq)
      | _ -> unexpected_token col x
    end

(* binop := binop op binop *)
and binop pre left seq =
  match seq () with
  | Seq.Nil -> left, Seq.empty
  | Seq.Cons ((col, 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 = mustbe (expr op_pre) seq in
          binop pre (Ast.binop left op right) seq
        else
          left, Seq.cons (col, x) seq

      | RParen -> left, Seq.cons (col, x) seq
      | Ident id when is_keyword id -> left, Seq.cons (col, x) seq
      | _ -> unexpected_token col x
    end

(* parse tokens *)
let parse ts =
  let ast, rest = try decl ts with Fatal e -> raise e in
  if rest () <> Seq.Nil then failwith "Parser.parse";
  ast