type token =
  | TBackslash
  | TDot
  | TLParen
  | TRParen
  | TLBrack
  | TRBrack
  | TColonEq
  | TVar of string

exception Parse_error of string

let is_var_start c = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')

let is_var_char c =
  is_var_start c || (c >= '0' && c <= '9') || c = '_' || c = '\'' || c = '-'

let lex s =
  let n = String.length s in
  let rec go i acc =
    if i >= n then List.rev acc
    else
      match s.[i] with
      | ' ' | '\t' | '\n' | '\r' -> go (i + 1) acc
      | '\\' -> go (i + 1) (TBackslash :: acc)
      | '.' -> go (i + 1) (TDot :: acc)
      | '(' -> go (i + 1) (TLParen :: acc)
      | ')' -> go (i + 1) (TRParen :: acc)
      | '[' -> go (i + 1) (TLBrack :: acc)
      | ']' -> go (i + 1) (TRBrack :: acc)
      | ':' ->
          if i + 1 < n && s.[i + 1] = '=' then go (i + 2) (TColonEq :: acc)
          else raise (Parse_error ("unexpected ':' at position " ^ string_of_int i))
      | c when is_var_start c ->
          let j = ref (i + 1) in
          while !j < n && is_var_char s.[!j] do
            incr j
          done;
          go !j (TVar (String.sub s i (!j - i)) :: acc)
      | c ->
          raise
            (Parse_error
               (Printf.sprintf "unexpected character %C at position %d" c i))
  in
  go 0 []

let parse s =
  let toks = Array.of_list (lex s) in
  let n = Array.length toks in
  let pos = ref 0 in
  let peek () = if !pos < n then Some toks.(!pos) else None in
  let next () =
    let t = peek () in
    incr pos;
    t
  in
  let expect tok what =
    match next () with
    | Some t when t = tok -> ()
    | _ -> raise (Parse_error ("expected " ^ what))
  in
  let starts_atom = function
    | Some (TVar _) | Some TLParen | Some TBackslash | Some TLBrack -> true
    | _ -> false
  in
  let rec term () =
    match peek () with
    | Some TBackslash -> lam ()
    | Some TLBrack -> subst ()
    | _ -> app ()
  and lam () =
    ignore (next ());
    match next () with
    | Some (TVar v) ->
        expect TDot "'.'";
        let body = term () in
        Named.Lam (v, body)
    | _ -> raise (Parse_error "expected variable after '\\'")
  and subst () =
    ignore (next ());
    match next () with
    | Some (TVar v) ->
        expect TColonEq "':='";
        let s = term () in
        expect TRBrack "']'";
        let body = term () in
        Named.Subst (v, s, body)
    | _ -> raise (Parse_error "expected variable after '['")
  and app () =
    let lhs = atom () in
    let rec args lhs =
      if starts_atom (peek ()) then args (Named.App (lhs, atom ())) else lhs
    in
    args lhs
  and atom () =
    match next () with
    | Some (TVar v) -> Named.Var v
    | Some TLParen ->
        let t = term () in
        expect TRParen "')'";
        t
    | Some TBackslash -> begin
        match next () with
        | Some (TVar v) ->
            expect TDot "'.'";
            Named.Lam (v, term ())
        | _ -> raise (Parse_error "expected variable after '\\'")
      end
    | Some TLBrack -> begin
        match next () with
        | Some (TVar v) ->
            expect TColonEq "':='";
            let s = term () in
            expect TRBrack "']'";
            Named.Subst (v, s, term ())
        | _ -> raise (Parse_error "expected variable after '['")
      end
    | _ -> raise (Parse_error "expected a term")
  in
  let t = term () in
  if !pos <> n then raise (Parse_error "trailing input");
  t

let parse_result s =
  match parse s with t -> Ok t | exception Parse_error m -> Error m