open Odoc_utils

module Location = Location_
module Ast = Odoc_parser.Ast

type internal_tags_removed =
  [ `Tag of Ast.ocamldoc_tag
  | `Heading of Ast.heading
  | `Media of
    Ast.reference_kind * Ast.media_href Ast.with_location * string * Ast.media
  | Ast.nestable_block_element ]
(** {!Ast.block_element} without internal tags. *)

type _ handle_internal_tags =
  | Expect_status :
      [ `Default | `Inline | `Open | `Closed ] handle_internal_tags
  | Expect_canonical : Reference.path option handle_internal_tags
  | Expect_none : unit handle_internal_tags
  | Expect_page_tags : Frontmatter.t handle_internal_tags

let describe_internal_tag = function
  | `Canonical _ -> "@canonical"
  | `Inline -> "@inline"
  | `Open -> "@open"
  | `Closed -> "@closed"
  | `Hidden -> "@hidden"
  | `Children_order _ -> "@children_order"
  | `Toc_status _ -> "@toc_status"
  | `Short_title _ -> "@short_title"
  | `Order_category _ -> "@order_category"

let warn_unexpected_tag { Location.value; location } =
  Error.raise_warning
  @@ Error.make "Unexpected tag '%s' at this location."
       (describe_internal_tag value)
       location

let warn_root_canonical location =
  Error.raise_warning
  @@ Error.make "Canonical paths must contain a dot, eg. X.Y." location

let rec find_tag ~filter = function
  | [] -> None
  | hd :: tl -> (
      match filter hd.Location.value with
      | Some x -> Some (x, hd.location)
      | None ->
          warn_unexpected_tag hd;
          find_tag ~filter tl)

let rec find_tags acc ~filter = function
  | [] -> List.rev acc
  | hd :: tl -> (
      match filter hd.Location.value with
      | Some x -> find_tags ((x, hd.location) :: acc) ~filter tl
      | None ->
          warn_unexpected_tag hd;
          find_tags acc ~filter tl)

(* Errors *)
let invalid_raw_markup_target : string -> Location.span -> Error.t =
  Error.make ~suggestion:"try '{%html:...%}'."
    "'{%%%s:': bad raw markup target."

let default_raw_markup_target_not_supported : Location.span -> Error.t =
  Error.make ~suggestion:"try '{%html:...%}'."
    "'{%%...%%}' (raw markup) needs a target language."

let bad_heading_level : int -> Location.span -> Error.t =
  Error.make "'%d': bad heading level (0-5 allowed)."

let heading_level_should_be_lower_than_top_level :
    int -> int -> Location.span -> Error.t =
 fun this_heading_level top_heading_level ->
  Error.make "%s: heading level should be lower than top heading level '%d'."
    (Printf.sprintf "'{%i'" this_heading_level)
    top_heading_level

let page_heading_required : string -> Error.t =
  Error.filename_only "Pages (.mld files) should start with a heading."

let tags_not_allowed : Location.span -> Error.t =
  Error.make "Tags are not allowed in pages."

let not_allowed :
    ?suggestion:string ->
    what:string ->
    in_what:string ->
    Location.span ->
    Error.t =
 fun ?suggestion ~what ~in_what ->
  Error.make ?suggestion "%s is not allowed in %s."
    (Astring.String.Ascii.capitalize what)
    in_what

let describe_element = function
  | `Reference (`Simple, _, _) -> "'{!...}' (cross-reference)"
  | `Reference (`With_text, _, _) -> "'{{!...} ...}' (cross-reference)"
  | `Link (_, _) -> "'{{:...} ...}' (external link)"
  | `Heading (level, _, _) ->
      Printf.sprintf "'{%i ...}' (section heading)" level
  | `Specific s -> s

(* End of errors *)

type 'a with_location = 'a Location.with_location

type ast_leaf_inline_element =
  [ `Space of string
  | `Word of string
  | `Code_span of string
  | `Math_span of string
  | `Raw_markup of string option * string ]

type sections_allowed = [ `All | `No_titles | `None ]

type alerts =
  [ `Tag of [ `Alert of string * string option ] ] Location_.with_location list

type status = {
  tags_allowed : bool;
  parent_of_sections : Paths.Identifier.LabelParent.t;
}

let leaf_inline_element :
    ast_leaf_inline_element with_location ->
    Comment.leaf_inline_element with_location =
 fun element ->
  match element with
  | { value = `Word _ | `Code_span _ | `Math_span _; _ } as element -> element
  | { value = `Space _; _ } -> Location.same element `Space
  | { value = `Raw_markup (target, s); location } -> (
      match target with
      | Some invalid_target
        when String.trim invalid_target = ""
             || String.exists
                  (function '%' | '}' -> true | _ -> false)
                  invalid_target ->
          Error.raise_warning
            (invalid_raw_markup_target invalid_target location);

          Location.same element (`Code_span s)
      | None ->
          Error.raise_warning (default_raw_markup_target_not_supported location);
          Location.same element (`Code_span s)
      | Some target -> Location.same element (`Raw_markup (target, s)))

type surrounding =
  [ `Link of
    string * Odoc_parser.Ast.inline_element Location_.with_location list
  | `Reference of
    [ `Simple | `With_text ]
    * string Location_.with_location
    * Odoc_parser.Ast.inline_element Location_.with_location list
  | `Specific of string ]

let rec non_link_inline_element :
    surrounding:surrounding ->
    Odoc_parser.Ast.inline_element with_location ->
    Comment.non_link_inline_element with_location =
 fun ~surrounding element ->
  match element with
  | { value = #ast_leaf_inline_element; _ } as element ->
      (leaf_inline_element element
        :> Comment.non_link_inline_element with_location)
  | { value = `Styled (style, content); _ } ->
      `Styled (style, non_link_inline_elements ~surrounding content)
      |> Location.same element
  | ( { value = `Reference (_, _, content); _ }
    | { value = `Link (_, content); _ } ) as element ->
      not_allowed
        ~what:(describe_element element.value)
        ~in_what:(describe_element surrounding)
        element.location
      |> Error.raise_warning;

      `Styled (`Emphasis, non_link_inline_elements ~surrounding content)
      |> Location.same element

and non_link_inline_elements ~surrounding elements =
  List.map (non_link_inline_element ~surrounding) elements

let rec inline_element :
    Odoc_parser.Ast.inline_element with_location ->
    Comment.inline_element with_location =
 fun element ->
  match element with
  | { value = #ast_leaf_inline_element; _ } as element ->
      (leaf_inline_element element :> Comment.inline_element with_location)
  | { value = `Styled (style, content); location } ->
      `Styled (style, inline_elements content) |> Location.at location
  | { value = `Reference (kind, target, content) as value; location } -> (
      let { Location.value = target; location = target_location } = target in
      match Error.raise_warnings (Reference.parse target_location target) with
      | Ok target ->
          let content = non_link_inline_elements ~surrounding:value content in
          Location.at location (`Reference (target, content))
      | Error error ->
          Error.raise_warning error;
          let placeholder =
            match kind with
            | `Simple -> `Code_span target
            | `With_text -> `Styled (`Emphasis, content)
          in
          inline_element (Location.at location placeholder))
  | { value = `Link (target, content) as value; location } ->
      `Link (target, non_link_inline_elements ~surrounding:value content)
      |> Location.at location

and inline_elements elements = List.map inline_element elements

let rec nestable_block_element :
    Odoc_parser.Ast.nestable_block_element with_location ->
    Comment.nestable_block_element with_location =
 fun element ->
  match element with
  | { value = `Paragraph content; location } ->
      Location.at location (`Paragraph (inline_elements content))
  | { value = `Code_block { meta; delimiter; content; output }; location } ->
      let output =
        match output with
        | None -> None
        | Some l -> Some (List.map nestable_block_element l)
      in
      let trimmed_content, warnings =
        Odoc_parser.codeblock_content location content.value
      in
      let warnings = List.map Error.t_of_parser_t warnings in
      List.iter (Error.raise_warning ~non_fatal:true) warnings;
      let content = Location.at content.location trimmed_content in
      let code_block = { Comment.meta; delimiter; content; output } in
      Location.at location (`Code_block code_block)
  | { value = `Math_block s; location } -> Location.at location (`Math_block s)
  | { value = `Verbatim v; location } ->
      let v, warnings = Odoc_parser.codeblock_content location v in
      let warnings = List.map Error.t_of_parser_t warnings in
      List.iter (Error.raise_warning ~non_fatal:true) warnings;
      Location.at location (`Verbatim v)
  | { value = `Modules modules; location } ->
      let modules =
        List.fold_left
          (fun acc { Location.value; location } ->
            match
              Error.raise_warnings (Reference.read_mod_longident location value)
            with
            | Ok r ->
                { Comment.module_reference = r; module_synopsis = None } :: acc
            | Error error ->
                Error.raise_warning error;
                acc)
          [] modules
        |> List.rev
      in
      Location.at location (`Modules modules)
  | { value = `List (kind, _syntax, items); location } ->
      `List (kind, List.map nestable_block_elements items)
      |> Location.at location
  | { value = `Table ((grid, align), (`Heavy | `Light)); location } ->
      let data =
        List.map
          (List.map (fun (cell, cell_type) ->
               (nestable_block_elements cell, cell_type)))
          grid
      in
      `Table { Comment.data; align } |> Location.at location
  | { value = `Media (_, { value = `Link href; _ }, content, m); location } ->
      `Media (`Link href, m, content) |> Location.at location
  | {
   value =
     `Media
       (kind, { value = `Reference href; location = href_location }, content, m);
   location;
  } -> (
      let fallback error =
        Error.raise_warning error;
        let placeholder =
          match kind with
          | `Simple -> `Code_span href
          | `With_text ->
              `Styled (`Emphasis, [ `Word content |> Location.at location ])
        in
        `Paragraph (inline_elements [ placeholder |> Location.at location ])
        |> Location.at location
      in
      match Error.raise_warnings (Reference.parse_asset href_location href) with
      | Ok target ->
          `Media (`Reference target, m, content) |> Location.at location
      | Error error -> fallback error)

and nestable_block_elements elements = List.map nestable_block_element elements

let tag :
    location:Location.span ->
    status ->
    Ast.ocamldoc_tag ->
    ( Comment.block_element with_location,
      internal_tags_removed with_location )
    result =
 fun ~location status tag ->
  if not status.tags_allowed then
    (* Trigger a warning but do not remove the tag. Avoid turning tags into
       text that would render the same. *)
    Error.raise_warning (tags_not_allowed location);
  let ok t = Ok (Location.at location (`Tag t)) in
  match tag with
  | (`Author _ | `Since _ | `Version _) as tag -> ok tag
  | `Custom (name, content) ->
      ok (`Custom (name, nestable_block_elements content))
  | `Deprecated content -> ok (`Deprecated (nestable_block_elements content))
  | `Param (name, content) ->
      ok (`Param (name, nestable_block_elements content))
  | `Raise (name, content) -> (
      match Error.raise_warnings (Reference.parse location name) with
      (* TODO: location for just name *)
      | Ok target ->
          ok (`Raise (`Reference (target, []), nestable_block_elements content))
      | Error error ->
          Error.raise_warning error;
          let placeholder = `Code_span name in
          ok (`Raise (placeholder, nestable_block_elements content)))
  | `Return content -> ok (`Return (nestable_block_elements content))
  | `See (kind, target, content) ->
      ok (`See (kind, target, nestable_block_elements content))
  | `Before (version, content) ->
      ok (`Before (version, nestable_block_elements content))

(* When the user does not give a section heading a label (anchor), we generate
   one from the text in the heading. This is the common case. This involves
   simply scanning the AST for words, lowercasing them, and joining them with
   hyphens.

   This must be done in the parser (i.e. early, not at HTML/other output
   generation time), so that the cross-referencer can see these anchors. *)
let generate_heading_label : Comment.inline_element with_location list -> string
    =
 fun content ->
  (* Code spans can contain spaces, so we need to replace them with hyphens. We
     also lowercase all the letters, for consistency with the rest of this
     procedure. *)
  let replace_spaces_with_hyphens_and_lowercase s =
    let result = Bytes.create (String.length s) in
    s
    |> String.iteri (fun index c ->
           let c =
             match c with
             | ' ' | '\t' | '\r' | '\n' -> '-'
             | _ -> Astring.Char.Ascii.lowercase c
           in
           Bytes.set result index c);
    Bytes.unsafe_to_string result
  in

  let strip_locs li = List.map (fun ele -> ele.Location.value) li in
  (* Perhaps this should be done using a [Buffer.t]; we can switch to that as
     needed. *)
  let rec scan_inline_elements anchor = function
    | [] -> anchor
    | element :: more ->
        let anchor =
          match (element : Comment.inline_element) with
          | `Space -> anchor ^ "-"
          | `Word w -> anchor ^ Astring.String.Ascii.lowercase w
          | `Code_span c | `Math_span c ->
              anchor ^ replace_spaces_with_hyphens_and_lowercase c
          | `Raw_markup _ ->
              (* TODO Perhaps having raw markup in a section heading should be an
                 error? *)
              anchor
          | `Styled (_, content) ->
              content |> strip_locs |> scan_inline_elements anchor
          | `Reference (_, content) | `Link (_, content) ->
              content |> strip_locs
              |> List.map (fun (ele : Comment.non_link_inline_element) ->
                     (ele :> Comment.inline_element))
              |> scan_inline_elements anchor
        in
        scan_inline_elements anchor more
  in
  content |> List.map (fun ele -> ele.Location.value) |> scan_inline_elements ""

let section_heading :
    status ->
    top_heading_level:int option ->
    Location.span ->
    [ `Heading of _ ] ->
    int option * Comment.block_element with_location =
 fun status ~top_heading_level location heading ->
  let (`Heading (level, label, content)) = heading in

  let text = inline_elements content in

  let heading_label_explicit, label =
    match label with
    | Some label -> (true, label)
    | None -> (false, generate_heading_label text)
  in
  let label =
    Paths.Identifier.Mk.label
      (status.parent_of_sections, Names.LabelName.make_std label)
  in

  let mk_heading heading_level =
    let attrs = { Comment.heading_level; heading_label_explicit } in
    let element = Location.at location (`Heading (attrs, label, text)) in
    let top_heading_level =
      match top_heading_level with None -> Some level | some -> some
    in
    (top_heading_level, element)
  in
  let level' =
    match level with
    | 0 -> `Title
    | 1 -> `Section
    | 2 -> `Subsection
    | 3 -> `Subsubsection
    | 4 -> `Paragraph
    | 5 -> `Subparagraph
    | _ ->
        Error.raise_warning (bad_heading_level level location);
        (* Implicitly promote to level-5. *)
        `Subparagraph
  in
  let () =
    match top_heading_level with
    | Some top_level when level <= top_level && level <= 5 ->
        Error.raise_warning
          (heading_level_should_be_lower_than_top_level level top_level location)
    | _ -> ()
  in
  mk_heading level'

let validate_first_page_heading status ast_element =
  match status.parent_of_sections.iv with
  | `Page (_, name) | `LeafPage (_, name) -> (
      match ast_element with
      | { Location.value = `Heading (_, _, _); _ } -> ()
      | _invalid_ast_element ->
          let filename = Names.PageName.to_string name ^ ".mld" in
          Error.raise_warning (page_heading_required filename))
  | _not_a_page -> ()

let top_level_block_elements status ast_elements =
  let rec traverse :
      top_heading_level:int option ->
      Comment.block_element with_location list ->
      internal_tags_removed with_location list ->
      Comment.block_element with_location list =
   fun ~top_heading_level comment_elements_acc ast_elements ->
    match ast_elements with
    | [] -> List.rev comment_elements_acc
    | ast_element :: ast_elements -> (
        (* The first [ast_element] in pages must be a title or section heading. *)
        if top_heading_level = None then
          validate_first_page_heading status ast_element;

        match ast_element with
        | { value = #Odoc_parser.Ast.nestable_block_element; _ } as element ->
            let element = nestable_block_element element in
            let element = (element :> Comment.block_element with_location) in
            traverse ~top_heading_level
              (element :: comment_elements_acc)
              ast_elements
        | { value = `Tag the_tag; location } -> (
            match tag ~location status the_tag with
            | Ok element ->
                traverse ~top_heading_level
                  (element :: comment_elements_acc)
                  ast_elements
            | Error placeholder ->
                traverse ~top_heading_level comment_elements_acc
                  (placeholder :: ast_elements))
        | { value = `Heading _ as heading; _ } ->
            let top_heading_level, element =
              section_heading status ~top_heading_level
                ast_element.Location.location heading
            in
            traverse ~top_heading_level
              (element :: comment_elements_acc)
              ast_elements)
  in
  let top_heading_level =
    (* Non-page documents have a generated title. *)
    match status.parent_of_sections.iv with
    | `Page _ | `LeafPage _ -> None
    | _parent_with_generated_title -> Some 0
  in
  traverse ~top_heading_level [] ast_elements

let strip_internal_tags ast : internal_tags_removed with_location list * _ =
  let rec loop ~start tags ast' = function
    | ({ Location.value = `Tag (#Ast.internal_tag as tag); _ } as wloc) :: tl
      -> (
        let next tag =
          loop ~start ({ wloc with value = tag } :: tags) ast' tl
        in
        match tag with
        | (`Inline | `Open | `Closed | `Hidden) as tag -> next tag
        | ( `Children_order _ | `Short_title _ | `Toc_status _
          | `Order_category _ ) as tag ->
            let tag_name = describe_internal_tag tag in
            if not start then
              Error.raise_warning
                (Error.make "%s tag has to be before any content" tag_name
                   wloc.location);
            next tag
        | `Canonical { Location.value = s; location = r_location } -> (
            match
              Error.raise_warnings (Reference.read_path_longident r_location s)
            with
            | Ok path -> next (`Canonical path)
            | Error e ->
                Error.raise_warning e;
                loop ~start tags ast' tl))
    | ({
         value =
           ( `Tag #Ast.ocamldoc_tag
           | `Heading _ | `Media _
           | #Ast.nestable_block_element );
         _;
       } as hd)
      :: tl ->
        loop ~start:false tags (hd :: ast') tl
    | [] -> (List.rev ast', List.rev tags)
  in
  loop ~start:true [] [] ast

(** Append alerts at the end of the comment. Tags are favoured in case of alerts
    of the same name. *)
let append_alerts_to_comment alerts
    (comment : Comment.block_element with_location list) =
  let alerts =
    List.filter
      (fun alert ->
        let (`Tag alert) = alert.Location_.value in
        List.for_all
          (fun elem ->
            match (elem.Location_.value, alert) with
            | `Tag (`Deprecated _), `Alert ("deprecated", _) -> false
            | _ -> true)
          comment)
      alerts
  in
  comment @ (alerts :> Comment.elements)

let handle_internal_tags (type a) tags : a handle_internal_tags -> a = function
  | Expect_status -> (
      match
        find_tag
          ~filter:(function
            | (`Inline | `Open | `Closed) as t -> Some t | _ -> None)
          tags
      with
      | Some (status, _) -> status
      | None -> `Default)
  | Expect_canonical -> (
      match
        find_tag ~filter:(function `Canonical p -> Some p | _ -> None) tags
      with
      | Some (`Root _, location) ->
          warn_root_canonical location;
          None
      | Some ((`Dot _ as p), _) -> Some p
      | None -> None)
  | Expect_page_tags ->
      let unparsed_lines =
        find_tags []
          ~filter:(function
            | ( `Children_order _ | `Toc_status _ | `Short_title _
              | `Order_category _ ) as p ->
                Some p
            | _ -> None)
          tags
      in
      let lines =
        let do_ parse loc els =
          let els = nestable_block_elements els in
          match parse loc els with
          | Ok res -> Some res
          | Error e ->
              Error.raise_warning e;
              None
        in
        List.filter_map
          (function
            | `Children_order co, loc ->
                do_ Frontmatter.parse_children_order loc co
            | `Toc_status co, loc -> do_ Frontmatter.parse_toc_status loc co
            | `Short_title t, loc -> do_ Frontmatter.parse_short_title loc t
            | `Order_category t, loc ->
                do_ Frontmatter.parse_order_category loc t)
          unparsed_lines
      in
      Frontmatter.of_lines lines |> Error.raise_warnings
  | Expect_none ->
      (* Will raise warnings. *)
      ignore (find_tag ~filter:(fun _ -> None) tags);
      ()

let ast_to_comment ~internal_tags ~tags_allowed ~parent_of_sections
    (ast : Ast.t) alerts =
  Error.catch_warnings (fun () ->
      let status = { tags_allowed; parent_of_sections } in
      let ast, tags = strip_internal_tags ast in
      let elts =
        top_level_block_elements status ast |> append_alerts_to_comment alerts
      in
      (elts, handle_internal_tags tags internal_tags))

let parse_comment ~internal_tags ~tags_allowed ~containing_definition ~location
    ~text =
  Error.catch_warnings (fun () ->
      let ast =
        Odoc_parser.parse_comment ~location ~text |> Error.raise_parser_warnings
      in
      ast_to_comment ~internal_tags ~tags_allowed
        ~parent_of_sections:containing_definition ast []
      |> Error.raise_warnings)

let parse_reference text =
  let location =
    Location_.
      {
        file = "";
        start = { line = 0; column = 0 };
        end_ = { line = 0; column = String.length text };
      }
  in
  Reference.parse location text

let non_link_inline_element :
    context:string ->
    Odoc_parser.Ast.inline_element with_location list ->
    Comment.non_link_inline_element with_location list =
 fun ~context elements ->
  let surrounding = `Specific context in
  non_link_inline_elements ~surrounding elements