module SM = Map.Make(String)

let heading_to_string heading =
  Cmarkit.Block.Heading.inline heading
  |> Cmarkit.Inline.to_plain_text ~break_on_soft:false
  |> List.map (String.concat "")
  |> String.concat "\n"

let aggregate t doc =
  let rec aggregate t heading = function
    | [] -> t
    | Cmarkit.Block.Heading (heading, _) :: rest ->
      let heading = heading_to_string heading in
      let t =
        SM.update heading
          (function
            | None -> Some []
            | Some blocks -> Some blocks)
          t
      in
      aggregate t heading rest
    | block :: rest ->
      let t =
        SM.update heading
          (function
            | None ->
              Some [ block ]
            | Some blocks ->
              Some (block :: blocks))
          t
      in
      aggregate t heading rest
  in
  match Cmarkit.Doc.block doc with
  | Cmarkit.Block.Blocks (blocks, _) ->
    aggregate t "" blocks
  | Cmarkit.Block.Heading (heading, _) ->
    let heading = heading_to_string heading in
    SM.update heading
      (function
        | None -> Some []
        | Some blocks -> Some blocks)
      t
  | block ->
    SM.update ""
      (function
        | None ->
          Some [ block ]
        | Some blocks ->
          Some (block :: blocks))
      t

let kortfat docs =
  let rev_t = List.fold_left aggregate SM.empty docs in
  let t = SM.map List.rev rev_t in
  let blocks =
    SM.fold
      (fun heading blocks acc ->
         let heading =
           Cmarkit.Block.Heading
             (Cmarkit.Block.Heading.make
                ~level:2
                (Cmarkit.Inline.Text
                   (heading, Cmarkit.Meta.none)),
              Cmarkit.Meta.none)
         in
         heading :: blocks @ acc)
      t
      []
  in
  Cmarkit.Block.Blocks (blocks, Cmarkit.Meta.none)
  |> Cmarkit.Doc.make