module Y = Hl_yaml.Unix
module StringMap = Map.Make (String)

(* Type for frontmatter data *)
type t = Yojson.Safe.t StringMap.t

(* Extract YAML frontmatter from markdown content *)
let extract_frontmatter (content : string) : string * string =
  let content = String.trim content in
  let len = String.length content in

  (* Check if content starts with --- *)
  if len < 6 || not (String.starts_with ~prefix:"---" content) then ("", content)
  else
    (* Find the end of the first line (after opening ---) *)
    match String.index_from_opt content 3 '\n' with
    | None -> ("", content) (* No newline after opening --- *)
    | Some first_newline -> (
        let yaml_start = first_newline + 1 in
        if yaml_start >= len then ("", content)
        else
          let yaml_section = String.sub content yaml_start (len - yaml_start) in
          (* Find the closing --- at the start of a line *)
          let closing_pos_opt =
            try
              Some
                (Str.search_forward (Str.regexp "^---[ \t]*$") yaml_section 0)
            with Not_found -> None
          in
          match closing_pos_opt with
          | None -> ("", content)
          | Some closing_pos ->
              let yaml_content = String.sub yaml_section 0 closing_pos in

              (* Find content after the closing --- *)
              let after_closing = yaml_start + closing_pos in
              let remaining_content =
                match String.index_from_opt content after_closing '\n' with
                | None -> ""
                | Some next_line ->
                    let content_start = next_line + 1 in
                    if content_start >= len then ""
                    else String.sub content content_start (len - content_start)
              in

              (String.trim yaml_content, String.trim remaining_content))

(* Convert Yojson object to StringMap *)
let yojson_to_string_map = function
  | `Assoc assoc ->
      Ok
        (List.fold_left
           (fun acc (key, value) -> StringMap.add key value acc)
           StringMap.empty assoc)
  | _ -> Error "Frontmatter must be a YAML object/map"

let parse_frontmatter (yaml_content : string) : (t, string) Result.t =
  let options = Y.make_options ~enable_imports:false () in
  match Y.parse ~options ~of_yojson:yojson_to_string_map yaml_content with
  | Ok fm -> Ok fm
  | Error error_list ->
      Error
        (List.map Hl_yaml.Spec.error_to_string error_list |> String.concat "\n")