open Fmlib_parse
module StringSet = Set.Make (String)
module StringMap = Map.Make (String)

module State = struct
  type t = {
    binders : string list;
    constants : StringSet.t;
    functions : int StringMap.t;
  }
end

module Token = Lexer.Token

module Final = struct
  type t = Kernel.Term.t
end

module Semantic = struct
  type t = ArityMismatch
end

module C = struct
  include Token_parser.Make (State) (Token) (Final) (Semantic)
  open Kernel.Term
  open State

  let ident_p =
    step "Expected identifier" (fun state _ -> function
      | Token.Ident id -> Some (id, state) | _ -> None)

  let expect tt s =
    step ("Expected " ^ s) (fun state _ tt' ->
        if tt = tt' then Some ((), state) else None)

  let rec args_p () =
    let* _ = expect Token.LParen "`('" in
    let* args =
      one_or_more_separated
        (fun x -> return [ x ])
        (fun acc _ t -> return (t :: acc))
        (term_p ()) (expect Comma "comma")
      |> map List.rev
    in
    let* _ = expect Token.RParen "`)'" in
    return args

  and term_p () =
    let* ident = ident_p in
    let* { binders; constants; functions } = get in
    match List.find_index (( = ) ident) binders with
    | Some x -> return @@ Var x
    | None ->
        if StringSet.mem ident constants then return @@ Const ident
        else if StringMap.mem ident functions then
          let* args = args_p () in
          if List.length args <> StringMap.find ident functions then
            fail ArityMismatch
          else return @@ Function (ident, args)
        else return @@ Free ident

  let final = term_p
end

include C.Parser

let token_parser state = C.make state @@ C.final ()