{
open Parser
open Location

let error_at lexbuf message =
  let span = span_of_lexing (Lexing.lexeme_start_p lexbuf) (Lexing.lexeme_end_p lexbuf) in
  Diagnostic.error span "%s" message

let error_between start stop message =
  Diagnostic.error (span_of_lexing start stop) "%s" message

let keyword = function
  | "let" -> LET
  | "in" -> IN
  | "fun" -> FUN
  | "if" -> IF
  | "then" -> THEN
  | "else" -> ELSE
  | "type" -> TYPE
  | "input" -> INPUT
  | "collection" -> COLLECTION
  | "query" -> QUERY
  | "filter" -> FILTER
  | "map" -> MAP
  | "sum" -> SUM
  | "count" -> COUNT
  | "true" -> BOOL true
  | "false" -> BOOL false
  | text -> IDENT text
}

let digit = ['0'-'9']
let alpha = ['a'-'z' 'A'-'Z' '_']
let identifier = alpha (alpha | digit | '\'')*

rule token = parse
  | [' ' '\t' '\r'] { token lexbuf }
  | '\n' { Lexing.new_line lexbuf; token lexbuf }
  | "(*" { comment 1 (Lexing.lexeme_start_p lexbuf) lexbuf }
  | digit+ {
      let text = Lexing.lexeme lexbuf in
      try INT (int_of_string text)
      with Failure _ -> error_at lexbuf (Printf.sprintf "integer literal %s does not fit in a machine integer" text)
    }
  | '"' {
      let start = Lexing.lexeme_start_p lexbuf in
      string_literal (Buffer.create 16) start lexbuf
    }
  | identifier { keyword (Lexing.lexeme lexbuf) }
  | "->" { ARROW }
  | "|>" { PIPE }
  | "&&" { AND }
  | "||" { OR }
  | "<=" { LE }
  | ">=" { GE }
  | "<>" { NE }
  | "<" { LT }
  | ">" { GT }
  | "=" { EQ }
  | "+" { PLUS }
  | "-" { MINUS }
  | "*" { STAR }
  | "/" { SLASH }
  | "(" { LPAREN }
  | ")" { RPAREN }
  | "," { COMMA }
  | "." { DOT }
  | "{" { LBRACE }
  | "}" { RBRACE }
  | ":" { COLON }
  | ";" { SEMI }
  | eof { EOF }
  | _ { error_at lexbuf (Printf.sprintf "unexpected character %C" (Lexing.lexeme_char lexbuf 0)) }

and comment depth start = parse
  | "(*" { comment (depth + 1) start lexbuf }
  | "*)" { if depth = 1 then token lexbuf else comment (depth - 1) start lexbuf }
  | '\n' { Lexing.new_line lexbuf; comment depth start lexbuf }
  | eof { error_between start (Lexing.lexeme_end_p lexbuf) "unterminated comment" }
  | _ { comment depth start lexbuf }

and string_literal buffer start = parse
  | '"' { STRING (Buffer.contents buffer) }
  | '\\' { escape buffer lexbuf; string_literal buffer start lexbuf }
  | '\n' { error_between start (Lexing.lexeme_start_p lexbuf) "unterminated string literal" }
  | eof { error_between start (Lexing.lexeme_end_p lexbuf) "unterminated string literal" }
  | _ { Buffer.add_string buffer (Lexing.lexeme lexbuf); string_literal buffer start lexbuf }

and escape buffer = parse
  | 'n' { Buffer.add_char buffer '\n' }
  | 't' { Buffer.add_char buffer '\t' }
  | 'r' { Buffer.add_char buffer '\r' }
  | '"' { Buffer.add_char buffer '"' }
  | '\\' { Buffer.add_char buffer '\\' }
  | _ { error_at lexbuf (Printf.sprintf "unknown escape sequence \\%s" (Lexing.lexeme lexbuf)) }