(* odoc-docsite: A docs.rs-inspired shell plugin for odoc.
   Registers as the "docsite" shell, usable with --shell docsite. *)

open Odoc_utils
module Html = Tyxml.Html
module Url = Odoc_document.Url

(* Register CSS and JS as support files *)
let () =
  Odoc_extension_registry.register_support_file ~prefix:"docsite"
    {
      filename = "extensions/docsite.css";
      content = Inline Odoc_docsite_css.css;
    };
  Odoc_extension_registry.register_support_file ~prefix:"docsite"
    {
      filename = "extensions/docsite.js";
      content = Inline Odoc_docsite_js.js;
    }

(* Site title from env var, default "Documentation" *)
let site_title =
  match Sys.getenv_opt "ODOC_SITE_TITLE" with
  | Some t -> t
  | None -> "Documentation"

(* --- Helpers --- *)

let file_uri ~config ~url (base : Odoc_html.Types.uri) file =
  match base with
  | Odoc_html.Types.Absolute uri -> uri ^ "/" ^ file
  | Relative uri ->
      let page = Url.Path.{ kind = `File; parent = uri; name = file } in
      Odoc_html.Link.href ~config ~resolve:(Current url) (Url.from_path page)

(* TyXML helper for generating TOC *)
let html_of_toc toc =
  let open Odoc_html.Types in
  let rec section (s : toc) =
    let link = Html.a ~a:[ Html.a_href s.href ] s.title in
    match s.children with [] -> [ link ] | cs -> [ link; sections cs ]
  and sections the_sections =
    the_sections
    |> List.map (fun s -> Html.li (section s))
    |> Html.ul
  in
  match toc with [] -> [] | _ -> [ sections toc ]

let toc_section toc =
  match toc with
  | [] -> []
  | _ ->
      [
        Html.div
          ~a:[ Html.a_class [ "odoc-tocs" ] ]
          [
            Html.nav
              ~a:[ Html.a_class [ "odoc-toc"; "odoc-local-toc" ] ]
              (Html.div
                 ~a:[ Html.a_class [ "odoc-toc-title" ] ]
                 [ Html.txt "On This Page" ]
               :: html_of_toc toc);
          ];
      ]

(* Breadcrumbs *)
let html_of_breadcrumbs (breadcrumbs : Odoc_html.Types.breadcrumbs) =
  let space = Html.txt " " in
  let sep = [ space; Html.entity "#x00BB"; space ] in
  let html =
    List.concat_map_sep ~sep
      ~f:(fun (breadcrumb : Odoc_html.Types.breadcrumb) ->
        match breadcrumb.href with
        | Some href ->
            [
              [
                Html.a
                  ~a:[ Html.a_href href ]
                  (breadcrumb.name
                    :> Html_types.flow5_without_interactive Html.elt list);
              ];
            ]
        | None ->
            [ (breadcrumb.name :> Html_types.nav_content_fun Html.elt list) ])
      breadcrumbs.parents
    |> List.flatten
  in
  let current_name :> Html_types.nav_content_fun Html.elt list =
    breadcrumbs.current.name
  in
  let rest =
    if List.is_empty breadcrumbs.parents then current_name
    else html @ sep @ current_name
  in
  (rest :> [< Html_types.nav_content_fun > `A `PCDATA `Wbr ] Html.elt list)

(* Serialize sidebar data to JSON for inline embedding *)
let sidebar_json_script sidebar_data =
  match sidebar_data with
  | None -> []
  | Some data ->
      let json = Odoc_html.Sidebar.to_json data in
      let json_str = Json.to_string json in
      [
        Html.script
          (Html.cdata_script
             (Printf.sprintf "window.__DOCSITE_SIDEBAR_DATA__ = %s;" json_str));
      ]

(* --- Page assembly --- *)

let page_creator ~config ~url ~uses_katex ~resources ~sidebar_data ~toc header
    breadcrumbs content =
  let support_uri = Odoc_html.Config.support_uri config in
  let path = Odoc_html.Link.Path.for_printing url in
  let file_uri = file_uri ~config ~url in
  let docsite_css_uri = file_uri support_uri "extensions/docsite.css" in
  let docsite_js_uri = file_uri support_uri "extensions/docsite.js" in

  (* Compute BASE_URL - relative path from current page to root *)
  let base_url =
    let page =
      Url.Path.{ kind = `File; parent = None; name = "" }
    in
    Odoc_html.Link.href ~config ~resolve:(Current url) (Url.from_path page)
  in

  (* Current URL as relative path from root *)
  let current_url =
    let filename = Odoc_html.Link.Path.as_filename ~config url in
    Fpath.to_string filename
  in

  (* Deduplicate resources *)
  let deduplicate_resources resources =
    let rec aux seen acc = function
      | [] -> List.rev acc
      | r :: rest ->
          if List.mem r seen then aux seen acc rest
          else aux (r :: seen) (r :: acc) rest
    in
    aux [] [] resources
  in

  (* Extension resources *)
  let extension_head_elements =
    let open Odoc_extension_registry in
    let is_absolute_url url =
      String.is_prefix ~affix:"http://" url
      || String.is_prefix ~affix:"https://" url
    in
    let resources = deduplicate_resources resources in
    List.concat_map
      (function
        | Js_url js_url ->
            let resolved =
              if is_absolute_url js_url then js_url
              else file_uri support_uri js_url
            in
            [ Html.script ~a:[ Html.a_src resolved ] (Html.txt "") ]
        | Css_url css_url ->
            let resolved =
              if is_absolute_url css_url then css_url
              else file_uri support_uri css_url
            in
            [ Html.link ~rel:[ `Stylesheet ] ~href:resolved () ]
        | Js_inline code ->
            let id = Printf.sprintf "%x" (Hashtbl.hash code land 0x7FFFFFFF) in
            [ Html.script ~a:[ Html.a_user_data "spa-inline" id ]
                (Html.cdata_script code) ]
        | Css_inline code -> [ Html.style [ Html.cdata_style code ] ])
      resources
  in

  (* KaTeX support *)
  let katex_elements =
    if uses_katex then
      let theme_uri = Odoc_html.Config.theme_uri config in
      let katex_css_uri = file_uri theme_uri "katex.min.css" in
      let katex_js_uri = file_uri support_uri "katex.min.js" in
      [
        Html.link ~rel:[ `Stylesheet ] ~href:katex_css_uri ();
        Html.script ~a:[ Html.a_src katex_js_uri ] (Html.txt "");
        Html.script
          (Html.cdata_script
             {|
          document.addEventListener("DOMContentLoaded", function () {
            var macros = {};
            var elements = Array.from(document.getElementsByClassName("odoc-katex-math"));
            for (var i = 0; i < elements.length; i++) {
              var el = elements[i];
              var content = el.textContent;
              var new_el = document.createElement("span");
              new_el.setAttribute("class", "odoc-katex-math-rendered");
              var display = el.classList.contains("display");
              katex.render(content, new_el, { throwOnError: false, displayMode: display, macros });
              el.replaceWith(new_el);
            }
          });
        |});
      ]
    else []
  in

  let head : Html_types.head Html.elt =
    let title_string =
      Printf.sprintf "%s (%s)" url.name (String.concat ~sep:"." path)
    in
    let meta_elements =
      [
        Html.meta ~a:[ Html.a_charset "utf-8" ] ();
        Html.link ~rel:[ `Stylesheet ] ~href:docsite_css_uri ();
        Html.meta
          ~a:
            [
              Html.a_name "viewport";
              Html.a_content "width=device-width,initial-scale=1.0";
            ]
          ();
        Html.meta
          ~a:[ Html.a_name "generator"; Html.a_content "odoc %%VERSION%%" ]
          ();
        (* Inject BASE_URL and CURRENT_URL for the JS *)
        Html.script
          (Html.Unsafe.data
             (Printf.sprintf "window.BASE_URL = %S; window.CURRENT_URL = %S;"
                base_url current_url));
      ]
      @ katex_elements @ extension_head_elements
      @ sidebar_json_script sidebar_data
    in
    Html.head (Html.title (Html.txt title_string)) meta_elements
  in

  (* Compute the root-relative href for the brand link *)
  let brand_href =
    let index_path =
      Url.Path.{ kind = `LeafPage; parent = None; name = "index" }
    in
    Odoc_html.Link.href ~config ~resolve:(Current url) (Url.from_path index_path)
  in

  let body =
    [
      (* Header bar *)
      Html.header
        ~a:[ Html.a_class [ "docsite-header" ] ]
        [
          (* Mobile menu toggle *)
          Html.Unsafe.data
            {|<button class="menu-toggle" aria-label="Toggle menu">
              <svg width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                <line x1="3" y1="12" x2="21" y2="12"></line>
                <line x1="3" y1="6" x2="21" y2="6"></line>
                <line x1="3" y1="18" x2="21" y2="18"></line>
              </svg>
            </button>|};
          Html.a
            ~a:
              [
                Html.a_href brand_href;
                Html.a_class [ "docsite-header-brand" ];
              ]
            [ Html.txt site_title ];
          Html.div
            ~a:[ Html.a_class [ "docsite-search-container" ] ]
            [
              Html.div
                ~a:[ Html.a_class [ "search-wrapper" ] ]
                [
                  Html.Unsafe.data
                    {|<svg class="search-icon" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                      <circle cx="11" cy="11" r="8"></circle>
                      <line x1="21" y1="21" x2="16.65" y2="16.65"></line>
                    </svg>|};
                  Html.input
                    ~a:
                      [
                        Html.a_input_type `Text;
                        Html.a_class [ "search-input" ];
                        Html.a_placeholder "Search documentation...";
                      ]
                    ();
                  Html.span
                    ~a:[ Html.a_class [ "search-shortcut" ] ]
                    [ Html.txt "/" ];
                  Html.div ~a:[ Html.a_class [ "search-results" ] ] [];
                ];
            ];
        ];
      (* Layout: sidebar + main *)
      Html.div
        ~a:[ Html.a_class [ "docsite-layout" ] ]
        [
          (* Sidebar *)
          Html.nav
            ~a:[ Html.a_class [ "docsite-sidebar" ] ]
            [
              Html.div
                ~a:[ Html.a_class [ "package-selector" ] ]
                [
                  Html.Unsafe.data
                    {|<select id="package-select">
                      <option value="">All packages</option>
                    </select>|};
                ];
              Html.div
                ~a:
                  [
                    Html.a_class [ "sidebar-nav" ];
                    Html.a_id "sidebar-content";
                  ]
                [ Html.div ~a:[ Html.a_class [ "sidebar-loading" ] ] [ Html.txt "Loading..." ] ];
            ];
          (* Main content area *)
          Html.main
            ~a:[ Html.a_class [ "docsite-main" ] ]
            ([ Html.nav ~a:[ Html.a_class [ "odoc-nav" ] ] (html_of_breadcrumbs breadcrumbs) ]
             @ [ Html.div ~a:[ Html.a_class [ "odoc-content" ] ]
                   ((header :> Html_types.div_content Html.elt list)
                    @ content) ]
             @ toc_section toc);
        ];
      (* Docsite JS *)
      Html.script ~a:[ Html.a_src docsite_js_uri; Html.a_defer () ] (Html.txt "");
    ]
  in

  let htmlpp = Html.pp ~indent:(Odoc_html.Config.indent config) () in
  let html =
    Html.html head (Html.body ~a:[ Html.a_class [ "odoc"; "odoc-docsite" ] ] body)
  in
  let content ppf =
    htmlpp ppf html;
    Format.pp_force_newline ppf ()
  in
  content

let make ~config ~url ~header ~breadcrumbs ~sidebar_data ~toc ~uses_katex
    ~resources ~assets content children =
  let filename = Odoc_html.Link.Path.as_filename ~config url in
  let content =
    page_creator ~config ~url ~uses_katex ~resources ~sidebar_data ~toc header
      breadcrumbs content
  in
  { Odoc_document.Renderer.filename; content; children; path = url; assets }

let src_page_creator ~breadcrumbs ~config ~url ~header ~sidebar_data title
    content =
  let support_uri = Odoc_html.Config.support_uri config in
  let file_uri = file_uri ~config ~url in
  let docsite_css_uri = file_uri support_uri "extensions/docsite.css" in
  let docsite_js_uri = file_uri support_uri "extensions/docsite.js" in

  let base_url =
    let page =
      Url.Path.{ kind = `File; parent = None; name = "" }
    in
    Odoc_html.Link.href ~config ~resolve:(Current url) (Url.from_path page)
  in
  let current_url =
    let filename = Odoc_html.Link.Path.as_filename ~config url in
    Fpath.to_string filename
  in

  let brand_href =
    let index_path =
      Url.Path.{ kind = `LeafPage; parent = None; name = "index" }
    in
    Odoc_html.Link.href ~config ~resolve:(Current url) (Url.from_path index_path)
  in

  let head : Html_types.head Html.elt =
    let title_string = Printf.sprintf "Source: %s" title in
    let meta_elements =
      [
        Html.meta ~a:[ Html.a_charset "utf-8" ] ();
        Html.link ~rel:[ `Stylesheet ] ~href:docsite_css_uri ();
        Html.meta
          ~a:
            [
              Html.a_name "viewport";
              Html.a_content "width=device-width,initial-scale=1.0";
            ]
          ();
        Html.script
          (Html.Unsafe.data
             (Printf.sprintf "window.BASE_URL = %S; window.CURRENT_URL = %S;"
                base_url current_url));
      ]
      @ sidebar_json_script sidebar_data
    in
    Html.head (Html.title (Html.txt title_string)) meta_elements
  in

  let body =
    [
      Html.header
        ~a:[ Html.a_class [ "docsite-header" ] ]
        [
          Html.a
            ~a:
              [
                Html.a_href brand_href;
                Html.a_class [ "docsite-header-brand" ];
              ]
            [ Html.txt site_title ];
          Html.div
            ~a:[ Html.a_class [ "docsite-search-container" ] ]
            [
              Html.div
                ~a:[ Html.a_class [ "search-wrapper" ] ]
                [
                  Html.input
                    ~a:
                      [
                        Html.a_input_type `Text;
                        Html.a_class [ "search-input" ];
                        Html.a_placeholder "Search documentation...";
                      ]
                    ();
                  Html.div ~a:[ Html.a_class [ "search-results" ] ] [];
                ];
            ];
        ];
      Html.div
        ~a:[ Html.a_class [ "docsite-layout" ] ]
        [
          Html.nav
            ~a:[ Html.a_class [ "docsite-sidebar" ] ]
            [
              Html.div
                ~a:[ Html.a_class [ "package-selector" ] ]
                [
                  Html.Unsafe.data
                    {|<select id="package-select">
                      <option value="">All packages</option>
                    </select>|};
                ];
              Html.div
                ~a:
                  [
                    Html.a_class [ "sidebar-nav" ];
                    Html.a_id "sidebar-content";
                  ]
                [ Html.div ~a:[ Html.a_class [ "sidebar-loading" ] ] [ Html.txt "Loading..." ] ];
            ];
          Html.main
            ~a:[ Html.a_class [ "docsite-main" ] ]
            ([ Html.nav ~a:[ Html.a_class [ "odoc-nav" ] ] (html_of_breadcrumbs breadcrumbs) ]
             @ [ Html.header
                   ~a:[ Html.a_class [ "odoc-preamble" ] ]
                   (header :> Html_types.flow5_without_header_footer Html.elt list) ]
             @ content);
        ];
      Html.script ~a:[ Html.a_src docsite_js_uri; Html.a_defer () ] (Html.txt "");
    ]
  in

  let htmlpp = Html.pp ~indent:false () in
  let html =
    Html.html head
      (Html.body ~a:[ Html.a_class [ "odoc-src"; "odoc-docsite" ] ] body)
  in
  let content ppf =
    htmlpp ppf html;
    Format.pp_force_newline ppf ()
  in
  content

let make_src ~config ~url ~breadcrumbs ~header ~sidebar_data title content =
  let filename = Odoc_html.Link.Path.as_filename ~config url in
  let content =
    src_page_creator ~breadcrumbs ~config ~url ~header ~sidebar_data title
      content
  in
  {
    Odoc_document.Renderer.filename;
    content;
    children = [];
    path = url;
    assets = [];
  }

(* Register the shell *)
let () =
  Odoc_html.Html_shell.register
    (module struct
      let name = "docsite"

      let make ~config (data : Odoc_html.Html_shell.page_data) =
        make ~config ~url:data.url
          ~header:(data.header @ data.preamble)
          ~breadcrumbs:data.breadcrumbs ~sidebar_data:data.sidebar_data
          ~toc:data.toc ~uses_katex:data.uses_katex ~resources:data.resources
          ~assets:data.assets data.content data.children

      let make_src ~config (data : Odoc_html.Html_shell.src_page_data) =
        make_src ~config ~url:data.url ~breadcrumbs:data.breadcrumbs
          ~header:data.header ~sidebar_data:data.sidebar_data data.title
          data.content
    end)