(** "Handwritten" structural line lexer.

    Classifies physical lines into high-level structural tokens that the Menhir
    grammar will consume. The lexer is context-aware: it knows the TODO keywords
    from the pre-scan to assist heading classification (though heading title
    parsing is deferred to normalization).

    The lexer operates line-by-line. Each call to [next_token] returns the
    classification of the next line. Tokens are [Parser.token] values directly.
*)

type state = {
  lines : string array;
  mutable line_idx : int;
  mutable line_number : int;  (** 1-based line number *)
  config : Config.t;
}
(** Lexer state. *)

(** Create a new lexer state from input text. *)
let create ~config input =
  let lines = String.split_on_char '\n' input |> Array.of_list in
  { lines; line_idx = 0; line_number = 1; config }

(** Get the current (1-based) line number. *)
let line_number st = st.line_number

(** Check if a line is entirely whitespace. *)
let is_blank_line s =
  let len = String.length s in
  let rec loop i =
    if i >= len then true
    else match s.[i] with ' ' | '\t' | '\r' -> loop (i + 1) | _ -> false
  in
  loop 0

(** Count leading spaces (tabs = 1 for simplicity in indentation). *)
let leading_indent s =
  let len = String.length s in
  let rec loop i =
    if i >= len then i
    else match s.[i] with ' ' -> loop (i + 1) | '\t' -> loop (i + 8) | _ -> i
  in
  loop 0

(** Check if string starts with prefix at position [pos]. *)
let starts_at s pos prefix =
  let plen = String.length prefix in
  if pos + plen > String.length s then false
  else
    let rec loop i =
      if i >= plen then true
      else if s.[pos + i] <> prefix.[i] then false
      else loop (i + 1)
    in
    loop 0

(** Case-insensitive prefix check at position. *)
let starts_at_ci s pos prefix =
  let plen = String.length prefix in
  if pos + plen > String.length s then false
  else
    let rec loop i =
      if i >= plen then true
      else
        let c1 = Char.lowercase_ascii s.[pos + i] in
        let c2 = Char.lowercase_ascii prefix.[i] in
        if c1 <> c2 then false else loop (i + 1)
    in
    loop 0

let rstrip = String_util.rstrip

(** Affiliated keyword names per Org spec defaults. *)
let affiliated_keywords =
  [ "CAPTION"; "DATA"; "HEADER"; "HEADERS"; "NAME"; "PLOT"; "RESULTS" ]

let attr_prefix = "ATTR_"

(** Check if a keyword key is an affiliated keyword. Affiliated keywords:
    CAPTION, DATA, HEADER, HEADERS, NAME, PLOT, RESULTS, ATTR_* *)
let is_affiliated_key key =
  let ukey = String.uppercase_ascii key in
  List.mem ukey affiliated_keywords
  || (String.length ukey > 5 && String.sub ukey 0 5 = attr_prefix)

(** Try to classify a line as a keyword: #+KEY: VALUE Returns Some (key, value)
    or None. *)
let try_keyword line start =
  let len = String.length line in
  if start + 1 >= len || line.[start] <> '#' || line.[start + 1] <> '+' then
    None
  else
    let rest_start = start + 2 in
    (* Find the colon *)
    let rec find_colon i =
      if i >= len then None
      else if line.[i] = ':' then Some i
      else if line.[i] = ' ' || line.[i] = '\t' then
        None (* space before colon = not keyword *)
      else find_colon (i + 1)
    in
    match find_colon rest_start with
    | None -> None
    | Some colon_pos ->
        let key = String.sub line rest_start (colon_pos - rest_start) in
        let value =
          if colon_pos + 1 < len then
            let v = String.sub line (colon_pos + 1) (len - colon_pos - 1) in
            (* Strip leading space from value *)
            if String.length v > 0 && v.[0] = ' ' then
              String.sub v 1 (String.length v - 1)
            else v
          else ""
        in
        Some (key, rstrip value)

(** Try to parse a #+begin_NAME line. Returns Some (name_lowercase,
    params_option) or None. *)
let try_begin_block line start =
  let len = String.length line in
  if not (starts_at_ci line start "#+begin_") then None
  else
    let name_start = start + 8 in
    (* Find end of name (next space or end of line) *)
    let rec find_end i =
      if i >= len then i
      else if line.[i] = ' ' || line.[i] = '\t' then i
      else find_end (i + 1)
    in
    let name_end = find_end name_start in
    if name_end = name_start then None (* empty name *)
    else
      let name =
        String.lowercase_ascii
          (String.sub line name_start (name_end - name_start))
      in
      let params =
        if name_end < len then
          let p = rstrip (String.sub line name_end (len - name_end)) in
          let p =
            if String.length p > 0 && p.[0] = ' ' then
              String.sub p 1 (String.length p - 1)
            else p
          in
          if String.length p > 0 then Some p else None
        else None
      in
      Some (name, params)

(** Try to parse a #+end_NAME line. *)
let try_end_block line start =
  let len = String.length line in
  if not (starts_at_ci line start "#+end_") then None
  else
    let name_start = start + 6 in
    let rec find_end i =
      if i >= len then i
      else if line.[i] = ' ' || line.[i] = '\t' then i
      else find_end (i + 1)
    in
    let name_end = find_end name_start in
    if name_end = name_start then None
    else
      Some
        (String.lowercase_ascii
           (String.sub line name_start (name_end - name_start)))

(** Try to parse a drawer begin line: :NAME: *)
let try_drawer_begin line start =
  let len = String.length line in
  if start >= len || line.[start] <> ':' then None
  else
    (* Must end with : and contain only word chars, hyphens, underscores between *)
    let rec find_end i =
      if i >= len then None
      else if line.[i] = ':' then
        (* Check rest is blank *)
        let rest = rstrip (String.sub line (i + 1) (len - i - 1)) in
        if String.length rest = 0 then Some i else None
      else if
        (line.[i] >= 'a' && line.[i] <= 'z')
        || (line.[i] >= 'A' && line.[i] <= 'Z')
        || (line.[i] >= '0' && line.[i] <= '9')
        || line.[i] = '-'
        || line.[i] = '_'
      then find_end (i + 1)
      else None
    in
    match find_end (start + 1) with
    | None -> None
    | Some end_pos ->
        let name = String.sub line (start + 1) (end_pos - start - 1) in
        if String.length name = 0 then None else Some name

(** Check if a line is :end: (drawer end). *)
let is_drawer_end line start =
  starts_at_ci line start ":end:"
  &&
  let rest = String.sub line (start + 5) (String.length line - start - 5) in
  is_blank_line rest

(** Check if a line is a property: :NAME: VALUE or :NAME+: VALUE *)
let try_property line start =
  let len = String.length line in
  if start >= len || line.[start] <> ':' then None
  else
    let rec find_colon i =
      if i >= len then None
      else if line.[i] = ':' then Some i
      else if
        (line.[i] >= 'a' && line.[i] <= 'z')
        || (line.[i] >= 'A' && line.[i] <= 'Z')
        || (line.[i] >= '0' && line.[i] <= '9')
        || line.[i] = '-'
        || line.[i] = '_'
        || line.[i] = '+'
      then find_colon (i + 1)
      else None
    in
    match find_colon (start + 1) with
    | None -> None
    | Some colon_pos ->
        let name = String.sub line (start + 1) (colon_pos - start - 1) in
        if String.length name = 0 then None
        else
          let value_start = colon_pos + 1 in
          let value =
            if value_start < len then
              let v =
                rstrip (String.sub line value_start (len - value_start))
              in
              let v =
                if String.length v > 0 && v.[0] = ' ' then
                  String.sub v 1 (String.length v - 1)
                else v
              in
              if String.length v > 0 then Some v else None
            else None
          in
          Some (name, value)

(** Try to classify a line as a list item. List items start with: - item, +
    item, * item (not at column 0), 1. item, 1) item, a. item, a) item *)
let try_list_item line =
  let len = String.length line in
  let indent = leading_indent line in
  if indent >= len then None
  else
    let pos = ref indent in
    (* Try bullet markers *)
    let bullet_result =
      if !pos >= len then None
      else
        let c = line.[!pos] in
        if (c = '-' || c = '+') && !pos + 1 < len && line.[!pos + 1] = ' ' then begin
          pos := !pos + 2;
          Some (String.make 1 c ^ " ")
        end
        else if c = '*' && indent > 0 && !pos + 1 < len && line.[!pos + 1] = ' '
        then begin
          (* * at column >0 is a list bullet; at column 0 it's a heading *)
          pos := !pos + 2;
          Some "* "
        end
        else begin
          (* Try ordered: digit(s) followed by . or ) then space *)
          let start = !pos in
          let rec scan_digits i =
            if i >= len then None
            else if line.[i] >= '0' && line.[i] <= '9' then scan_digits (i + 1)
            else if
              i > start
              && (line.[i] = '.' || line.[i] = ')')
              && i + 1 < len
              && line.[i + 1] = ' '
            then begin
              let bullet_text = String.sub line start (i - start + 1) ^ " " in
              pos := i + 2;
              Some bullet_text
            end
            else if
              i = start
              && ((line.[i] >= 'a' && line.[i] <= 'z')
                 || (line.[i] >= 'A' && line.[i] <= 'Z'))
              && i + 1 < len
              && (line.[i + 1] = '.' || line.[i + 1] = ')')
              && i + 2 < len
              && line.[i + 2] = ' '
            then begin
              let bullet_text = String.sub line start 2 ^ " " in
              pos := i + 3;
              Some bullet_text
            end
            else None
          in
          scan_digits start
        end
    in
    match bullet_result with
    | None -> None
    | Some bullet ->
        (* Check for counter set: [@N] *)
        let counter_set =
          if !pos + 2 < len && line.[!pos] = '[' && line.[!pos + 1] = '@' then begin
            let start = !pos + 2 in
            let rec find_bracket i =
              if i >= len then None
              else if line.[i] = ']' then
                let num_str = String.sub line start (i - start) in
                try
                  pos := i + 1;
                  if !pos < len && line.[!pos] = ' ' then incr pos;
                  Some (int_of_string num_str)
                with Failure _ -> None
              else find_bracket (i + 1)
            in
            find_bracket start
          end
          else None
        in
        (* Check for checkbox: [ ], [X], [-] *)
        let checkbox =
          if
            !pos + 2 < len
            && line.[!pos] = '['
            && (line.[!pos + 1] = ' '
               || line.[!pos + 1] = 'X'
               || line.[!pos + 1] = 'x'
               || line.[!pos + 1] = '-')
            && line.[!pos + 2] = ']'
          then begin
            let cb = String.sub line !pos 3 in
            pos := !pos + 3;
            if !pos < len && line.[!pos] = ' ' then incr pos;
            Some cb
          end
          else None
        in
        let rest =
          if !pos < len then String.sub line !pos (len - !pos) else ""
        in
        Some
          {
            Token.lid_indent = indent;
            lid_bullet = bullet;
            lid_counter_set = counter_set;
            lid_checkbox = checkbox;
            lid_rest = rstrip rest;
          }

(** Check if a line is a horizontal rule: 5+ hyphens, nothing else. *)
let is_horizontal_rule line start =
  let len = String.length line in
  let rec count i =
    if i >= len then i - start
    else if line.[i] = '-' then count (i + 1)
    else if line.[i] = ' ' || line.[i] = '\t' || line.[i] = '\r' then
      (* trailing whitespace ok *)
      let rest = rstrip (String.sub line i (len - i)) in
      if String.length rest = 0 then i - start else 0
    else 0
  in
  count start >= 5

(** Check if a line is a planning line (starts with DEADLINE:, SCHEDULED:, or
    CLOSED:). *)
let is_planning_line line start =
  starts_at line start "DEADLINE:"
  || starts_at line start "SCHEDULED:"
  || starts_at line start "CLOSED:"

(** Check if a line is a table row (starts with |). *)
let is_table_row line start = start < String.length line && line.[start] = '|'

(** Check if a line is a comment: # followed by space or end of line. *)
let try_comment_line line start =
  let len = String.length line in
  if start >= len || line.[start] <> '#' then None
  else if start + 1 >= len then Some ""
  else if line.[start + 1] = ' ' then
    Some (rstrip (String.sub line (start + 2) (len - start - 2)))
  else if line.[start + 1] = '+' then None (* keyword, not comment *)
  else None (* # followed by non-space non-+ is not a comment *)

(** Check if a line is fixed-width: ": " or ":" at end of line. *)
let try_fixed_width line start =
  let len = String.length line in
  if start >= len || line.[start] <> ':' then None
  else if start + 1 >= len then Some ""
  else if line.[start + 1] = ' ' then
    Some (rstrip (String.sub line (start + 2) (len - start - 2)))
  else None

(** Check if a line is a heading: starts with one or more * followed by space.
*)
let try_heading line =
  let len = String.length line in
  if len = 0 || line.[0] <> '*' then None
  else
    let rec count_stars i =
      if i >= len then i else if line.[i] = '*' then count_stars (i + 1) else i
    in
    let stars = count_stars 0 in
    if stars >= len then
      (* Line is all stars with no space — it's a heading with empty title *)
      None (* Actually per spec, space after stars is mandatory *)
    else if line.[stars] = ' ' then
      let rest = String.sub line (stars + 1) (len - stars - 1) in
      Some (stars, rstrip rest)
    else None

(** Classify the next line and return its token. *)
let classify_line _st line =
  let line = rstrip line in
  if is_blank_line line then Parser.BLANK
  else
    let indent = leading_indent line in
    let start = indent in
    (* Headings must be at column 0 *)
    if indent = 0 then
      begin match try_heading line with
      | Some (level, rest) -> Parser.HEADING (level, rest)
      | None -> (
          (* Try keyword/block *)
          match try_begin_block line start with
          | Some (name, params) -> Parser.BEGIN_BLOCK (name, params)
          | None -> (
              match try_end_block line start with
              | Some name -> Parser.END_BLOCK name
              | None -> (
                  match try_keyword line start with
                  | Some (key, value) ->
                      if is_affiliated_key key then
                        Parser.AFFILIATED_KEYWORD
                          (String.uppercase_ascii key, None, value)
                      else Parser.KEYWORD (String.uppercase_ascii key, value)
                  | None -> (
                      if is_drawer_end line start then Parser.DRAWER_END
                      else
                        match try_drawer_begin line start with
                        | Some name -> Parser.DRAWER_BEGIN name
                        | None -> (
                            match try_property line start with
                            | Some (name, value) -> Parser.PROPERTY (name, value)
                            | None -> (
                                if is_horizontal_rule line start then
                                  Parser.HORIZONTAL_RULE
                                else if is_planning_line line start then
                                  Parser.PLANNING line
                                else if is_table_row line start then
                                  Parser.TABLE_ROW line
                                else
                                  match try_comment_line line start with
                                  | Some content -> Parser.COMMENT_LINE content
                                  | None -> (
                                      match try_fixed_width line start with
                                      | Some content ->
                                          Parser.FIXED_WIDTH content
                                      | None -> (
                                          match try_list_item line with
                                          | Some data -> Parser.LIST_ITEM data
                                          | None -> Parser.TEXT_LINE line)))))))
          )
      end
    else
      (* Indented lines *)
      begin if is_planning_line line start then Parser.PLANNING line
      else if is_table_row line start then Parser.TABLE_ROW line
      else if is_drawer_end line start then Parser.DRAWER_END
      else
        match try_drawer_begin line start with
        | Some name -> Parser.DRAWER_BEGIN name
        | None -> (
            match try_begin_block line start with
            | Some (name, params) -> Parser.BEGIN_BLOCK (name, params)
            | None -> (
                match try_end_block line start with
                | Some name -> Parser.END_BLOCK name
                | None -> (
                    match try_keyword line start with
                    | Some (key, value) ->
                        if is_affiliated_key key then
                          Parser.AFFILIATED_KEYWORD
                            (String.uppercase_ascii key, None, value)
                        else Parser.KEYWORD (String.uppercase_ascii key, value)
                    | None -> (
                        match try_property line start with
                        | Some (name, value) -> Parser.PROPERTY (name, value)
                        | None -> (
                            match try_comment_line line start with
                            | Some content -> Parser.COMMENT_LINE content
                            | None -> (
                                match try_fixed_width line start with
                                | Some content -> Parser.FIXED_WIDTH content
                                | None -> (
                                    match try_list_item line with
                                    | Some data -> Parser.LIST_ITEM data
                                    | None -> Parser.TEXT_LINE line)))))))
      end

(** Get the next token from the lexer. *)
let next_token st =
  if st.line_idx >= Array.length st.lines then Parser.EOF
  else begin
    let line = st.lines.(st.line_idx) in
    st.line_idx <- st.line_idx + 1;
    st.line_number <- st.line_number + 1;
    classify_line st line
  end

(** Tokenize the entire input into a list of (token, line_number) pairs. Useful
    for testing and for feeding Menhir. *)
let tokenize ~config input =
  let st = create ~config input in
  let rec loop acc =
    let ln = st.line_number in
    let tok = next_token st in
    match tok with
    | Parser.EOF -> List.rev ((Parser.EOF, ln) :: acc)
    | _ -> loop ((tok, ln) :: acc)
  in
  loop []