(* 렉서. grammar.ebnf의 어휘 절을 그대로 구현한다.
   유일하게 문법과 얽히는 부분은 NEWLINE 삽입인데, 그 판정은
   Token.can_end_statement 하나로 끝나며 선읽기를 요구하지 않는다. *)

type error = { pos : Token.pos; msg : string }

exception Error of error

type state = {
  src : string;
  len : int;
  mutable i : int;
  mutable line : int;
  mutable bol : int; (* 현재 줄의 시작 offset *)
  mutable out : Token.t list; (* 역순 누적 *)
}

let pos st : Token.pos = { line = st.line; col = st.i - st.bol + 1 }
let fail st msg = raise (Error { pos = pos st; msg })
let fail_at p msg = raise (Error { pos = p; msg })
let peek st = if st.i < st.len then Some st.src.[st.i] else None
let peek2 st = if st.i + 1 < st.len then Some st.src.[st.i + 1] else None
let emit st p kind = st.out <- { Token.kind; pos = p } :: st.out
let last_kind st = match st.out with [] -> None | t :: _ -> Some t.Token.kind

(* 줄 끝에 도달했을 때만 호출된다. 직전 토큰이 값으로 끝날 수 있으면 삽입.
   이미 Newline을 넣었거나 아무것도 없으면 넣지 않으므로,
   빈 줄과 주석만 있는 줄은 구분자를 만들지 않는다. *)
let maybe_newline st =
  match last_kind st with
  | Some k when Token.can_end_statement k -> emit st (pos st) Token.Newline
  | _ -> ()

let is_letter c = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
let is_digit c = c >= '0' && c <= '9'
let is_ident_rest c = is_letter c || is_digit c || c = '_'

let lex_ident st =
  let p = pos st in
  let start = st.i in
  while st.i < st.len && is_ident_rest st.src.[st.i] do
    st.i <- st.i + 1
  done;
  let s = String.sub st.src start (st.i - start) in
  match Token.keyword s with
  | Some k -> emit st p k
  | None -> emit st p (Token.Ident s)

let lex_int st =
  let p = pos st in
  let start = st.i in
  while st.i < st.len && (is_digit st.src.[st.i] || st.src.[st.i] = '_') do
    st.i <- st.i + 1
  done;
  (* 숫자 바로 뒤에 글자가 붙으면 오타일 가능성이 높다. 조용히 넘기지 않는다. *)
  if st.i < st.len && is_letter st.src.[st.i] then
    fail st "정수 리터럴 뒤에 글자가 올 수 없습니다";
  emit st p (Token.Int (String.sub st.src start (st.i - start)))

let lex_string st =
  let p = pos st in
  st.i <- st.i + 1;
  let buf = Buffer.create 16 in
  let rec go () =
    if st.i >= st.len then fail_at p "문자열이 닫히지 않았습니다"
    else
      match st.src.[st.i] with
      | '"' -> st.i <- st.i + 1
      | '\n' -> fail_at p "문자열이 닫히지 않았습니다 (줄바꿈 전에 닫아야 합니다)"
      | '\\' ->
          st.i <- st.i + 1;
          if st.i >= st.len then fail_at p "문자열이 닫히지 않았습니다";
          let c = st.src.[st.i] in
          let decoded =
            match c with
            | 'n' -> '\n'
            | 't' -> '\t'
            | '\\' -> '\\'
            | '"' -> '"'
            | _ -> fail st (Printf.sprintf "알 수 없는 이스케이프 \\%c" c)
          in
          Buffer.add_char buf decoded;
          st.i <- st.i + 1;
          go ()
      | c ->
          Buffer.add_char buf c;
          st.i <- st.i + 1;
          go ()
  in
  go ();
  emit st p (Token.Str (Buffer.contents buf))

let skip_line_comment st =
  while st.i < st.len && st.src.[st.i] <> '\n' do
    st.i <- st.i + 1
  done

(* 한 글자 또는 두 글자 연산자. 두 글자를 항상 먼저 본다. *)
let lex_punct st =
  let p = pos st in
  let one k =
    st.i <- st.i + 1;
    emit st p k
  in
  let two k =
    st.i <- st.i + 2;
    emit st p k
  in
  match (st.src.[st.i], peek2 st) with
  | '-', Some '>' -> two Token.Arrow
  | '=', Some '>' -> two Token.FatArrow
  | '=', Some '=' -> two Token.EqEq
  | '!', Some '=' -> two Token.BangEq
  | '<', Some '=' -> two Token.Le
  | '>', Some '=' -> two Token.Ge
  | '&', Some '&' -> two Token.AmpAmp
  | '|', Some '|' -> two Token.PipePipe
  | '(', _ -> one Token.LParen
  | ')', _ -> one Token.RParen
  | '{', _ -> one Token.LBrace
  | '}', _ -> one Token.RBrace
  | '[', _ -> one Token.LBracket
  | ']', _ -> one Token.RBracket
  | ',', _ -> one Token.Comma
  | ':', _ -> one Token.Colon
  | '.', _ -> one Token.Dot
  | '?', _ -> one Token.Question
  | '=', _ -> one Token.Eq
  | '!', _ -> one Token.Bang
  | '<', _ -> one Token.Lt
  | '>', _ -> one Token.Gt
  | '+', _ -> one Token.Plus
  | '-', _ -> one Token.Minus
  | '*', _ -> one Token.Star
  | '/', _ -> one Token.Slash
  | '%', _ -> one Token.Percent
  | '|', _ -> one Token.Pipe
  | '&', _ -> fail st "& 연산자는 없습니다 (&&를 의도했습니까?)"
  | c, _ -> fail st (Printf.sprintf "예상치 못한 문자 %C" c)

let lex (src : string) : Token.t list =
  let st =
    { src; len = String.length src; i = 0; line = 1; bol = 0; out = [] }
  in
  let rec go () =
    if st.i >= st.len then (
      maybe_newline st;
      emit st (pos st) Token.Eof)
    else
      let c = st.src.[st.i] in
      match c with
      | ' ' | '\t' | '\r' ->
          st.i <- st.i + 1;
          go ()
      | '\n' ->
          maybe_newline st;
          st.i <- st.i + 1;
          st.line <- st.line + 1;
          st.bol <- st.i;
          go ()
      | '/' when peek2 st = Some '/' ->
          skip_line_comment st;
          go ()
      | '_' ->
          let p = pos st in
          if st.i + 1 < st.len && is_ident_rest st.src.[st.i + 1] then
            fail st "이름은 문자로 시작해야 합니다"
          else (
            st.i <- st.i + 1;
            emit st p Token.Underscore;
            go ())
      | c when is_letter c ->
          lex_ident st;
          go ()
      | c when is_digit c ->
          lex_int st;
          go ()
      | '"' ->
          lex_string st;
          go ()
      | _ ->
          lex_punct st;
          go ()
  in
  go ();
  List.rev st.out

let lex_result src =
  match lex src with tokens -> Ok tokens | exception Error e -> Error e

(* peek는 lex_punct 안에서만 쓰이지 않으므로 경고를 피한다 *)
let _ = peek