(*
 * Copyright (c) 2016, 2017 Thomas Refis <trefis@janestreet.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *)

open Odoc_utils
open Types
module Comment = Odoc_model.Comment
open Odoc_model.Names

let default_lang_tag = "ocaml"

(** Resource collection for extension handlers.
    Resources are collected during document generation and retrieved when
    building the final page. *)
module Resources = struct
  let collected : Odoc_extension_registry.resource list ref = ref []

  let add resources =
    collected := !collected @ resources

  let take () =
    let result = !collected in
    collected := [];
    result

  let clear () =
    collected := []
end

(** Asset collection for extension handlers.
    Assets (binary files like PNGs) are collected during document generation
    and written alongside the HTML output. *)
module Assets = struct
  let collected : Odoc_extension_registry.asset list ref = ref []

  let add assets =
    collected := !collected @ assets

  let take () =
    let result = !collected in
    collected := [];
    result

  let clear () =
    collected := []
end

let source_of_code s =
  if s = "" then [] else [ Source.Elt [ inline @@ Inline.Text s ] ]

module Reference = struct
  open Odoc_model.Paths

  let rec render_resolved : Reference.Resolved.t -> string =
   fun r ->
    let open Reference.Resolved in
    match r with
    | `Identifier id -> Identifier.name id
    | `Alias (_, r) -> render_resolved (r :> t)
    | `AliasModuleType (_, r) -> render_resolved (r :> t)
    | `Module (r, s) -> render_resolved (r :> t) ^ "." ^ ModuleName.to_string s
    | `Hidden p -> render_resolved (p :> t)
    | `ModuleType (r, s) ->
        render_resolved (r :> t) ^ "." ^ ModuleTypeName.to_string s
    | `Type (r, s) -> render_resolved (r :> t) ^ "." ^ TypeName.to_string s
    | `Constructor (r, s) ->
        render_resolved (r :> t) ^ "." ^ ConstructorName.to_string s
    | `PolyConstructor (r, s) ->
        render_resolved (r :> t) ^ ".`" ^ ConstructorName.to_string s
    | `Field (r, s) -> render_resolved (r :> t) ^ "." ^ FieldName.to_string s
    | `UnboxedField (r, s) -> render_resolved (r :> t) ^ "." ^ UnboxedFieldName.to_string s
    | `Extension (r, s) ->
        render_resolved (r :> t) ^ "." ^ ExtensionName.to_string s
    | `ExtensionDecl (r, _, s) ->
        render_resolved (r :> t) ^ "." ^ ExtensionName.to_string s
    | `Exception (r, s) ->
        render_resolved (r :> t) ^ "." ^ ExceptionName.to_string s
    | `Value (r, s) -> render_resolved (r :> t) ^ "." ^ ValueName.to_string s
    | `Class (r, s) -> render_resolved (r :> t) ^ "." ^ TypeName.to_string s
    | `ClassType (r, s) -> render_resolved (r :> t) ^ "." ^ TypeName.to_string s
    | `Method (r, s) ->
        (* CR trefis: do we really want to print anything more than [s] here?  *)
        render_resolved (r :> t) ^ "." ^ MethodName.to_string s
    | `InstanceVariable (r, s) ->
        (* CR trefis: the following makes no sense to me... *)
        render_resolved (r :> t) ^ "." ^ InstanceVariableName.to_string s
    | `Label (_, s) -> LabelName.to_string s

  let render_path (tag, cs) =
    let tag =
      match tag with
      | `TRelativePath -> "./"
      | `TAbsolutePath -> "/"
      | `TCurrentPackage -> "//"
    in
    tag ^ String.concat ~sep:"/" cs

  let rec render_unresolved : Reference.t -> string =
    let open Reference in
    function
    | `Resolved r -> render_resolved r
    | `Root (n, _) -> n
    | `Dot (p, f) -> render_unresolved (p :> t) ^ "." ^ f
    | `Page_path p -> render_path p
    | `Asset_path p -> render_path p
    | `Module_path p -> render_path p
    | `Any_path p -> render_path p
    | `Module (p, f) ->
        render_unresolved (p :> t) ^ "." ^ ModuleName.to_string f
    | `ModuleType (p, f) ->
        render_unresolved (p :> t) ^ "." ^ ModuleTypeName.to_string f
    | `Type (p, f) -> render_unresolved (p :> t) ^ "." ^ TypeName.to_string f
    | `Constructor (p, f) ->
        render_unresolved (p :> t) ^ "." ^ ConstructorName.to_string f
    | `Field (p, f) -> render_unresolved (p :> t) ^ "." ^ FieldName.to_string f
    | `UnboxedField (p, f) -> render_unresolved (p :> t) ^ "." ^ UnboxedFieldName.to_string f
    | `Extension (p, f) ->
        render_unresolved (p :> t) ^ "." ^ ExtensionName.to_string f
    | `ExtensionDecl (p, f) ->
        render_unresolved (p :> t) ^ "." ^ ExtensionName.to_string f
    | `Exception (p, f) ->
        render_unresolved (p :> t) ^ "." ^ ExceptionName.to_string f
    | `Value (p, f) -> render_unresolved (p :> t) ^ "." ^ ValueName.to_string f
    | `Class (p, f) -> render_unresolved (p :> t) ^ "." ^ TypeName.to_string f
    | `ClassType (p, f) ->
        render_unresolved (p :> t) ^ "." ^ TypeName.to_string f
    | `Method (p, f) ->
        render_unresolved (p :> t) ^ "." ^ MethodName.to_string f
    | `InstanceVariable (p, f) ->
        render_unresolved (p :> t) ^ "." ^ InstanceVariableName.to_string f
    | `Label (p, f) -> render_unresolved (p :> t) ^ "." ^ LabelName.to_string f

  (* This is the entry point. *)
  let to_ir : ?text:Inline.t -> Reference.t -> Inline.t =
   fun ?text ref ->
    match ref with
    | `Resolved r ->
        (* IDENTIFIER MUST BE RENAMED TO DEFINITION. *)
        let id = Reference.Resolved.identifier r in
        let rendered = render_resolved r in
        let content =
          match text with
          | None -> [ inline @@ Inline.Source (source_of_code rendered) ]
          | Some s -> s
        and tooltip =
          (* Add a tooltip if the content is not the rendered reference. *)
          match text with
          | None -> None
          | Some _ -> Some rendered
        in
        let target =
          match id with
          | Some id ->
              let url = Url.from_identifier ~stop_before:false id in
              Target.Internal (Resolved url)
          | None -> Internal Unresolved
        in
        let link = { Link.target; content; tooltip } in
        [ inline @@ Inline.Link link ]
    | _ -> (
        let s = render_unresolved ref in
        match text with
        | None ->
            let s = source_of_code s in
            [ inline @@ Inline.Source s ]
        | Some content ->
            let link =
              { Link.target = Internal Unresolved; content; tooltip = Some s }
            in
            [ inline @@ Inline.Link link ])
end

let leaf_inline_element : Comment.leaf_inline_element -> Inline.one = function
  | `Space -> inline @@ Text " "
  | `Word s -> inline @@ Text s
  | `Code_span s -> inline @@ Source (source_of_code s)
  | `Math_span s -> inline @@ Math s
  | `Raw_markup (target, s) -> inline @@ Raw_markup (target, s)

let rec non_link_inline_element : Comment.non_link_inline_element -> Inline.one
    = function
  | #Comment.leaf_inline_element as e -> leaf_inline_element e
  | `Styled (style, content) ->
      inline @@ Styled (style, non_link_inline_element_list content)

and non_link_inline_element_list : _ -> Inline.t =
 fun elements ->
  List.map
    (fun elt -> non_link_inline_element elt.Odoc_model.Location_.value)
    elements

let link_content = non_link_inline_element_list

let rec inline_element : Comment.inline_element -> Inline.t = function
  | #Comment.leaf_inline_element as e -> [ leaf_inline_element e ]
  | `Styled (style, content) ->
      [ inline @@ Styled (style, inline_element_list content) ]
  | `Reference (path, content) ->
      (* TODO Rework that ugly function. *)
      (* TODO References should be set in code style, if they are to code
              elements. *)
      let content =
        match content with
        | [] -> None
        | _ -> Some (non_link_inline_element_list content)
        (* XXX Span *)
      in
      Reference.to_ir ?text:content path
  | `Link (target, content) ->
      let content =
        match content with
        | [] -> [ inline @@ Text target ]
        | _ -> non_link_inline_element_list content
      in
      [ inline @@ Link { target = External target; content; tooltip = None } ]

and inline_element_list elements =
  List.concat
  @@ List.map
       (fun elt -> inline_element elt.Odoc_model.Location_.value)
       elements

let module_references ms =
  let module_reference (m : Comment.module_reference) =
    let reference =
      Reference.to_ir (m.module_reference :> Odoc_model.Paths.Reference.t)
    and synopsis =
      match m.module_synopsis with
      | Some synopsis ->
          [
            block ~attr:[ "synopsis" ] @@ Inline (inline_element_list synopsis);
          ]
      | None -> []
    in
    { Description.attr = []; key = reference; definition = synopsis }
  in
  let items = List.map module_reference ms in
  block ~attr:[ "modules" ] @@ Description items

let rec nestable_block_element :
    Comment.nestable_block_element -> Block.one list =
 fun content ->
  match content with
  | `Paragraph p -> [ paragraph p ]
  | `Code_block c ->
      let lang_tag, other_tags =
        match c.meta with
        | Some { language = { Odoc_parser.Loc.value; _ }; tags } -> (value, tags)
        | None -> (default_lang_tag, [])
      in
      let prefix = Odoc_extension_registry.prefix_of_language lang_tag in
      (* Check for a registered code block handler *)
      let handler_result =
        match Odoc_extension_registry.find_code_block_handler ~prefix with
        | Some handler ->
            let meta = { Odoc_extension_registry.language = lang_tag; tags = other_tags } in
            handler meta (Odoc_model.Location_.value c.content)
        | None ->
            None
      in
      (match handler_result with
      | Some result ->
          (* Handler produced a result, collect resources/assets and use content *)
          Resources.add result.resources;
          Assets.add result.assets;
          result.content
      | None ->
          (* No handler or handler declined, use default rendering *)
          let rest =
            match c.output with
            | Some xs -> nestable_block_element_list xs
            | None -> []
          in
          let value : 'a Odoc_parser.Loc.with_location -> 'a = fun x -> x.value in
          let classes =
            List.filter_map
              (function `Binding (_, _) -> None | `Tag t -> Some (value t))
              other_tags
          in
          let data =
            List.filter_map
              (function
                | `Binding (k, v) -> Some (value k, value v) | `Tag _ -> None)
              other_tags
          in
          [
            block
            @@ Source
                 ( lang_tag,
                   classes,
                   data,
                   source_of_code (Odoc_model.Location_.value c.content),
                   rest );
          ]
          @ rest)
  | `Math_block s -> [ block @@ Math s ]
  | `Verbatim s -> [ block @@ Verbatim s ]
  | `Modules ms -> [ module_references ms ]
  | `List (kind, items) ->
      let kind =
        match kind with
        | `Unordered -> Block.Unordered
        | `Ordered -> Block.Ordered
      in
      let f = function
        | [ { Odoc_model.Location_.value = `Paragraph content; _ } ] ->
            [ block @@ Block.Inline (inline_element_list content) ]
        | item -> nestable_block_element_list item
      in
      let items = List.map f items in
      [ block @@ Block.List (kind, items) ]
  | `Table { data; align } ->
      let data =
        List.map
          (List.map (fun (cell, cell_type) ->
               (nestable_block_element_list cell, cell_type)))
          data
      in
      let generate_align data =
        let max (a : int) b = if a < b then b else a in
        (* Length of the longest line of the table *)
        let max_length =
          List.fold_left (fun m l -> max m (List.length l)) 0 data
        in
        let rec list_init i =
          if i <= 0 then [] else Table.Default :: list_init (i - 1)
        in
        list_init max_length
      in
      let align =
        match align with
        | None -> generate_align data
        | Some align ->
            List.map
              (function
                | None -> Table.Default
                | Some `Right -> Right
                | Some `Left -> Left
                | Some `Center -> Center)
              align
        (* We should also check wellness of number of table cells vs alignment,
           and raise warnings *)
      in
      [ block @@ Table { data; align } ]
  | `Media (href, media, content) ->
      let content =
        match (content, href) with
        | "", `Reference path ->
            Reference.render_unresolved (path :> Comment.Reference.t)
        | "", `Link href -> href
        | _ -> content
      in
      let url =
        match href with
        | `Reference (`Resolved r) -> (
            let id =
              Odoc_model.Paths.Reference.Resolved.Asset.(identifier (r :> t))
            in
            match Url.from_asset_identifier id with
            | url -> Target.Internal (Resolved url))
        | `Reference _ -> Internal Unresolved
        | `Link href -> External href
      in
      let i =
        match media with
        | `Audio -> Block.Audio (url, content)
        | `Video -> Video (url, content)
        | `Image -> Image (url, content)
      in
      [ block i ]

and paragraph : Comment.paragraph -> Block.one = function
  | [ { value = `Raw_markup (target, s); _ } ] ->
      block @@ Block.Raw_markup (target, s)
  | p -> block @@ Block.Paragraph (inline_element_list p)

and nestable_block_element_list :
    Comment.nestable_block_element Comment.with_location list -> Block.one list
    =
 fun elements ->
  elements
  |> List.map Odoc_model.Location_.value
  |> List.map nestable_block_element
  |> List.concat

let tag : Comment.tag -> Description.one =
 fun t ->
  let sp = inline (Text " ") in
  let item ?value ~tag definition =
    let tag_name = inline ~attr:[ "at-tag" ] (Text tag) in
    let tag_value = match value with None -> [] | Some t -> sp :: t in
    let key = tag_name :: tag_value in
    { Description.attr = [ tag ]; key; definition }
  in
  let mk_value desc = [ inline ~attr:[ "value" ] desc ] in
  let text_def s = [ block (Block.Inline [ inline @@ Text s ]) ] in
  let content_to_inline ?(prefix = []) content =
    match content with
    | None -> []
    | Some content -> prefix @ [ inline @@ Text content ]
  in
  match t with
  | `Author s -> item ~tag:"author" (text_def s)
  | `Deprecated content ->
      item ~tag:"deprecated" (nestable_block_element_list content)
  | `Param (name, content) ->
      let value = mk_value (Inline.Text name) in
      item ~tag:"parameter" ~value (nestable_block_element_list content)
  | `Raise (kind, content) ->
      let value = inline_element (kind :> Comment.inline_element) in
      item ~tag:"raises" ~value (nestable_block_element_list content)
  | `Return content -> item ~tag:"returns" (nestable_block_element_list content)
  | `See (kind, target, content) ->
      let value =
        match kind with
        | `Url ->
            mk_value
              (Inline.Link
                 {
                   target = External target;
                   content = [ inline @@ Text target ];
                   tooltip = None;
                 })
        | `File -> mk_value (Inline.Source (source_of_code target))
        | `Document -> mk_value (Inline.Text target)
      in
      item ~tag:"see" ~value (nestable_block_element_list content)
  | `Since s -> item ~tag:"since" (text_def s)
  | `Before (version, content) ->
      let value = mk_value (Inline.Text version) in
      item ~tag:"before" ~value (nestable_block_element_list content)
  | `Version s -> item ~tag:"version" (text_def s)
  | `Alert ("deprecated", content) ->
      let content = content_to_inline content in
      item ~tag:"deprecated" [ block (Block.Inline content) ]
  | `Alert (tag, content) ->
      let content = content_to_inline ~prefix:[ sp ] content in
      item ~tag:"alert"
        [ block (Block.Inline ([ inline @@ Text tag ] @ content)) ]
  | `Custom (name, content) ->
      (* Check if there's a registered extension for this tag *)
      let prefix = Odoc_extension_registry.prefix_of_tag name in
      (match Odoc_extension_registry.find_handler ~prefix with
      | Some handler ->
          (match handler name content with
          | Some result ->
              (* Extension handled the tag - collect resources/assets and use output *)
              Resources.add result.Odoc_extension_registry.resources;
              Assets.add result.Odoc_extension_registry.assets;
              { Description.attr = [ name ];
                key = [];
                definition = result.Odoc_extension_registry.content }
          | None ->
              (* Extension declined to handle this tag variant *)
              item ~tag:name (nestable_block_element_list content))
      | None ->
          (* No extension registered - use default handling *)
          item ~tag:name (nestable_block_element_list content))

let attached_block_element : Comment.attached_block_element -> Block.t =
  function
  | #Comment.nestable_block_element as e -> nestable_block_element e
  | `Tag t ->
      let t = tag t in
      if t.Description.key = [] && t.Description.definition = [] then
        (* Extension tag with no visible output (e.g. config-only tags
           that only inject resources). Emit nothing. *)
        []
      else
        [ block ~attr:[ "at-tags" ] @@ Description [ t ] ]

(* TODO collaesce tags *)

let block_element : Comment.block_element -> Block.t = function
  | #Comment.attached_block_element as e -> attached_block_element e
  | `Heading (_, _, text) ->
      (* We are not supposed to receive Heading in this context.
         TODO: Remove heading in attached documentation in the model *)
      [ block @@ Paragraph (inline_element_list text) ]

let heading_level_to_int = function
  | `Title -> 0
  | `Section -> 1
  | `Subsection -> 2
  | `Subsubsection -> 3
  | `Paragraph -> 4
  | `Subparagraph -> 5

let heading
    (attrs, { Odoc_model.Paths.Identifier.iv = `Label (_, label); _ }, text) =
  let label = Odoc_model.Names.LabelName.to_string label in
  let title = inline_element_list text in
  let level = heading_level_to_int attrs.Comment.heading_level in
  let label = Some label in
  let source_anchor = None in
  Item.Heading { label; level; title; source_anchor }

let item_element : Comment.block_element -> Item.t list = function
  | #Comment.attached_block_element as e ->
      [ Item.Text (attached_block_element e) ]
  | `Heading h -> [ heading h ]

(** The documentation of the expansion is used if there is no comment attached
    to the declaration. *)
let synopsis ~decl_doc ~expansion_doc =
  let ([], Some docs | docs, _) = (decl_doc, expansion_doc) in
  match Comment.synopsis docs with Some p -> [ paragraph p ] | None -> []

let standalone docs =
  List.concat_map item_element
  @@ List.map (fun x -> x.Odoc_model.Location_.value) docs

let to_ir (docs : Comment.elements) =
  List.concat_map block_element
  @@ List.map (fun x -> x.Odoc_model.Location_.value) docs

let has_doc docs = docs <> []