(** Handwritten inline lexer.

    Tokenizes bounded text (heading titles, paragraph bodies, etc.) into tokens
    suitable for the inline Menhir grammar.

    Key design: emphasis delimiters are only emitted as OPEN/CLOSE tokens when
    they form valid matched pairs. Unmatched delimiters become TEXT. This
    eliminates ambiguity in the Menhir grammar. Tokens are [Inline_parser.token]
    values directly. *)

(** PRE characters that allow an emphasis marker to open. *)
let is_pre ch =
  match ch with
  | ' ' | '\t' | '-' | '(' | '{' | '\'' | '"' | '\n' -> true
  | _ -> false

(** POST characters that allow an emphasis marker to close. *)
let is_post ch =
  match ch with
  | ' ' | '\t' | '-' | '.' | ',' | ';' | ':' | '!' | '?' | ')' | '}' | '\''
  | '"' | '\n' | '[' ->
      true
  | _ -> false

(** Common link protocols for plain link detection. *)
let link_protocols = [ "https://"; "http://"; "ftp://"; "file://"; "mailto:" ]

(** Check if a string starts with a prefix at position [pos]. *)
let starts_with_at s pos prefix =
  let plen = String.length prefix in
  if pos + plen > String.length s then false else String.sub s pos plen = prefix

(** Pre-scan for matched emphasis pairs in the input. Returns a set of positions
    where open/close markers are valid. *)
let find_emphasis_pairs input delim =
  let len = String.length input in
  let opens = ref [] in
  let pairs = Hashtbl.create 16 in
  let i = ref 0 in
  while !i < len do
    if input.[!i] = delim then begin
      (* Check if this could be an opener *)
      let pre_ok = !i = 0 || is_pre input.[!i - 1] in
      let next_not_space =
        !i + 1 < len
        && match input.[!i + 1] with ' ' | '\t' | '\n' -> false | _ -> true
      in
      (* Check if this could be a closer *)
      let post_ok = !i + 1 >= len || is_post input.[!i + 1] in
      let prev_not_space =
        !i > 0
        && match input.[!i - 1] with ' ' | '\t' | '\n' -> false | _ -> true
      in
      if prev_not_space && post_ok && !opens <> [] then begin
        (* This is a closer — match with most recent opener *)
        let open_pos = List.hd !opens in
        opens := List.tl !opens;
        Hashtbl.replace pairs open_pos `Open;
        Hashtbl.replace pairs !i `Close
      end
      else if pre_ok && next_not_space then begin
        (* This is a potential opener *)
        opens := !i :: !opens
      end
    end;
    incr i
  done;
  pairs

type state = {
  input : string;
  len : int;
  mutable pos : int;
  star_pairs : (int, [ `Open | `Close ]) Hashtbl.t;
  slash_pairs : (int, [ `Open | `Close ]) Hashtbl.t;
  under_pairs : (int, [ `Open | `Close ]) Hashtbl.t;
  plus_pairs : (int, [ `Open | `Close ]) Hashtbl.t;
}
(** Lexer state. *)

let create input =
  {
    input;
    len = String.length input;
    pos = 0;
    star_pairs = find_emphasis_pairs input '*';
    slash_pairs = find_emphasis_pairs input '/';
    under_pairs = find_emphasis_pairs input '_';
    plus_pairs = find_emphasis_pairs input '+';
  }

(** Try to match a verbatim/code span. *)
let try_verbatim_code st delim =
  let start = st.pos in
  let pre_ok = start = 0 || is_pre st.input.[start - 1] in
  if not pre_ok then None
  else begin
    let rec find_close i =
      if i >= st.len then None
      else if st.input.[i] = '\n' then None
      else if st.input.[i] = delim then
        let post_ok = i + 1 >= st.len || is_post st.input.[i + 1] in
        if post_ok && i > start + 1 then
          Some (String.sub st.input (start + 1) (i - start - 1), i + 1)
        else find_close (i + 1)
      else find_close (i + 1)
    in
    find_close (start + 1)
  end

(** Try to match a complete timestamp. *)
let try_timestamp st =
  let c = st.input.[st.pos] in
  let closer = if c = '<' then '>' else if c = '[' then ']' else '\x00' in
  if closer = '\x00' then None
  else begin
    let next_pos = st.pos + 1 in
    if next_pos >= st.len then None
    else if not (st.input.[next_pos] >= '0' && st.input.[next_pos] <= '9') then
      None
    else begin
      let rec find_close i =
        if i >= st.len then None
        else if st.input.[i] = closer then Some i
        else if st.input.[i] = '\n' then None
        else find_close (i + 1)
      in
      match find_close next_pos with
      | None -> None
      | Some end_pos ->
          let ts = String.sub st.input st.pos (end_pos - st.pos + 1) in
          (* Check for range *)
          let after = end_pos + 1 in
          if
            after + 2 < st.len
            && st.input.[after] = '-'
            && st.input.[after + 1] = '-'
          then begin
            let range_start = after + 2 in
            if range_start < st.len && st.input.[range_start] = c then
              let next2 = range_start + 1 in
              if
                next2 < st.len
                && st.input.[next2] >= '0'
                && st.input.[next2] <= '9'
              then
                match find_close next2 with
                | Some end2 ->
                    let full = String.sub st.input st.pos (end2 - st.pos + 1) in
                    Some (full, end2 + 1)
                | None -> Some (ts, after)
              else Some (ts, after)
            else Some (ts, after)
          end
          else Some (ts, after)
    end
  end

(** Try to match an angle link: <protocol:path> *)
let try_angle_link st =
  if st.input.[st.pos] <> '<' then None
  else begin
    let start = st.pos + 1 in
    let rec find_close i =
      if i >= st.len then None
      else if st.input.[i] = '>' then
        let content = String.sub st.input start (i - start) in
        if String.contains content ':' && not (String.contains content ' ') then
          Some (content, i + 1)
        else None
      else if st.input.[i] = '\n' || st.input.[i] = ' ' then None
      else find_close (i + 1)
    in
    find_close start
  end

(** Try to match a plain link. *)
let try_plain_link st =
  let rec try_protocols = function
    | [] -> None
    | proto :: rest ->
        if starts_with_at st.input st.pos proto then begin
          let plen = String.length proto in
          let rec find_end i =
            if i >= st.len then i
            else
              match st.input.[i] with
              | ' ' | '\t' | '\n' | '>' | '<' | ']' -> i
              | _ -> find_end (i + 1)
          in
          let end_pos = find_end (st.pos + plen) in
          let end_pos = ref end_pos in
          while
            !end_pos > st.pos + plen
            &&
            match st.input.[!end_pos - 1] with
            | '.' | ',' | ';' | ':' | '!' | '?' | ')' -> true
            | _ -> false
          do
            decr end_pos
          done;
          if !end_pos > st.pos + plen then
            Some (String.sub st.input st.pos (!end_pos - st.pos), !end_pos)
          else None
        end
        else try_protocols rest
  in
  try_protocols link_protocols

(** Try to match a line break: \\\\ at end of line *)
let try_line_break st =
  if
    st.pos + 1 < st.len
    && st.input.[st.pos] = '\\'
    && st.input.[st.pos + 1] = '\\'
  then begin
    let rec skip_space i =
      if i >= st.len then Some i
      else
        match st.input.[i] with
        | ' ' | '\t' -> skip_space (i + 1)
        | '\n' -> Some (i + 1)
        | _ -> None
    in
    skip_space (st.pos + 2)
  end
  else None

(** Get the next token. *)
let next_token st =
  if st.pos >= st.len then Inline_parser.EOF
  else
    let c = st.input.[st.pos] in
    (* Try line break *)
    match try_line_break st with
    | Some new_pos ->
        st.pos <- new_pos;
        Inline_parser.LINE_BREAK
    | None ->
        (* Try verbatim/code *)
        if c = '~' then (
          match try_verbatim_code st '~' with
          | Some (content, new_pos) ->
              st.pos <- new_pos;
              Inline_parser.TILDE content
          | None ->
              st.pos <- st.pos + 1;
              Inline_parser.TEXT "~")
        else if c = '=' then (
          match try_verbatim_code st '=' with
          | Some (content, new_pos) ->
              st.pos <- new_pos;
              Inline_parser.EQUALS content
          | None ->
              st.pos <- st.pos + 1;
              Inline_parser.TEXT "=")
        else if
          (* Try link brackets *)
          c = '[' && st.pos + 1 < st.len && st.input.[st.pos + 1] = '['
        then begin
          st.pos <- st.pos + 2;
          Inline_parser.LINK_OPEN
        end
        else if c = ']' && st.pos + 1 < st.len && st.input.[st.pos + 1] = ']'
        then begin
          st.pos <- st.pos + 2;
          Inline_parser.LINK_CLOSE
        end
        else if c = ']' && st.pos + 1 < st.len && st.input.[st.pos + 1] = '['
        then begin
          st.pos <- st.pos + 2;
          Inline_parser.LINK_SEP
        end
        else if
          (* Try timestamp *)
          c = '<'
          || c = '['
             && st.pos + 1 < st.len
             && st.input.[st.pos + 1] >= '0'
             && st.input.[st.pos + 1] <= '9'
        then
          begin match try_timestamp st with
          | Some (ts, new_pos) ->
              st.pos <- new_pos;
              Inline_parser.TIMESTAMP ts
          | None ->
              if c = '<' then (
                match try_angle_link st with
                | Some (content, new_pos) ->
                    st.pos <- new_pos;
                    Inline_parser.ANGLE_LINK content
                | None ->
                    st.pos <- st.pos + 1;
                    Inline_parser.TEXT "<")
              else begin
                st.pos <- st.pos + 1;
                Inline_parser.TEXT "["
              end
          end
        else
          (* Try plain link *)
          begin match try_plain_link st with
          | Some (link, new_pos) ->
              st.pos <- new_pos;
              Inline_parser.PLAIN_LINK link
          | None ->
              (* Try emphasis markers — only emit OPEN/CLOSE if pre-scan found a pair *)
              if c = '*' then begin
                let tok =
                  match Hashtbl.find_opt st.star_pairs st.pos with
                  | Some `Open -> Inline_parser.STAR_OPEN
                  | Some `Close -> Inline_parser.STAR_CLOSE
                  | None -> Inline_parser.TEXT "*"
                in
                st.pos <- st.pos + 1;
                tok
              end
              else if c = '/' then begin
                let tok =
                  match Hashtbl.find_opt st.slash_pairs st.pos with
                  | Some `Open -> Inline_parser.SLASH_OPEN
                  | Some `Close -> Inline_parser.SLASH_CLOSE
                  | None -> Inline_parser.TEXT "/"
                in
                st.pos <- st.pos + 1;
                tok
              end
              else if c = '_' then begin
                let tok =
                  match Hashtbl.find_opt st.under_pairs st.pos with
                  | Some `Open -> Inline_parser.UNDER_OPEN
                  | Some `Close -> Inline_parser.UNDER_CLOSE
                  | None -> Inline_parser.TEXT "_"
                in
                st.pos <- st.pos + 1;
                tok
              end
              else if c = '+' then begin
                let tok =
                  match Hashtbl.find_opt st.plus_pairs st.pos with
                  | Some `Open -> Inline_parser.PLUS_OPEN
                  | Some `Close -> Inline_parser.PLUS_CLOSE
                  | None -> Inline_parser.TEXT "+"
                in
                st.pos <- st.pos + 1;
                tok
              end
              else if c = '\n' then begin
                st.pos <- st.pos + 1;
                Inline_parser.NEWLINE
              end
              else begin
                (* Accumulate plain text *)
                let start = st.pos in
                let rec scan i =
                  if i >= st.len then i
                  else
                    match st.input.[i] with
                    | '*' | '/' | '_' | '+' | '~' | '=' | '[' | ']' | '<' | '\n'
                    | '\\' ->
                        i
                    | _ ->
                        if
                          List.exists
                            (fun proto -> starts_with_at st.input i proto)
                            link_protocols
                        then i
                        else scan (i + 1)
                in
                let end_pos = scan (start + 1) in
                st.pos <- end_pos;
                Inline_parser.TEXT (String.sub st.input start (end_pos - start))
              end
          end

(** Tokenize an entire string into a list of tokens. *)
let tokenize input =
  let st = create input in
  let rec loop acc =
    let tok = next_token st in
    match tok with
    | Inline_parser.EOF -> List.rev (Inline_parser.EOF :: acc)
    | _ -> loop (tok :: acc)
  in
  loop []