(* Second round of resolution tackles references and forward paths *)
open Odoc_model
open Lang
module Id = Paths.Identifier

module Opt = struct
  let map f = function Some x -> Some (f x) | None -> None
end

(* omg. Our current warning system is spread on different system. Hence this
   atrocity. *)
let maybe_suppress env warnings_tag =
  if Env.should_suppress_warnings env warnings_tag then fun f ->
    Lookup_failures.catch_failures ~filename:"" (fun () ->
        Error.catch_warnings f |> fun x ->
        Error.unpack_warnings x |> fst |> Error.unpack_warnings |> fst)
    |> Error.unpack_warnings |> fst
  else fun f -> f () |> Error.raise_warnings

let source_loc env id loc =
  let id = (id :> Id.NonSrc.t) in
  match loc with Some _ as loc -> loc | None -> Shape_tools.lookup_def env id

(** Equivalent to {!Comment.synopsis}. *)
let synopsis_from_comment (docs : Component.CComment.docs) =
  match docs.elements with
  | ({ value = #Comment.nestable_block_element; _ } as e) :: _ ->
      (* Only the first element is considered. *)
      Comment.synopsis [ e ]
  | _ -> None

let synopsis_of_module env (m : Component.Module.t) =
  let open Odoc_utils.ResultMonad in
  match synopsis_from_comment m.doc with
  | Some _ as s -> s
  | None -> (
      let rec handle_expansion : Tools.expansion -> _ = function
        | Functor (_, expr) -> (
            match Tools.expansion_of_module_type_expr env expr with
            | Ok e -> handle_expansion e
            | Error _ as e -> e)
        | Signature sg -> Ok sg
      in
      (* If there is no doc, look at the expansion. *)
      match
        Tools.expansion_of_module env m >>= fun exp -> handle_expansion exp
      with
      | Ok sg -> synopsis_from_comment (Component.extract_signature_doc sg)
      | Error _ -> None)

let ambiguous_label_warning label_name labels =
  let pp_label_loc fmt (`Label (_, x)) =
    Location_.pp_span_start fmt x.Component.Label.location
  in
  Lookup_failures.report_warning
    "@[<2>Label '%s' is ambiguous. The other occurences are:@ %a@]" label_name
    (Format.pp_print_list ~pp_sep:Format.pp_force_newline pp_label_loc)
    labels

(** Raise a warning when a label explicitly set by the user collides. This
    warning triggers even if one of the colliding labels have been automatically
    generated. *)
let check_ambiguous_label ~loc env
    ( attrs,
      ({ Odoc_model.Paths.Identifier.iv = `Label (_, label_name); _ } as id),
      _ ) =
  if attrs.Comment.heading_label_explicit then
    (* Looking for an identical identifier but a different location. *)
    let conflicting (`Label (id', comp)) =
      Id.equal id id'
      && not (Location_.span_equal comp.Component.Label.location loc)
    in
    let label_name = Names.LabelName.to_string label_name in
    match Env.lookup_by_name Env.s_label label_name env with
    | Ok lbl when conflicting lbl -> ambiguous_label_warning label_name [ lbl ]
    | Error (`Ambiguous (hd, tl)) -> (
        match List.filter conflicting (hd :: tl) with
        | [] -> ()
        | xs -> ambiguous_label_warning label_name xs)
    | Ok _ | Error `Not_found -> ()

let expansion_needed self target =
  let self = (self :> Paths.Path.Resolved.t) in
  let hidden_alias = Paths.Path.Resolved.is_hidden self
  and self_canonical =
    let i = Paths.Path.Resolved.identifier self in
    i = Some (target :> Paths.Identifier.t)
  in

  self_canonical || hidden_alias

exception Loop

let rec is_forward : Paths.Path.Module.t -> bool = function
  | `Resolved _ -> false
  | `Root _ -> false
  | `Forward _ -> true
  | `Identifier _ -> false
  | `Dot (p, _) -> is_forward p
  | `Apply (p1, p2) -> is_forward p1 || is_forward p2
  | `Substituted s -> is_forward s

let rec should_reresolve : Paths.Path.Resolved.t -> bool =
 fun p ->
  let open Paths.Path.Resolved in
  match p with
  | `CoreType _ -> false
  | `Identifier _ -> false
  | `Subst (x, y) -> should_reresolve (x :> t) || should_reresolve (y :> t)
  | `Hidden p -> should_reresolve (p :> t)
  | `Canonical (x, y) ->
      should_reresolve (x :> t) || should_resolve (y :> Paths.Path.t)
  | `CanonicalModuleType (x, y) ->
      should_reresolve (x :> t) || should_resolve (y :> Paths.Path.t)
  | `CanonicalType (x, y) ->
      should_reresolve (x :> t) || should_resolve (y :> Paths.Path.t)
  | `Apply (x, y) ->
      should_reresolve (x :> t) || should_reresolve (y :> Paths.Path.Resolved.t)
  | `SubstT (x, y) -> should_reresolve (x :> t) || should_reresolve (y :> t)
  | `Alias (y, x) ->
      should_resolve (x :> Paths.Path.t) || should_reresolve (y :> t)
  | `AliasModuleType (x, y) ->
      should_reresolve (x :> t) || should_reresolve (y :> t)
  | `Type (p, _)
  | `Value (p, _)
  | `Class (p, _)
  | `ClassType (p, _)
  | `ModuleType (p, _)
  | `Module (p, _) ->
      should_reresolve (p :> t)
  | `OpaqueModule m -> should_reresolve (m :> t)
  | `OpaqueModuleType m -> should_reresolve (m :> t)
  | `Substituted m -> should_reresolve (m :> t)
  | `SubstitutedMT m -> should_reresolve (m :> t)
  | `SubstitutedT m -> should_reresolve (m :> t)
  | `SubstitutedCT m -> should_reresolve (m :> t)

and should_resolve : Paths.Path.t -> bool =
 fun p -> match p with `Resolved p -> should_reresolve p | _ -> true

let type_path : Env.t -> Paths.Path.Type.t -> Paths.Path.Type.t =
 fun env p ->
  if not (should_resolve (p :> Paths.Path.t)) then p
  else
    let cp = Component.Of_Lang.(type_path (empty ()) p) in
    match cp with
    | `Resolved p ->
        let result = Tools.reresolve_type env p in
        `Resolved Lang_of.(Path.resolved_type (empty ()) result)
    | _ -> (
        match Tools.resolve_type_path env cp with
        | Ok p' ->
            let result = Tools.reresolve_type env p' in
            `Resolved Lang_of.(Path.resolved_type (empty ()) result)
        | Error e ->
            Errors.report ~what:(`Type_path cp) ~tools_error:e `Lookup;
            p)

let value_path : Env.t -> Paths.Path.Value.t -> Paths.Path.Value.t =
 fun env p ->
  if not (should_resolve (p :> Paths.Path.t)) then p
  else
    let cp = Component.Of_Lang.(value_path (empty ()) p) in
    match cp with
    | `Resolved p ->
        let result = Tools.reresolve_value env p in
        `Resolved Lang_of.(Path.resolved_value (empty ()) result)
    | _ -> (
        match Tools.resolve_value_path env cp with
        | Ok p' ->
            let result = Tools.reresolve_value env p' in
            `Resolved Lang_of.(Path.resolved_value (empty ()) result)
        | Error e ->
            Errors.report ~what:(`Value_path cp) ~tools_error:e `Lookup;
            p)

let class_type_path : Env.t -> Paths.Path.ClassType.t -> Paths.Path.ClassType.t
    =
 fun env p ->
  if not (should_resolve (p :> Paths.Path.t)) then p
  else
    let cp = Component.Of_Lang.(class_type_path (empty ()) p) in
    match cp with
    | `Resolved p ->
        let result = Tools.reresolve_class_type env p in
        `Resolved Lang_of.(Path.resolved_class_type (empty ()) result)
    | _ -> (
        match Tools.resolve_class_type_path env cp with
        | Ok p' ->
            let result = Tools.reresolve_class_type env p' in
            `Resolved Lang_of.(Path.resolved_class_type (empty ()) result)
        | Error e ->
            Errors.report ~what:(`Class_type_path cp) ~tools_error:e `Lookup;
            p)

and module_type_path :
    Env.t -> Paths.Path.ModuleType.t -> Paths.Path.ModuleType.t =
 fun env p ->
  if not (should_resolve (p :> Paths.Path.t)) then p
  else
    let cp = Component.Of_Lang.(module_type_path (empty ()) p) in
    match cp with
    | `Resolved p ->
        let result = Tools.reresolve_module_type env p in
        `Resolved Lang_of.(Path.resolved_module_type (empty ()) result)
    | _ -> (
        match Tools.resolve_module_type_path env cp with
        | Ok p' ->
            let result = Tools.reresolve_module_type env p' in
            `Resolved Lang_of.(Path.resolved_module_type (empty ()) result)
        | Error e ->
            Errors.report ~what:(`Module_type_path cp) ~tools_error:e `Resolve;
            p)

and module_path : Env.t -> Paths.Path.Module.t -> Paths.Path.Module.t =
 fun env p ->
  if not (should_resolve (p :> Paths.Path.t)) then p
  else
    let cp = Component.Of_Lang.(module_path (empty ()) p) in
    match cp with
    | `Resolved p ->
        let after = Tools.reresolve_module env p in
        `Resolved Lang_of.(Path.resolved_module (empty ()) after)
    | _ -> (
        match Tools.resolve_module_path env cp with
        | Ok p' ->
            let result = Tools.reresolve_module env p' in
            `Resolved Lang_of.(Path.resolved_module (empty ()) result)
        | Error _ when is_forward p -> p
        | Error e ->
            Errors.report ~what:(`Module_path cp) ~tools_error:e `Resolve;
            p)

let rec comment_inline_element :
    loc:_ ->
    Env.t ->
    string option ->
    Comment.inline_element ->
    Comment.inline_element =
 fun ~loc:_ env warnings_tag x ->
  match x with
  | `Styled (s, ls) ->
      `Styled
        ( s,
          List.map (with_location (comment_inline_element env warnings_tag)) ls
        )
  | `Reference (r, content) as orig -> (
      match
        maybe_suppress env warnings_tag (fun () ->
            Ref_tools.resolve_reference env r)
      with
      | Ok (ref_, c) ->
          let content =
            (* In case of labels, use the heading text as reference text if
               it's not specified. *)
            match (content, c) with
            | [], Some content ->
                Comment.link_content_of_inline_elements content
            | content, _ -> content
          in
          `Reference (`Resolved ref_, content)
      | Error e ->
          if not (Env.should_suppress_warnings env warnings_tag) then
            Errors.report ~what:(`Reference r) ~tools_error:(`Reference e)
              `Resolve;
          orig)
  | y -> y

and paragraph env warnings_tag elts =
  List.map (with_location (comment_inline_element env warnings_tag)) elts

and resolve_external_synopsis env warnings_tag synopsis =
  let env = Env.inherit_resolver env in
  paragraph env warnings_tag synopsis

and comment_nestable_block_element env warnings_tag parent ~loc:_
    (x : Comment.nestable_block_element) =
  match x with
  | `Paragraph elts -> `Paragraph (paragraph env warnings_tag elts)
  | (`Code_block _ | `Math_block _ | `Verbatim _) as x -> x
  | `List (x, ys) ->
      `List
        ( x,
          List.rev_map
            (comment_nestable_block_element_list env warnings_tag parent)
            ys
          |> List.rev )
  | `Table { data; align } ->
      let data =
        let map f x = List.rev_map f x |> List.rev in
        map
          (map (fun (cell, cell_type) ->
               ( comment_nestable_block_element_list env warnings_tag parent cell,
                 cell_type )))
          data
      in
      `Table { Comment.data; align }
  | `Modules refs ->
      let refs =
        List.rev_map
          (fun (r : Comment.module_reference) ->
            match
              maybe_suppress env warnings_tag (fun () ->
                  Ref_tools.resolve_module_reference env r.module_reference)
            with
            | Ok (r, _, m) ->
                let module_synopsis =
                  Opt.map
                    (resolve_external_synopsis env warnings_tag)
                    (synopsis_of_module env m)
                in
                { Comment.module_reference = `Resolved r; module_synopsis }
            | Error e ->
                if not (Env.should_suppress_warnings env warnings_tag) then
                  Errors.report
                    ~what:(`Reference (r.module_reference :> Paths.Reference.t))
                    ~tools_error:(`Reference e) `Resolve;
                r)
          refs
        |> List.rev
      in
      `Modules refs
  | `Media (`Reference r, m, content) as orig -> (
      match
        maybe_suppress env warnings_tag (fun () ->
            Ref_tools.resolve_asset_reference env r)
      with
      | Ok x -> `Media (`Reference (`Resolved x), m, content)
      | Error e ->
          if not (Env.should_suppress_warnings env warnings_tag) then
            Errors.report
              ~what:(`Reference (r :> Paths.Reference.t))
              ~tools_error:(`Reference e) `Resolve;
          orig)
  | `Media _ as orig -> orig

and comment_nestable_block_element_list env warnings_tag parent
    (xs : Comment.nestable_block_element Comment.with_location list) =
  List.rev_map
    (with_location (comment_nestable_block_element env warnings_tag parent))
    xs
  |> List.rev

and comment_tag env warnings_tag parent ~loc:_ (x : Comment.tag) =
  match x with
  | `Deprecated content ->
      `Deprecated
        (comment_nestable_block_element_list env warnings_tag parent content)
  | `Param (name, content) ->
      `Param
        ( name,
          comment_nestable_block_element_list env warnings_tag parent content )
  | `Raise ((`Reference (r, reference_content) as orig), content) -> (
      match
        maybe_suppress env warnings_tag (fun () ->
            Ref_tools.resolve_reference env r)
      with
      | Ok (x, _) ->
          `Raise
            ( `Reference (`Resolved x, reference_content),
              comment_nestable_block_element_list env warnings_tag parent
                content )
      | Error e ->
          if not (Env.should_suppress_warnings env warnings_tag) then
            Errors.report ~what:(`Reference r) ~tools_error:(`Reference e)
              `Resolve;
          `Raise
            ( orig,
              comment_nestable_block_element_list env warnings_tag parent
                content ))
  | `Raise ((`Code_span _ as orig), content) ->
      `Raise
        ( orig,
          comment_nestable_block_element_list env warnings_tag parent content )
  | `Return content ->
      `Return
        (comment_nestable_block_element_list env warnings_tag parent content)
  | `See (kind, target, content) ->
      `See
        ( kind,
          target,
          comment_nestable_block_element_list env warnings_tag parent content )
  | `Before (version, content) ->
      `Before
        ( version,
          comment_nestable_block_element_list env warnings_tag parent content )
  | `Author _ | `Since _ | `Alert _ | `Version _ ->
      x (* only contain primitives *)

and comment_block_element env warnings_tag parent ~loc
    (x : Comment.block_element) =
  match x with
  | #Comment.nestable_block_element as x ->
      (comment_nestable_block_element env warnings_tag parent ~loc x
        :> Comment.block_element)
  | `Heading (attrs, label, elems) ->
      let cie = comment_inline_element env warnings_tag in
      let elems =
        List.rev_map (fun ele -> with_location cie ele) elems |> List.rev
      in
      let h = (attrs, label, elems) in
      check_ambiguous_label ~loc env h;
      `Heading h
  | `Tag t -> `Tag (comment_tag env warnings_tag parent ~loc t)

and with_location : type a.
    (loc:_ -> a -> a) -> a Location_.with_location -> a Location_.with_location
    =
 fun fn { value; location = loc } ->
  let value = Lookup_failures.with_location loc (fun () -> fn ~loc value) in
  { value; location = loc }

and comment_docs env parent d =
  {
    Comment.elements =
      List.rev_map
        (with_location
           (comment_block_element env d.Comment.warnings_tag
              (parent :> Id.LabelParent.t)))
        d.Comment.elements
      |> List.rev;
    warnings_tag = d.warnings_tag;
  }

and comment env parent = function
  | `Stop -> `Stop
  | `Docs d -> `Docs (comment_docs env parent d)

and open_ env parent = function
  | { Odoc_model__Lang.Open.doc; _ } as open_ ->
      { open_ with doc = comment_docs env parent doc }

let warn_on_hidden_representation (id : Id.Type.t)
    (r : Lang.TypeDecl.Representation.t) =
  let open Paths.Identifier in
  let rec internal_typ_exp typ_expr =
    let open Lang.TypeExpr in
    let open Paths.Path in
    match typ_expr with
    | Constr (p, ts) ->
        is_hidden (p :> Paths.Path.t)
        || List.exists (fun t -> internal_typ_exp t) ts
    | Poly (_, t) | Alias (t, _) -> internal_typ_exp t
    | Arrow (_, t, t2) -> internal_typ_exp t || internal_typ_exp t2
    | Tuple ts -> List.exists (fun (_, t) -> internal_typ_exp t) ts
    | Class (_, ts) -> List.exists (fun t -> internal_typ_exp t) ts
    | _ -> false
  in

  let internal_cstr_arg t =
    let open Lang.TypeDecl.Constructor in
    let open Lang.TypeDecl.Field in
    match t.args with
    | Tuple type_exprs ->
        List.exists (fun type_expr -> internal_typ_exp type_expr) type_exprs
    | Record fields ->
        List.exists (fun field -> internal_typ_exp field.type_) fields
  in

  let internal_field t =
    let open Lang.TypeDecl.Field in
    internal_typ_exp t.type_
  in

  let fmt_cfg = Component.Fmt.{ default with short_paths = true } in
  match r with
  | Variant constructors ->
      if List.exists internal_cstr_arg constructors then
        Lookup_failures.report_warning "@[<2>Hidden constructors in type '%a'@]"
          Component.Fmt.(model_identifier fmt_cfg)
          (id :> Id.any)
  | Record fields ->
      if List.exists internal_field fields then
        Lookup_failures.report_warning "@[<2>Hidden fields in type '%a'@]"
          Component.Fmt.(model_identifier fmt_cfg)
          (id :> Id.any)
  | Extensible -> ()

let rec unit env t =
  let open Compilation_unit in
  let content =
    if t.hidden then t.content
    else
      match t.content with
      | Module sg ->
          let sg = signature env (t.id :> Id.Signature.t) sg in
          Module sg
      | Pack _ as p -> p
  in
  let source_loc = source_loc env t.id t.source_loc in
  { t with content; linked = true; source_loc }

and value_ env parent t =
  let open Value in
  {
    t with
    source_loc = source_loc env t.id t.source_loc;
    doc = comment_docs env parent t.doc;
    type_ = type_expression env parent [] t.type_;
  }

and exception_ env parent e =
  let open Exception in
  let res = Opt.map (type_expression env parent []) e.res in
  let args = type_decl_constructor_argument env parent e.args in
  let source_loc = source_loc env e.id e.source_loc in
  let doc = comment_docs env parent e.doc in
  { e with source_loc; res; args; doc }

and extension env parent t =
  let open Extension in
  let constructor c =
    let open Constructor in
    {
      c with
      source_loc = source_loc env c.id c.source_loc;
      args = type_decl_constructor_argument env parent c.args;
      res = Opt.map (type_expression env parent []) c.res;
      doc = comment_docs env parent c.doc;
    }
  in
  let type_path = type_path env t.type_path in
  let constructors = List.map constructor t.constructors in
  let doc = comment_docs env parent t.doc in
  { t with type_path; constructors; doc }

and class_type_expr env parent =
  let open ClassType in
  function
  | Constr (path, texps) ->
      Constr (path, List.map (type_expression env parent []) texps)
  | Signature s -> Signature (class_signature env parent s)

and class_type env parent c =
  let open ClassType in
  let doc = comment_docs env parent c.doc in
  {
    c with
    source_loc = source_loc env c.id c.source_loc;
    expr = class_type_expr env parent c.expr;
    doc;
  }

and class_signature env parent c =
  let open ClassSignature in
  let env = Env.open_class_signature c env in
  let map_item = function
    | Method m -> Method (method_ env parent m)
    | InstanceVariable i -> InstanceVariable (instance_variable env parent i)
    | Constraint cst -> Constraint (constraint_ env parent cst)
    | Inherit c -> Inherit (inherit_ env parent c)
    | Comment c -> Comment c
  in
  {
    self = Opt.map (type_expression env parent []) c.self;
    items = List.map map_item c.items;
    doc = comment_docs env parent c.doc;
  }

and method_ env parent m =
  let open Method in
  let doc = comment_docs env parent m.doc in
  { m with type_ = type_expression env parent [] m.type_; doc }

and instance_variable env parent i =
  let open InstanceVariable in
  let doc = comment_docs env parent i.doc in
  { i with type_ = type_expression env parent [] i.type_; doc }

and constraint_ env parent cst =
  let open ClassSignature.Constraint in
  let left = type_expression env parent [] cst.left
  and right = type_expression env parent [] cst.right
  and doc = comment_docs env parent cst.doc in
  { left; right; doc }

and inherit_ env parent ih =
  let open ClassSignature.Inherit in
  let expr = class_type_expr env parent ih.expr
  and doc = comment_docs env parent ih.doc in
  { expr; doc }

and class_ env parent c =
  let open Class in
  let rec map_decl = function
    | ClassType expr -> ClassType (class_type_expr env parent expr)
    | Arrow (lbl, expr, decl) ->
        Arrow (lbl, type_expression env parent [] expr, map_decl decl)
  in
  let doc = comment_docs env parent c.doc in
  let source_loc = source_loc env c.id c.source_loc in
  let type_ = map_decl c.type_ in
  { c with source_loc; type_; doc }

and module_substitution env parent m =
  let open ModuleSubstitution in
  let doc = comment_docs env parent m.doc in
  { m with manifest = module_path env m.manifest; doc }

and signature : Env.t -> Id.Signature.t -> Signature.t -> _ =
 fun env id s ->
  let env = Env.open_signature s env |> Env.add_docs s.doc in
  let items = signature_items env id s.items
  and doc = comment_docs env id s.doc in
  { s with items; doc }

and signature_items :
    Env.t -> Id.Signature.t -> Signature.item list -> Signature.item list =
 fun env id s ->
  let open Signature in
  let items, _ =
    List.fold_left
      (fun (items, env) item ->
        let std i = (i :: items, env) in
        match item with
        | Module (r, m) -> std @@ Module (r, module_ env m)
        | ModuleSubstitution m ->
            let env' = Env.open_module_substitution m env in
            (ModuleSubstitution (module_substitution env id m) :: items, env')
        | Type (r, t) -> std @@ Type (r, type_decl env id t)
        | TypeSubstitution t ->
            let env' = Env.open_type_substitution t env in
            (TypeSubstitution (type_decl env id t) :: items, env')
        | ModuleType mt -> std @@ ModuleType (module_type env mt)
        | ModuleTypeSubstitution mts ->
            let env' = Env.open_module_type_substitution mts env in
            ( ModuleTypeSubstitution (module_type_substitution env mts) :: items,
              env' )
        | Value v -> std @@ Value (value_ env id v)
        | Comment c -> std @@ Comment (comment env id c)
        | TypExt t -> std @@ TypExt (extension env id t)
        | Exception e -> std @@ Exception (exception_ env id e)
        | Class (r, c) -> std @@ Class (r, class_ env id c)
        | ClassType (r, c) -> std @@ ClassType (r, class_type env id c)
        | Include i -> std @@ Include (include_ env i)
        | Open o -> std @@ Open (open_ env id o))
      ([], env) s
  in
  List.rev items

and simple_expansion :
    Env.t ->
    Id.Signature.t ->
    ModuleType.simple_expansion ->
    ModuleType.simple_expansion =
 fun env id m ->
  match m with
  | Signature sg -> Signature (signature env id sg)
  | Functor (arg, sg) ->
      let env' = Env.add_functor_parameter arg env in
      Functor (functor_argument env arg, simple_expansion env' id sg)

and module_ : Env.t -> Module.t -> Module.t =
 fun env m ->
  let open Module in
  let open Odoc_utils.ResultMonad in
  let sg_id = (m.id :> Id.Signature.t) in
  if m.hidden then m
  else
    let type_ = module_decl env sg_id m.type_ in
    let type_ =
      match type_ with
      | Alias (`Resolved p, _) ->
          if expansion_needed p m.id then
            let cp = Component.Of_Lang.(resolved_module_path (empty ()) p) in
            match
              Tools.expansion_of_module_path ~strengthen:false env
                (`Resolved cp)
              >>= Expand_tools.handle_expansion env (m.id :> Id.Signature.t)
            with
            | Ok (_, e) ->
                let le = Lang_of.(simple_expansion (empty ()) sg_id e) in
                Alias (`Resolved p, Some (simple_expansion env sg_id le))
            | Error _ -> type_
          else type_
      | Alias _ | ModuleType _ -> type_
    in
    let source_loc = source_loc env m.id m.source_loc in
    let doc = comment_docs env sg_id m.doc in
    { m with source_loc; doc; type_ }

and module_decl : Env.t -> Id.Signature.t -> Module.decl -> Module.decl =
 fun env id decl ->
  let open Module in
  match decl with
  | ModuleType expr -> ModuleType (module_type_expr env id expr)
  | Alias (p, e) ->
      Alias (module_path env p, Opt.map (simple_expansion env id) e)

and include_decl : Env.t -> Id.Signature.t -> Include.decl -> Include.decl =
 fun env id decl ->
  let open Include in
  match decl with
  | ModuleType expr -> ModuleType (u_module_type_expr env id expr)
  | Alias p -> Alias (module_path env p)

and module_type : Env.t -> ModuleType.t -> ModuleType.t =
 fun env m ->
  let sg_id = (m.id :> Id.Signature.t) in
  let open ModuleType in
  let expr' =
    match m.expr with
    | None -> None
    | Some expr -> Some (module_type_expr env sg_id expr)
  in
  (* let self_canonical =
       match m.expr with
       | Some (Path (`Resolved p)) when Paths.Path.Resolved.ModuleType.canonical_ident p = Some m.id ->
         true
       | _ -> false
     in*)
  let doc = comment_docs env sg_id m.doc in
  let source_loc = (source_loc env m.id) m.source_loc in
  { m with source_loc; expr = expr'; doc }

and module_type_substitution :
    Env.t -> ModuleTypeSubstitution.t -> ModuleTypeSubstitution.t =
 fun env m ->
  let sg_id = (m.id :> Id.Signature.t) in
  let open ModuleTypeSubstitution in
  let manifest' = module_type_expr env sg_id m.manifest in
  let doc = comment_docs env sg_id m.doc in
  { m with manifest = manifest'; doc }

and include_ : Env.t -> Include.t -> Include.t =
 fun env i ->
  let open Include in
  let decl = include_decl env i.parent i.decl in
  let doc = comment_docs env i.parent i.doc in
  let expansion =
    (* Don't call {!signature} to avoid adding the content of the expansion to
       the environment, which is already done recursively by
       {!Env.open_signature}. *)
    let content =
      (* Add context around errors from the expansion. *)
      Lookup_failures.with_context
        "While resolving the expansion of include at %a" Location_.pp_span_start
        i.loc (fun () ->
          let { content; _ } = i.expansion in
          let items = signature_items env i.parent content.items
          and doc = comment_docs env i.parent content.doc in
          { content with items; doc })
    in
    { i.expansion with content }
  in
  { i with decl; expansion; doc }

and functor_parameter_parameter :
    Env.t -> FunctorParameter.parameter -> FunctorParameter.parameter =
 fun env a ->
  let sg_id = (a.id :> Id.Signature.t) in
  let expr = module_type_expr env sg_id a.expr in
  { a with expr }

and functor_argument env a =
  match a with
  | FunctorParameter.Unit -> FunctorParameter.Unit
  | Named arg -> Named (functor_parameter_parameter env arg)

and handle_fragments env id sg subs =
  let open ModuleType in
  List.fold_left
    (fun (sg_res, subs) lsub ->
      match (sg_res, lsub) with
      | Ok sg, ModuleEq (frag, decl) ->
          let frag' =
            match frag with
            | `Resolved f ->
                let cfrag =
                  Component.Of_Lang.(resolved_module_fragment (empty ()) f)
                in
                `Resolved
                  (Tools.reresolve_module_fragment env cfrag
                  |> Lang_of.(Path.resolved_module_fragment (empty ())))
            | _ -> frag
          in
          let sg' =
            Tools.fragmap env
              Component.Of_Lang.(with_module_type_substitution (empty ()) lsub)
              sg
          in
          (sg', ModuleEq (frag', module_decl env id decl) :: subs)
      | Ok sg, TypeEq (frag, eqn) ->
          let frag' =
            match frag with
            | `Resolved f ->
                let cfrag =
                  Component.Of_Lang.(resolved_type_fragment (empty ()) f)
                in
                `Resolved
                  (Tools.reresolve_type_fragment env cfrag
                  |> Lang_of.(Path.resolved_type_fragment (empty ())))
            | _ -> frag
          in
          let sg' =
            Tools.fragmap env
              Component.Of_Lang.(with_module_type_substitution (empty ()) lsub)
              sg
          in
          (sg', TypeEq (frag', type_decl_equation env id eqn) :: subs)
      | Ok sg, ModuleTypeEq (frag, eqn) ->
          let frag' =
            match frag with
            | `Resolved f ->
                let cfrag =
                  Component.Of_Lang.(resolved_module_type_fragment (empty ()) f)
                in
                `Resolved
                  (Tools.reresolve_module_type_fragment env cfrag
                  |> Lang_of.(Path.resolved_module_type_fragment (empty ())))
            | _ -> frag
          in
          let sg' =
            Tools.fragmap env
              Component.Of_Lang.(with_module_type_substitution (empty ()) lsub)
              sg
          in
          (sg', ModuleTypeEq (frag', module_type_expr env id eqn) :: subs)
      | Ok sg, ModuleSubst (frag, mpath) ->
          let frag' =
            match frag with
            | `Resolved f ->
                let cfrag =
                  Component.Of_Lang.(resolved_module_fragment (empty ()) f)
                in
                `Resolved
                  (Tools.reresolve_module_fragment env cfrag
                  |> Lang_of.(Path.resolved_module_fragment (empty ())))
            | _ -> frag
          in
          let sg' =
            Tools.fragmap env
              Component.Of_Lang.(with_module_type_substitution (empty ()) lsub)
              sg
          in
          (sg', ModuleSubst (frag', module_path env mpath) :: subs)
      | Ok sg, TypeSubst (frag, eqn) ->
          let frag' =
            match frag with
            | `Resolved f ->
                let cfrag =
                  Component.Of_Lang.(resolved_type_fragment (empty ()) f)
                in
                `Resolved
                  (Tools.reresolve_type_fragment env cfrag
                  |> Lang_of.(Path.resolved_type_fragment (empty ())))
            | _ -> frag
          in
          let sg' =
            Tools.fragmap env
              Component.Of_Lang.(with_module_type_substitution (empty ()) lsub)
              sg
          in
          (sg', TypeSubst (frag', type_decl_equation env id eqn) :: subs)
      | Ok sg, ModuleTypeSubst (frag, eqn) ->
          let frag' =
            match frag with
            | `Resolved f ->
                let cfrag =
                  Component.Of_Lang.(resolved_module_type_fragment (empty ()) f)
                in
                `Resolved
                  (Tools.reresolve_module_type_fragment env cfrag
                  |> Lang_of.(Path.resolved_module_type_fragment (empty ())))
            | _ -> frag
          in
          let sg' =
            Tools.fragmap env
              Component.Of_Lang.(with_module_type_substitution (empty ()) lsub)
              sg
          in
          (sg', ModuleTypeSubst (frag', module_type_expr env id eqn) :: subs)
      | (Error _ as e), lsub -> (e, lsub :: subs))
    (Ok sg, []) subs
  |> snd |> List.rev

and u_module_type_expr :
    Env.t -> Id.Signature.t -> ModuleType.U.expr -> ModuleType.U.expr =
 fun env id expr ->
  match expr with
  | Signature s -> Signature s
  (* No need to link 'unexpanded' module type expressions that are actually expanded... *)
  | Path p -> Path (module_type_path env p)
  | With (subs, expr) as unresolved -> (
      let cexpr = Component.Of_Lang.(u_module_type_expr (empty ()) expr) in
      match Tools.signature_of_u_module_type_expr env cexpr with
      | Ok sg ->
          With (handle_fragments env id sg subs, u_module_type_expr env id expr)
      | Error e ->
          Errors.report ~what:(`Module_type_U cexpr) ~tools_error:e `Resolve;
          unresolved)
  | TypeOf (StructInclude p, original_path) ->
      TypeOf (StructInclude (module_path env p), original_path)
  | TypeOf (ModPath p, original_path) ->
      TypeOf (ModPath (module_path env p), original_path)

and module_type_expr :
    Env.t -> Id.Signature.t -> ModuleType.expr -> ModuleType.expr =
 fun env id expr ->
  let open ModuleType in
  let open Odoc_utils.ResultMonad in
  let do_expn cur (e : Paths.Path.ModuleType.t option) =
    match (cur, e) with
    | Some e, _ ->
        Some (simple_expansion env (id :> Paths.Identifier.Signature.t) e)
    | None, Some (`Resolved p_path) ->
        if expansion_needed p_path id then
          let cp =
            Component.Of_Lang.(resolved_module_type_path (empty ()) p_path)
          in
          match
            Tools.expansion_of_module_type_expr env
              (Path { p_path = `Resolved cp; p_expansion = None })
            >>= Expand_tools.handle_expansion env (id :> Id.Signature.t)
          with
          | Ok (_, e) ->
              let le = Lang_of.(simple_expansion (empty ()) id e) in
              Some (simple_expansion env id le)
          | Error _ -> None
        else None
    | None, _ -> None
  in
  match expr with
  | Signature s -> Signature (signature env id s)
  | Path { p_path; p_expansion } ->
      let p_path = module_type_path env p_path in
      Path { p_path; p_expansion = do_expn p_expansion (Some p_path) }
  | With { w_substitutions; w_expansion; w_expr } as unresolved -> (
      let cexpr = Component.Of_Lang.(u_module_type_expr (empty ()) w_expr) in
      match Tools.signature_of_u_module_type_expr env cexpr with
      | Ok sg ->
          With
            {
              w_substitutions = handle_fragments env id sg w_substitutions;
              w_expansion = do_expn w_expansion None;
              w_expr = u_module_type_expr env id w_expr;
            }
      | Error e ->
          Errors.report ~what:(`Module_type_U cexpr) ~tools_error:e `Expand;
          unresolved)
  | Functor (arg, res) ->
      let arg' = functor_argument env arg in
      let env = Env.add_functor_parameter arg env in
      let res' = module_type_expr env (Paths.Identifier.Mk.result id) res in
      Functor (arg', res')
  | TypeOf { t_desc = StructInclude p; t_expansion; t_original_path } ->
      TypeOf
        {
          t_desc = StructInclude (module_path env p);
          t_expansion = do_expn t_expansion None;
          t_original_path;
        }
  | TypeOf { t_desc = ModPath p; t_expansion; t_original_path } ->
      TypeOf
        {
          t_desc = ModPath (module_path env p);
          t_expansion = do_expn t_expansion None;
          t_original_path;
        }

and type_decl_representation :
    Env.t ->
    Id.Signature.t ->
    TypeDecl.Representation.t ->
    TypeDecl.Representation.t =
 fun env parent r ->
  let open TypeDecl.Representation in
  match r with
  | Variant cs -> Variant (List.map (type_decl_constructor env parent) cs)
  | Record fs -> Record (List.map (type_decl_field env parent) fs)
  | Extensible -> Extensible

and type_decl : Env.t -> Id.Signature.t -> TypeDecl.t -> TypeDecl.t =
 fun env parent t ->
  let open TypeDecl in
  let equation = type_decl_equation env parent t.equation in
  let doc = comment_docs env parent t.doc in
  let source_loc = source_loc env t.id t.source_loc in
  let hidden_path =
    match equation.Equation.manifest with
    | Some (Constr (`Resolved path, params))
      when Paths.Path.Resolved.(is_hidden (path :> t))
           || Paths.Path.Resolved.(identifier (path :> t))
              = Some (t.id :> Paths.Identifier.t) ->
        Some (path, params)
    | _ -> None
  in
  let representation =
    Opt.map
      (fun r ->
        let r' = type_decl_representation env parent r in
        warn_on_hidden_representation t.id r';
        r')
      t.representation
  in
  let default = { t with source_loc; equation; doc; representation } in
  match hidden_path with
  | Some (p, params) -> (
      let p' = Component.Of_Lang.(resolved_type_path (empty ()) p) in
      match Tools.lookup_type env p' with
      | Ok (`FType (_, t')) ->
          let equation =
            try
              Expand_tools.collapse_eqns default.equation
                (Lang_of.type_decl_equation (Lang_of.empty ())
                   (parent :> Id.FieldParent.t)
                   t'.equation)
                params
            with _ -> default.equation
          in
          { default with equation = type_decl_equation env parent equation }
      | Ok (`FClass _ | `FClassType _ | `FType_removed _ | `CoreType _)
      | Error _ ->
          default)
  | None -> default

and type_decl_equation env parent t =
  let open TypeDecl.Equation in
  let manifest = Opt.map (type_expression env parent []) t.manifest in
  let constraints =
    List.map
      (fun (tex1, tex2) ->
        (type_expression env parent [] tex1, type_expression env parent [] tex2))
      t.constraints
  in
  { t with manifest; constraints }

and type_decl_field env parent f =
  let open TypeDecl.Field in
  let doc = comment_docs env parent f.doc in
  { f with type_ = type_expression env parent [] f.type_; doc }

and type_decl_constructor_argument env parent c =
  let open TypeDecl.Constructor in
  match c with
  | Tuple ts -> Tuple (List.map (type_expression env parent []) ts)
  | Record fs -> Record (List.map (type_decl_field env parent) fs)

and type_decl_constructor env parent c =
  let open TypeDecl.Constructor in
  let doc = comment_docs env parent c.doc in
  let args = type_decl_constructor_argument env parent c.args in
  let res = Opt.map (type_expression env parent []) c.res in
  { c with doc; args; res }

and type_expression_polyvar env parent visited v =
  let open TypeExpr.Polymorphic_variant in
  let constructor c =
    let open Constructor in
    let doc = comment_docs env parent c.doc in
    {
      c with
      arguments = List.map (type_expression env parent visited) c.arguments;
      doc;
    }
  in
  let element = function
    | Type t ->
        Type
          (match type_expression env parent visited t with
          | Constr _ as x -> x
          | _ -> t)
        (* These have to remain Constrs *)
    | Constructor c -> Constructor (constructor c)
  in
  { v with elements = List.map element v.elements }

and type_expression_object env parent visited o =
  let open TypeExpr.Object in
  let method_ m =
    { m with type_ = type_expression env parent visited m.type_ }
  in
  let field = function
    | Method m -> Method (method_ m)
    | Inherit t -> Inherit (type_expression env parent visited t)
  in
  { o with fields = List.map field o.fields }

and type_expression_package env parent visited p =
  let open TypeExpr.Package in
  let substitution (frag, t) =
    let cfrag = Component.Of_Lang.(type_fragment (empty ()) frag) in
    let frag' =
      match cfrag with
      | `Resolved f -> `Resolved (Tools.reresolve_type_fragment env f)
      | _ -> cfrag
    in
    ( Lang_of.(Path.type_fragment (empty ()) frag'),
      type_expression env parent visited t )
  in
  {
    path = module_type_path env p.path;
    substitutions = List.map substitution p.substitutions;
  }

and type_expression : Env.t -> Id.Signature.t -> _ -> _ =
 fun env parent visited texpr ->
  let open TypeExpr in
  match texpr with
  | Var _ | Any -> texpr
  | Alias (t, str) -> Alias (type_expression env parent visited t, str)
  | Arrow (lbl, t1, t2) ->
      Arrow
        ( lbl,
          type_expression env parent visited t1,
          type_expression env parent visited t2 )
  | Tuple ts ->
      Tuple
        (List.map
           (fun (lbl, ty) -> (lbl, type_expression env parent visited ty))
           ts)
  | Constr (path', ts') -> (
      let path = type_path env path' in
      let ts = List.map (type_expression env parent visited) ts' in
      if not (Paths.Path.is_hidden (path :> Paths.Path.t)) then Constr (path, ts)
      else
        let cp = Component.Of_Lang.(type_path (empty ()) path') in
        match Tools.resolve_type env cp with
        | Ok (cp', `FType (_, t)) ->
            let cp' = Tools.reresolve_type env cp' in
            let p = Lang_of.(Path.resolved_type (empty ()) cp') in
            if List.mem p visited then raise Loop
            else if Cpath.is_resolved_type_hidden cp' then
              match t.Component.TypeDecl.equation with
              | { manifest = Some expr; params; _ } -> (
                  try
                    let map =
                      List.fold_left2
                        (fun acc param sub ->
                          match param.Lang.TypeDecl.desc with
                          | Lang.TypeDecl.Var x -> (x, sub) :: acc
                          | Any -> acc)
                        [] params ts
                    in
                    let t' =
                      Expand_tools.type_expr map
                        Lang_of.(
                          type_expr (empty ()) (parent :> Id.LabelParent.t) expr)
                    in
                    type_expression env parent (p :: visited) t'
                  with
                  | Loop -> Constr (`Resolved p, ts)
                  | e ->
                      Format.eprintf
                        "Caught unexpected exception when expanding type \
                         declaration (%s)\n\
                         %!"
                        (Printexc.to_string e);
                      Constr (`Resolved p, ts))
              | _ -> Constr (`Resolved p, ts)
            else Constr (`Resolved p, ts)
        | Ok (cp', (`FClass _ | `FClassType _ | `CoreType _)) ->
            let p = Lang_of.(Path.resolved_type (empty ()) cp') in
            Constr (`Resolved p, ts)
        | Ok (_cp, `FType_removed (_, x, _eq)) ->
            (* Type variables ? *)
            Lang_of.(type_expr (empty ()) (parent :> Id.LabelParent.t) x)
        | Error _ -> Constr (path', ts))
  | Polymorphic_variant v ->
      Polymorphic_variant (type_expression_polyvar env parent visited v)
  | Object o -> Object (type_expression_object env parent visited o)
  | Class (path', ts') -> (
      let path = class_type_path env path' in
      let ts = List.map (type_expression env parent visited) ts' in
      if not (Paths.Path.is_hidden (path :> Paths.Path.t)) then Class (path, ts)
      else
        let cp = Component.Of_Lang.(class_type_path (empty ()) path') in
        match Tools.resolve_class_type env cp with
        | Ok (cp', (`FClass _ | `FClassType _)) ->
            let cp' = Tools.reresolve_class_type env cp' in
            let p = Lang_of.(Path.resolved_class_type (empty ()) cp') in
            Class (`Resolved p, ts)
        | _ -> Class (path', ts))
  | Poly (strs, t) -> Poly (strs, type_expression env parent visited t)
  | Package p -> Package (type_expression_package env parent visited p)

let link ~filename x y =
  Lookup_failures.catch_failures ~filename (fun () ->
      if y.Lang.Compilation_unit.linked then y else unit x y)

let page env page =
  let () =
    List.iter
      (fun child ->
        match child with
        | Page.Page_child page -> (
            match Env.lookup_page_by_name page env with
            | Ok _ -> ()
            | Error `Not_found -> Errors.report ~what:(`Child_page page) `Lookup
            )
        | Page.Module_child mod_ -> (
            match
              Env.lookup_root_module
                (Odoc_model.Names.ModuleName.make_std mod_)
                env
            with
            | Some _ -> ()
            | None -> Errors.report ~what:(`Child_module mod_) `Lookup))
      page.Lang.Page.children
  in
  {
    page with
    Page.content = comment_docs env page.Page.name page.content;
    linked = true;
  }

let source_info env infos =
  let open Source_info in
  let jump_to v f_impl f_doc =
    let documentation =
      match v.documentation with Some p -> Some (f_doc p) | None -> None
    in
    let implementation =
      match v.implementation with
      | Some (Unresolved p) -> (
          match f_impl p with
          | Some x -> Some (Resolved x)
          | None -> v.implementation)
      | x -> x
    in
    { documentation; implementation }
  in
  List.map
    (fun (i, pos) ->
      let info =
        match i with
        | Value v ->
            Value
              (jump_to v (Shape_tools.lookup_value_path env) (value_path env))
        | Module v ->
            Module
              (jump_to v (Shape_tools.lookup_module_path env) (module_path env))
        | ModuleType v ->
            ModuleType
              (jump_to v
                 (Shape_tools.lookup_module_type_path env)
                 (module_type_path env))
        | Type v ->
            Type (jump_to v (Shape_tools.lookup_type_path env) (type_path env))
        | i -> i
      in
      (info, pos))
    infos

let impl env i =
  let open Implementation in
  { i with source_info = source_info env i.source_info; linked = true }

let resolve_page ~filename env p =
  Lookup_failures.catch_failures ~filename (fun () ->
      if p.Lang.Page.linked then p else page env p)

let resolve_impl ~filename env i =
  Lookup_failures.catch_failures ~filename (fun () ->
      if i.Lang.Implementation.linked then i else impl env i)
