open Common
open Parser

let lexeme = Sedlexing.Utf8.lexeme
let digit = [%sedlex.regexp? '0' .. '9']
let sign = [%sedlex.regexp? '-' | '+']
let int_literal = [%sedlex.regexp? Opt sign, Plus digit]
let frac = [%sedlex.regexp? Plus digit, '.', Star digit]
let ne_frac = [%sedlex.regexp? '.', Plus digit]

let float_literal =
  [%sedlex.regexp?
    ( Opt sign, frac
    | Opt sign, ne_frac
    | ( Opt sign,
        ('0' | Star digit),
        Opt ne_frac,
        ('e' | 'E'),
        Opt sign,
        Plus digit ) )]

let char_literal = [%sedlex.regexp? '\'', Star (Compl '\''), '\'']
let lower = [%sedlex.regexp? 'a' .. 'z']
let upper = [%sedlex.regexp? 'A' .. 'Z']
let alpha = [%sedlex.regexp? lower | upper]
let id0 = [%sedlex.regexp? lower | '_']
let id = [%sedlex.regexp? id0, Star (alpha | digit | '_')]
let case_id = [%sedlex.regexp? upper, Star (alpha | digit | '_')]
let whitespace = [%sedlex.regexp? Plus ('\t' | ' ')]
let newline = [%sedlex.regexp? '\r' | '\n' | "\r\n"]
let lparen = [%sedlex.regexp? '(']
let rparen = [%sedlex.regexp? ')']
let lbracket = [%sedlex.regexp? '[']
let rbracket = [%sedlex.regexp? ']']
let dot = [%sedlex.regexp? '.']
let plus = [%sedlex.regexp? '+']
let minus = [%sedlex.regexp? '-']
let times = [%sedlex.regexp? '*']
let div = [%sedlex.regexp? '/']
let mod_ = [%sedlex.regexp? "mod"]
let equal = [%sedlex.regexp? '=']
let fplus = [%sedlex.regexp? "+."]
let fminus = [%sedlex.regexp? "-."]
let ftimes = [%sedlex.regexp? "*."]
let fdiv = [%sedlex.regexp? "/."]
let fequal = [%sedlex.regexp? "=."]
let fless = [%sedlex.regexp? "<."]
let flesseq = [%sedlex.regexp? "<=."]
let fgreater = [%sedlex.regexp? ">."]
let fgreatereq = [%sedlex.regexp? ">=."]
let fmax = [%sedlex.regexp? "float_max"]
let fmin = [%sedlex.regexp? "float_min"]
let fnegate = [%sedlex.regexp? "float_negate"]
let fsqrt = [%sedlex.regexp? "sqrt"]
let fsin = [%sedlex.regexp? "sin"]
let fcos = [%sedlex.regexp? "cos"]
let fexp = [%sedlex.regexp? "exp"]
let flog = [%sedlex.regexp? "log"]
let fround = [%sedlex.regexp? "round"]
let float_of_int = [%sedlex.regexp? "float_of_int"]
let int_of_float = [%sedlex.regexp? "int_of_float"]
let and_ = [%sedlex.regexp? "and"]
let pipe = [%sedlex.regexp? "|"]
let barbar = [%sedlex.regexp? "||"]
let ampersandampersand = [%sedlex.regexp? "&&"]
let or_ = [%sedlex.regexp? "or"]
let not = [%sedlex.regexp? '!']
let less = [%sedlex.regexp? '<']
let lesseq = [%sedlex.regexp? "<="]
let greater = [%sedlex.regexp? '>']
let greatereq = [%sedlex.regexp? ">="]
let arrow = [%sedlex.regexp? "->"]
let if_ = [%sedlex.regexp? "if"]
let then_ = [%sedlex.regexp? "then"]
let else_ = [%sedlex.regexp? "else"]
let fun_ = [%sedlex.regexp? "fun"]
let let_ = [%sedlex.regexp? "let"]
let rec_ = [%sedlex.regexp? "rec"]
let in_ = [%sedlex.regexp? "in"]
let for_ = [%sedlex.regexp? "for"]
let to_ = [%sedlex.regexp? "to"]
let do_ = [%sedlex.regexp? "do"]
let done_ = [%sedlex.regexp? "done"]
let true_ = [%sedlex.regexp? "true"]
let false_ = [%sedlex.regexp? "false"]
let vec = [%sedlex.regexp? "vec"]
let vector = [%sedlex.regexp? "vector"]
let vector_get = [%sedlex.regexp? "vector_get"]
let vector_set = [%sedlex.regexp? "vector_set"]
let vector_length = [%sedlex.regexp? "vector_length"]
let print_ = [%sedlex.regexp? "print"]
let print_int = [%sedlex.regexp? "print_int"]
let print_bool = [%sedlex.regexp? "print_bool"]
let print_float = [%sedlex.regexp? "print_float"]
let print_char = [%sedlex.regexp? "print_char"]
let read_int = [%sedlex.regexp? "read_int"]
let read_float = [%sedlex.regexp? "read_float"]
let read_char = [%sedlex.regexp? "read_char"]
let time = [%sedlex.regexp? "time"]
let type_ = [%sedlex.regexp? "type"]
let of_ = [%sedlex.regexp? "of"]
let comma = [%sedlex.regexp? ',']
let dot = [%sedlex.regexp? '.']
let colon = [%sedlex.regexp? ':']
let coloncolon = [%sedlex.regexp? "::"]
let semicolon = [%sedlex.regexp? ';']
let bool = [%sedlex.regexp? "bool"]
let int = [%sedlex.regexp? "int"]
let float = [%sedlex.regexp? "float"]
let question_mark = [%sedlex.regexp? "any" | "?"]
let match_ = [%sedlex.regexp? "match"]
let with_ = [%sedlex.regexp? "with"]
let comment_open = [%sedlex.regexp? "(", Plus '*']
let comment_close = [%sedlex.regexp? Plus '*', ")"]

let rec token lexbuf =
  match%sedlex lexbuf with
  | lparen -> LPAREN
  | rparen -> RPAREN
  | lbracket -> LBRACKET
  | rbracket -> RBRACKET
  | fplus -> FPLUS
  | fminus -> FMINUS
  | ftimes -> FTIMES
  | fdiv -> FDIV
  | fequal -> FEQUAL
  | fless -> FLESS
  | flesseq -> FLESSEQ
  | fgreater -> FGREATER
  | fgreatereq -> FGREATEREQ
  | plus -> PLUS
  | minus -> MINUS
  | times -> TIMES
  | div -> DIV
  | mod_ -> MOD
  | equal -> EQUAL
  | and_ -> AND
  | pipe -> PIPE
  | barbar -> BARBAR
  | ampersandampersand -> AMPERSANDAMPERSAND
  | not -> NOT
  | less -> LESS
  | lesseq -> LESSEQ
  | greater -> GREATER
  | greatereq -> GREATEREQ
  | arrow -> ARROW
  | bool -> TBOOL
  | int -> TINT
  | float -> TFLOAT
  | question_mark -> TANY
  | if_ -> IF
  | then_ -> THEN
  | else_ -> ELSE
  | fun_ -> FUN
  | let_ -> LET
  | rec_ -> REC
  | in_ -> IN
  | for_ -> FOR
  | to_ -> TO
  | do_ -> DO
  | done_ -> DONE
  | true_ -> TRUE
  | false_ -> FALSE
  | vec -> VEC
  | vector -> VECTOR
  | vector_get -> VECTOR_GET
  | vector_set -> VECTOR_SET
  | vector_length -> VECTOR_LENGTH
  | print_ -> PRINT
  | print_int -> PRINT_INT
  | print_bool -> PRINT_BOOL
  | print_float -> PRINT_FLOAT
  | print_char -> PRINT_CHAR
  | read_int -> READ_INT
  | read_float -> READ_FLOAT
  | read_char -> READ_CHAR
  | fmax -> FMAX
  | fmin -> FMIN
  | fnegate -> FNEGATE
  | fsqrt -> FSQRT
  | fsin -> FSIN
  | fcos -> FCOS
  | fexp -> FEXP
  | flog -> FLOG
  | fround -> FROUND
  | float_of_int -> FLOAT_OF_INT
  | int_of_float -> INT_OF_FLOAT
  | time -> TIME
  | type_ -> TYPE
  | of_ -> OF
  | comma -> COMMA
  | dot -> DOT
  | colon -> COLON
  | coloncolon -> COLONCOLON
  | semicolon -> SEMICOLON
  | match_ -> MATCH
  | with_ -> WITH
  | id -> ID (Sedlexing.Utf8.lexeme lexbuf)
  | case_id -> CASE_ID (Sedlexing.Utf8.lexeme lexbuf)
  | int_literal -> INT (int_of_string (Sedlexing.Utf8.lexeme lexbuf))
  | float_literal -> FLOAT (float_of_string (Sedlexing.Utf8.lexeme lexbuf))
  | char_literal -> (
      let loc = Loc.loc_of_lexbuf lexbuf in
      let str = Sedlexing.Utf8.lexeme lexbuf in
      let str = String.sub str 1 (String.length str - 2) in
      let unescaped = Scanf.unescaped str in
      let char_opt =
        if String.length unescaped = 1 then Some unescaped.[0] else None
      in
      match char_opt with
      | Some c -> CHAR c
      | None ->
          raise (Error.error_with_loc (Syntax (InvalidCharLiteral str)) loc))
  | whitespace -> token lexbuf
  | newline -> token lexbuf
  | comment_open -> skip_comment 1 lexbuf
  | eof -> EOF
  | any ->
      let loc = Loc.loc_of_lexbuf lexbuf in
      let str = Sedlexing.Utf8.lexeme lexbuf in
      raise (Error.error_with_loc (Syntax (UnexpectedToken str)) loc)
  | _ -> failwith "unreachable"

and skip_comment depth lexbuf =
  match%sedlex lexbuf with
  | eof ->
      let loc = Loc.loc_of_lexbuf lexbuf in
      raise (Error.error_with_loc (Syntax UnterminatedComment) loc)
  | comment_open -> skip_comment (depth + 1) lexbuf
  | comment_close ->
      if depth = 1 then token lexbuf else skip_comment (depth - 1) lexbuf
  | any -> skip_comment depth lexbuf
  | _ -> failwith "unreachable"
