type state =
  | Token
  | String
  | DString
  | MString of int
  | EndString
  | Expression

type t = {
  buffer : Sedlexing.lexbuf;
  mutable start_p : Lexing.position option;
  state : state Stack.t;
  reset_line : bool;
}

let state : t -> state option = fun t -> Stack.top_opt t.state
let enter_state : t -> state -> unit = fun t state -> Stack.push state t.state
let leave_state : t -> unit = fun t -> ignore (Stack.pop_opt t.state)
let buffer : t -> Sedlexing.lexbuf = fun t -> t.buffer

let start : t -> unit =
 fun t ->
  let _start_pos, end_pos = Sedlexing.lexing_positions t.buffer in
  let () =
    if not t.reset_line then
      Sedlexing.set_position t.buffer { end_pos with Lexing.pos_lnum = 1 }
  in
  Stack.clear t.state;
  t.start_p <- None

let positions : t -> Lexing.position * Lexing.position =
 fun t ->
  let default, end_p = Sedlexing.lexing_positions t.buffer in
  let start_p = Option.value ~default t.start_p in
  (start_p, end_p)

let content : t -> string = fun t -> Sedlexing.Utf8.lexeme t.buffer

let from_lexbuf : ?reset_line:bool -> Sedlexing.lexbuf -> t =
 fun ?(reset_line = true) t ->
  { buffer = t; start_p = None; reset_line; state = Stack.create () }

let set_start_position : t -> Lexing.position -> unit =
 fun t position ->
  match t.start_p with
  | None -> t.start_p <- Some position
  | _ ->
      (* We are already inside a block code, don’t stack it *)
      ()

let tokenize : (t -> 'a) -> t -> unit -> 'a * Lexing.position * Lexing.position
    =
 fun f t ->
  let lexer () =
    (* Clear the previous registered start position if any *)
    t.start_p <- None;
    let token = f t in
    let default, curr_p = positions t in

    let start_p = Option.value ~default t.start_p in
    t.start_p <- None;

    (token, start_p, curr_p)
  in
  lexer

let rollback : t -> unit = fun t -> Sedlexing.rollback t.buffer