(** Raw_ast to Ast normalization.

    Transforms the flat [Raw_ast.document] produced by the structural parser
    into the semantic [Ast.document] tree. This includes:
    - building heading hierarchy from flat headings
    - parsing heading titles for TODO/priority/tags/COMMENT
    - inline parsing of text fields
    - list nesting by indentation
    - block classification
    - planning and property extraction
    - affiliated keyword attachment
    - directive collection *)

(* ------------------------------------------------------------------ *)
(* Utility helpers                                                     *)
(* ------------------------------------------------------------------ *)

let string_lowercase = String.lowercase_ascii

(** Trim leading and trailing whitespace from a string. *)
let trim = String.trim

(** Split a string on the first space. *)
let split_first_word s =
  let s = trim s in
  match String.index_opt s ' ' with
  | None -> (s, "")
  | Some i -> (String.sub s 0 i, String.sub s (i + 1) (String.length s - i - 1))

(* ------------------------------------------------------------------ *)
(* Heading title parsing                                               *)
(* ------------------------------------------------------------------ *)

(** Parse tags from the end of a heading title. Tags look like :tag1:tag2: at
    the end of the line. *)
let extract_tags title =
  let title = trim title in
  let len = String.length title in
  if len < 2 || title.[len - 1] <> ':' then (title, [])
  else begin
    (* Find the start of the tags section: look for a space followed by : *)
    let rec find_tag_start i =
      if i <= 0 then None
      else if title.[i] = ' ' || title.[i] = '\t' then
        if i + 1 < len && title.[i + 1] = ':' then Some (i + 1)
        else find_tag_start (i - 1)
      else find_tag_start (i - 1)
    in
    match find_tag_start (len - 2) with
    | None ->
        (* Check if the whole title is a tag string *)
        if title.[0] = ':' then
          let tag_str = String.sub title 1 (len - 2) in
          let tags = String.split_on_char ':' tag_str in
          if List.for_all (fun t -> String.length t > 0) tags then ("", tags)
          else (title, [])
        else (title, [])
    | Some tag_start ->
        let tag_str = String.sub title (tag_start + 1) (len - tag_start - 2) in
        let tags = String.split_on_char ':' tag_str in
        if List.for_all (fun t -> String.length t > 0) tags then
          let title_part = trim (String.sub title 0 tag_start) in
          (title_part, tags)
        else (title, [])
  end

(** Parse priority from the beginning of a heading title. Priority looks like
    [#A] at the start. *)
let extract_priority title =
  let title = trim title in
  if
    String.length title >= 4
    && title.[0] = '['
    && title.[1] = '#'
    && title.[3] = ']'
    && title.[2] >= 'A'
    && title.[2] <= 'Z'
  then
    let rest = trim (String.sub title 4 (String.length title - 4)) in
    (Some title.[2], rest)
  else (None, title)

(** Parse a heading raw title string into its components. Order: TODO
    [#PRIORITY] COMMENT title :tags: *)
let parse_heading_title config raw_title =
  let raw_title = trim raw_title in
  (* Extract tags from end first *)
  let title_no_tags, tags = extract_tags raw_title in
  (* Extract TODO keyword from front *)
  let todo, rest =
    let first_word, remainder = split_first_word title_no_tags in
    if first_word <> "" && Config.is_todo_keyword config first_word then
      (Some first_word, remainder)
    else (None, title_no_tags)
  in
  (* Extract priority *)
  let priority, rest2 = extract_priority rest in
  (* Extract COMMENT marker *)
  let commented, rest3 =
    let first_word, remainder = split_first_word rest2 in
    if first_word = "COMMENT" then (true, remainder) else (false, rest2)
  in
  (* Parse remaining text as inline content *)
  let title_inline =
    if trim rest3 = "" then [] else Parse.inline (trim rest3)
  in
  (todo, priority, commented, title_inline, tags)

(* ------------------------------------------------------------------ *)
(* Checkbox parsing                                                    *)
(* ------------------------------------------------------------------ *)

let parse_checkbox = function
  | Some "[ ]" -> Some Ast.Unchecked
  | Some "[X]" | Some "[x]" -> Some Ast.Checked
  | Some "[-]" -> Some Ast.Partial
  | _ -> None

(* ------------------------------------------------------------------ *)
(* Descriptive list tag extraction                                     *)
(* ------------------------------------------------------------------ *)

(** Extract a descriptive list tag from the first body line. A descriptive list
    item has "TAG :: rest" in its text. *)
let extract_list_tag text =
  match String.index_opt text ':' with
  | None -> (None, text)
  | Some i ->
      if i + 1 < String.length text && text.[i + 1] = ':' then
        let tag = trim (String.sub text 0 i) in
        let rest_start = i + 2 in
        let rest =
          if rest_start < String.length text then
            trim (String.sub text rest_start (String.length text - rest_start))
          else ""
        in
        (Some tag, rest)
      else (None, text)

(* ------------------------------------------------------------------ *)
(* List nesting                                                        *)
(* ------------------------------------------------------------------ *)

(** Determine list kind from the first item. *)
let determine_list_kind (items : Raw_ast.raw_list_item list) =
  match items with
  | [] -> Ast.Unordered
  | first :: _ ->
      let bullet = trim first.Raw_ast.rli_bullet in
      if first.Raw_ast.rli_tag <> None then Ast.Descriptive
      else begin
        (* Check if bullet indicates an ordered list *)
        let len = String.length bullet in
        if len >= 2 then
          let last = bullet.[len - 1] in
          if last = '.' || last = ')' then
            (* Check if prefix is digits or alpha *)
            let prefix = String.sub bullet 0 (len - 1) in
            let is_ordered =
              String.length prefix > 0
              &&
              let c = prefix.[0] in
              (c >= '0' && c <= '9')
              || (c >= 'a' && c <= 'z')
              || (c >= 'A' && c <= 'Z')
            in
            if is_ordered then Ast.Ordered else Ast.Unordered
          else Ast.Unordered
        else Ast.Unordered
      end

(** Check if a raw list item text contains a descriptive tag (::). *)
let item_has_tag (item : Raw_ast.raw_list_item) =
  item.Raw_ast.rli_tag <> None
  ||
  match item.Raw_ast.rli_body_lines with
  | first :: _ -> (
      match String.index_opt first ':' with
      | Some i -> i + 1 < String.length first && first.[i + 1] = ':'
      | None -> false)
  | [] -> false

(** Determine list kind considering descriptive tag detection. *)
let determine_list_kind_with_tags (items : Raw_ast.raw_list_item list) =
  match items with
  | [] -> Ast.Unordered
  | first :: _ ->
      if item_has_tag first then Ast.Descriptive else determine_list_kind items

(** Nest flat list items by indentation into a tree. Items with greater
    indentation than the first item at this level become sub-items of the
    preceding item. *)
let rec nest_list_items (items : Raw_ast.raw_list_item list) :
    Ast.list_item list =
  match items with
  | [] -> []
  | first :: _ ->
      let base_indent = first.Raw_ast.rli_indent in
      let rec group acc current_item remaining =
        match remaining with
        | [] ->
            let finished = finish_item current_item [] in
            List.rev (finished :: acc)
        | next :: _ when next.Raw_ast.rli_indent > base_indent -> (
            (* Collect all deeper items as sub-items of current_item *)
            let sub_items, remaining' =
              collect_sub_items base_indent remaining
            in
            let finished = finish_item current_item sub_items in
            match remaining' with
            | [] -> List.rev (finished :: acc)
            | next_top :: rest_top -> group (finished :: acc) next_top rest_top)
        | next :: rest ->
            (* Same level: finish current, start new *)
            let finished = finish_item current_item [] in
            group (finished :: acc) next rest
      in
      group [] first (List.tl items)

and collect_sub_items base_indent items =
  let rec collect acc = function
    | [] -> (List.rev acc, [])
    | item :: rest ->
        if item.Raw_ast.rli_indent > base_indent then collect (item :: acc) rest
        else (List.rev acc, item :: rest)
  in
  collect [] items

and finish_item (raw : Raw_ast.raw_list_item)
    (sub_items : Raw_ast.raw_list_item list) : Ast.list_item =
  let checkbox = parse_checkbox raw.Raw_ast.rli_checkbox in
  (* Join body lines *)
  let body_text = String.concat " " raw.Raw_ast.rli_body_lines in
  (* Check for descriptive tag *)
  let tag_opt, content_text =
    match raw.Raw_ast.rli_tag with
    | Some t -> (Some t, body_text)
    | None -> extract_list_tag body_text
  in
  let tag_inline = Option.map (fun t -> Parse.inline t) tag_opt in
  (* Build contents: paragraph from text + sub-list if any *)
  let contents =
    let text_elements =
      if trim content_text = "" then []
      else [ Ast.Paragraph (Parse.inline content_text, []) ]
    in
    let sub_list_elements =
      if sub_items = [] then []
      else
        let kind = determine_list_kind_with_tags sub_items in
        [ Ast.Plain_list (kind, nest_list_items sub_items) ]
    in
    text_elements @ sub_list_elements
  in
  {
    Ast.li_bullet = raw.Raw_ast.rli_bullet;
    li_counter_set = raw.Raw_ast.rli_counter_set;
    li_checkbox = checkbox;
    li_tag = tag_inline;
    li_contents = contents;
  }

(* ------------------------------------------------------------------ *)
(* Table normalization                                                  *)
(* ------------------------------------------------------------------ *)

(** Parse a table row string into cells. Row format: "| cell1 | cell2 | cell3 |"
*)
let parse_table_row_cells row_str =
  (* Strip leading and trailing | *)
  let s = trim row_str in
  let s =
    if String.length s > 0 && s.[0] = '|' then
      String.sub s 1 (String.length s - 1)
    else s
  in
  let s =
    if String.length s > 0 && s.[String.length s - 1] = '|' then
      String.sub s 0 (String.length s - 1)
    else s
  in
  let parts = String.split_on_char '|' s in
  List.map (fun cell -> Parse.inline (trim cell)) parts

let normalize_table_row (row : Raw_ast.raw_table_row) : Ast.table_row =
  match row with
  | Raw_ast.Raw_table_rule -> Ast.Table_row_rule
  | Raw_ast.Raw_table_standard s ->
      Ast.Table_row_standard (parse_table_row_cells s)

let normalize_table (rows : Raw_ast.raw_table_row list) : Ast.table =
  { Ast.rows = List.map normalize_table_row rows }

(* ------------------------------------------------------------------ *)
(* Block classification                                                *)
(* ------------------------------------------------------------------ *)

(** Parse src block params: first word is language, rest is switches/arguments.
*)
let parse_src_params params_opt =
  match params_opt with
  | None -> (None, None, None)
  | Some params ->
      let params = trim params in
      if params = "" then (None, None, None)
      else
        let lang, rest = split_first_word params in
        let language = if lang = "" then None else Some lang in
        let arguments = if trim rest = "" then None else Some (trim rest) in
        (language, None, arguments)

(** Normalize a block body that contains recursive elements. *)
let rec normalize_recursive_body (elements : Raw_ast.raw_element list) config :
    Ast.element list =
  normalize_elements config elements

(** Classify a raw block into the appropriate Ast.block variant. *)
and classify_block (raw : Raw_ast.raw_block) config
    (affiliated : Ast.affiliated list) : Ast.block =
  let name = string_lowercase raw.rb_name in
  match name with
  | "src" ->
      let body_str =
        match raw.rb_body with
        | Raw_ast.Opaque_body s -> s
        | Raw_ast.Recursive_body _ -> ""
      in
      let language, switches, arguments = parse_src_params raw.rb_params in
      Ast.Src_block
        {
          src_language = language;
          src_switches = switches;
          src_arguments = arguments;
          src_value = body_str;
          src_affiliated = affiliated;
        }
  | "example" ->
      let body_str =
        match raw.rb_body with
        | Raw_ast.Opaque_body s -> s
        | Raw_ast.Recursive_body _ -> ""
      in
      Ast.Example_block
        {
          ex_value = body_str;
          ex_switches = raw.rb_params;
          ex_affiliated = affiliated;
        }
  | "export" ->
      let body_str =
        match raw.rb_body with
        | Raw_ast.Opaque_body s -> s
        | Raw_ast.Recursive_body _ -> ""
      in
      let backend =
        match raw.rb_params with
        | None -> ""
        | Some p -> fst (split_first_word p)
      in
      Ast.Export_block
        {
          exp_backend = backend;
          exp_value = body_str;
          exp_affiliated = affiliated;
        }
  | "comment" ->
      let body_str =
        match raw.rb_body with
        | Raw_ast.Opaque_body s -> s
        | Raw_ast.Recursive_body _ -> ""
      in
      Ast.Comment_block { cb_value = body_str; cb_affiliated = affiliated }
  | "quote" ->
      let contents =
        match raw.rb_body with
        | Raw_ast.Recursive_body elems -> normalize_recursive_body elems config
        | Raw_ast.Opaque_body s ->
            if trim s = "" then [] else [ Ast.Paragraph (Parse.inline s, []) ]
      in
      Ast.Quote_block { qt_contents = contents; qt_affiliated = affiliated }
  | "center" ->
      let contents =
        match raw.rb_body with
        | Raw_ast.Recursive_body elems -> normalize_recursive_body elems config
        | Raw_ast.Opaque_body s ->
            if trim s = "" then [] else [ Ast.Paragraph (Parse.inline s, []) ]
      in
      Ast.Center_block { cn_contents = contents; cn_affiliated = affiliated }
  | "verse" ->
      let inline =
        match raw.rb_body with
        | Raw_ast.Opaque_body s -> Parse.inline s
        | Raw_ast.Recursive_body _ -> []
      in
      Ast.Verse_block { vs_contents = inline; vs_affiliated = affiliated }
  | _ ->
      let contents =
        match raw.rb_body with
        | Raw_ast.Recursive_body elems -> normalize_recursive_body elems config
        | Raw_ast.Opaque_body s ->
            if trim s = "" then [] else [ Ast.Paragraph (Parse.inline s, []) ]
      in
      Ast.Custom_block
        {
          cst_name = raw.rb_name;
          cst_params = raw.rb_params;
          cst_contents = contents;
          cst_affiliated = affiliated;
        }

(* ------------------------------------------------------------------ *)
(* Element normalization with affiliated keyword attachment            *)
(* ------------------------------------------------------------------ *)

(** Normalize a list of raw elements into Ast elements. Handles affiliated
    keyword collection and attachment. *)
and normalize_elements config (raw_elements : Raw_ast.raw_element list) :
    Ast.element list =
  let rec process acc affiliated = function
    | [] ->
        (* Any trailing affiliated keywords become plain keywords *)
        let trailing =
          List.rev_map
            (fun (aff : Ast.affiliated) ->
              Ast.Keyword { kw_key = aff.aff_name; kw_value = aff.aff_value })
            affiliated
        in
        List.rev (trailing @ acc)
    | Raw_ast.Raw_paragraph [] :: rest ->
        (* Blank paragraph: skip it but flush any pending affiliated as keywords *)
        if affiliated <> [] then begin
          let kws =
            List.rev_map
              (fun (aff : Ast.affiliated) ->
                Ast.Keyword { kw_key = aff.aff_name; kw_value = aff.aff_value })
              affiliated
          in
          process (kws @ acc) [] rest
        end
        else process acc [] rest
    | Raw_ast.Raw_affiliated (name, opt, value) :: rest ->
        let aff =
          { Ast.aff_name = name; aff_optional = opt; aff_value = value }
        in
        process acc (aff :: affiliated) rest
    | elem :: rest -> (
        let normalized =
          normalize_single_element config elem (List.rev affiliated)
        in
        match normalized with
        | Some e -> process (e :: acc) [] rest
        | None -> process acc [] rest)
  in
  process [] [] raw_elements

(** Normalize a single raw element into an Ast element. *)
and normalize_single_element config (elem : Raw_ast.raw_element)
    (affiliated : Ast.affiliated list) : Ast.element option =
  match elem with
  | Raw_ast.Raw_paragraph [] -> None
  | Raw_ast.Raw_paragraph lines ->
      let text = String.concat " " lines in
      let inline = Parse.inline text in
      Some (Ast.Paragraph (inline, affiliated))
  | Raw_ast.Raw_list_items items ->
      let kind = determine_list_kind_with_tags items in
      let nested = nest_list_items items in
      Some (Ast.Plain_list (kind, nested))
  | Raw_ast.Raw_table rows -> Some (Ast.Table (normalize_table rows))
  | Raw_ast.Raw_block raw_block ->
      Some (Ast.Block (classify_block raw_block config affiliated))
  | Raw_ast.Raw_drawer drawer ->
      let contents = normalize_elements config drawer.rd_contents in
      Some (Ast.Drawer (drawer.rd_name, contents))
  | Raw_ast.Raw_keyword (key, value) ->
      Some (Ast.Keyword { kw_key = key; kw_value = value })
  | Raw_ast.Raw_affiliated (name, _opt, value) ->
      (* Standalone affiliated keyword not attached to anything → keyword *)
      Some (Ast.Keyword { kw_key = name; kw_value = value })
  | Raw_ast.Raw_comment lines -> Some (Ast.Comment lines)
  | Raw_ast.Raw_fixed_width lines -> Some (Ast.Fixed_width lines)
  | Raw_ast.Raw_horizontal_rule -> Some Ast.Horizontal_rule
  | Raw_ast.Raw_planning _ ->
      (* Planning outside heading context becomes a paragraph *)
      None
  | Raw_ast.Raw_property_drawer _ ->
      (* Property drawer outside heading context: ignore *)
      None
  | Raw_ast.Raw_heading _ ->
      (* Should not appear in element lists *)
      None

(* ------------------------------------------------------------------ *)
(* Planning and property extraction from heading body                  *)
(* ------------------------------------------------------------------ *)

(** Extract a timestamp from a raw planning entry. Uses the inline parser which
    handles timestamps. *)
let extract_timestamp ts_str =
  let inlines = Parse.inline ts_str in
  let rec find_ts = function
    | [] -> None
    | Ast.Timestamp ts :: _ -> Some ts
    | _ :: rest -> find_ts rest
  in
  find_ts inlines

(** Extract planning information from a list of raw_planning entries. *)
let extract_planning (entries : Raw_ast.raw_planning list) : Ast.planning option
    =
  match entries with
  | [] -> None
  | _ ->
      let planning =
        List.fold_left
          (fun acc (entry : Raw_ast.raw_planning) ->
            match entry.rpl_keyword with
            | "DEADLINE" ->
                {
                  acc with
                  Ast.deadline = extract_timestamp entry.rpl_timestamp;
                }
            | "SCHEDULED" ->
                {
                  acc with
                  Ast.scheduled = extract_timestamp entry.rpl_timestamp;
                }
            | "CLOSED" ->
                { acc with Ast.closed = extract_timestamp entry.rpl_timestamp }
            | _ -> acc)
          { Ast.deadline = None; scheduled = None; closed = None }
          entries
      in
      if
        planning.deadline = None && planning.scheduled = None
        && planning.closed = None
      then None
      else Some planning

(** Extract properties from drawer contents. In a PROPERTIES drawer, each entry
    was parsed as Raw_keyword(name, value). *)
let extract_properties (elements : Raw_ast.raw_element list) : Ast.property list
    =
  List.filter_map
    (fun elem ->
      match elem with
      | Raw_ast.Raw_keyword (name, value) ->
          let prop_value = if trim value = "" then None else Some value in
          Some { Ast.prop_name = name; prop_value }
      | Raw_ast.Raw_paragraph [] -> None
      | _ -> None)
    elements

(** Process heading body to extract planning, properties, and remaining
    elements. *)
let process_heading_body config (body : Raw_ast.raw_element list) =
  let planning, rest1 =
    match body with
    | Raw_ast.Raw_planning entries :: rest -> (extract_planning entries, rest)
    | _ -> (None, body)
  in
  let properties, rest2 =
    match rest1 with
    | Raw_ast.Raw_drawer { rd_name; rd_contents } :: rest
      when string_lowercase rd_name = "properties" ->
        (extract_properties rd_contents, rest)
    | _ -> ([], rest1)
  in
  let elements = normalize_elements config rest2 in
  (planning, properties, elements)

(* ------------------------------------------------------------------ *)
(* Heading hierarchy                                                   *)
(* ------------------------------------------------------------------ *)

(** Build heading hierarchy from a flat list of raw headings. Each heading's
    children are subsequent headings with a greater level, until a heading of
    equal or lesser level is encountered. *)
let build_heading_tree config (raw_headings : Raw_ast.raw_heading list) :
    Ast.heading list =
  let rec process_headings headings =
    match headings with
    | [] -> []
    | raw :: rest ->
        let children_raw, siblings_raw =
          collect_children raw.Raw_ast.rh_level rest
        in
        let todo, priority, commented, title, tags =
          parse_heading_title config raw.Raw_ast.rh_raw_title
        in
        let planning, properties, contents =
          process_heading_body config raw.Raw_ast.rh_body
        in
        let children = process_headings children_raw in
        let heading =
          {
            Ast.level = raw.Raw_ast.rh_level;
            todo;
            priority;
            commented;
            title;
            tags;
            planning;
            properties;
            contents;
            children;
          }
        in
        heading :: process_headings siblings_raw
  and collect_children parent_level headings =
    (* Collect all headings that are deeper than parent_level *)
    let rec collect children remaining =
      match remaining with
      | [] -> (List.rev children, [])
      | h :: _ when h.Raw_ast.rh_level <= parent_level ->
          (List.rev children, remaining)
      | h :: rest -> collect (h :: children) rest
    in
    collect [] headings
  in
  process_headings raw_headings

(* ------------------------------------------------------------------ *)
(* Directive extraction                                                *)
(* ------------------------------------------------------------------ *)

(** Known directive keywords that should be extracted from the preamble. *)
let directive_keywords =
  [
    "TITLE";
    "AUTHOR";
    "DATE";
    "EMAIL";
    "LANGUAGE";
    "DESCRIPTION";
    "OPTIONS";
    "STARTUP";
    "FILETAGS";
    "CATEGORY";
    "PROPERTY";
    "ARCHIVE";
    "COLUMNS";
    "LINK";
    "PRIORITIES";
    "TAGS";
    "EXPORT_FILE_NAME";
  ]

(** Check if a keyword should be treated as a directive. *)
let is_directive_keyword key =
  List.mem (String.uppercase_ascii key) directive_keywords

(** Extract directives from the leading keywords of the preamble. Directives are
    keywords that appear before any non-keyword/non-blank element in the
    preamble. *)
let extract_directives (raw_elements : Raw_ast.raw_element list) :
    Ast.directive list * Raw_ast.raw_element list =
  let rec collect_leading directives remaining =
    match remaining with
    | Raw_ast.Raw_keyword (key, value) :: rest when is_directive_keyword key ->
        let dir = { Ast.dir_key = key; dir_value = value } in
        collect_leading (dir :: directives) rest
    | Raw_ast.Raw_paragraph [] :: rest ->
        (* Skip blank paragraphs while collecting directives *)
        collect_leading directives rest
    | Raw_ast.Raw_affiliated _ :: rest ->
        (* Skip affiliated keywords in directive zone *)
        collect_leading directives rest
    | _ -> (List.rev directives, remaining)
  in
  collect_leading [] raw_elements

(* ------------------------------------------------------------------ *)
(* Top-level normalization                                             *)
(* ------------------------------------------------------------------ *)

let normalize (config : Config.t) (raw : Raw_ast.document) : Ast.document =
  (* Extract directives from preamble *)
  let directives, remaining_preamble = extract_directives raw.rd_preamble in
  (* Normalize remaining preamble elements *)
  let preamble = normalize_elements config remaining_preamble in
  (* Build heading hierarchy *)
  let headings = build_heading_tree config raw.rd_headings in
  { Ast.directives; preamble; headings }