(*
 * SPDX-FileCopyrightText: 2024 The Forester Project Contributors AND The RedPRL Development Team
 * SPDX-FileCopyrightText: 2026 The Forester Project Contributors
 *
 * SPDX-License-Identifier: GPL-3.0-or-later
 * SPDX-License-Identifier: GPL-3.0-or-later OR Apache-2.0 WITH LLVM-exception
 *)

open Forester_core
open Forester_compiler

open struct
  module L = Lsp.Types
end

open Forest.Syntax

let addr_definition ~forest ~base ~position nodes =
  let@ {value = str; _} = Option.bind @@ Analysis.addr_at ~position nodes in
  let target_uri = URI.named_uri ~base str in
  let forest = Driver.ensure_loaded ~forest target_uri in
  Lsp_state.modify (fun lsp_state -> {lsp_state with forest});
  let lsp_uri = Option.bind forest.={target_uri} Tree.lsp_uri in
  let range =
    L.Range.create ~start:{character = 1; line = 0}
      ~end_:{character = 1; line = 0}
  in
  Option.map (fun uri -> L.Location.{uri; range}) lsp_uri

let path_definition ~forest ~position nodes =
  let@ node = Option.bind @@ Analysis.node_at_code ~position nodes in
  let@ paths, _ = Option.bind @@ Analysis.paths node in
  let visible = Analysis.get_visible ~forest ~position nodes in
  let@ path = List.find_map @~ paths in
  match Trie.find_singleton path visible with
  | Some (_, Some range) -> Lsp_shims.location_of_range range
  | _ -> None

let compute (params : L.DefinitionParams.t) =
  let Lsp_state.{forest; _} = Lsp_state.get () in
  let base = Forest.base forest in
  let uri = URI.of_lsp_uri ~base params.textDocument.uri in
  let@ tree = Option.bind forest.={uri} in
  let@ nodes = Option.bind Tree.(Option.map Parsed.nodes @@ code tree) in
  match addr_definition ~forest ~base ~position:params.position nodes with
  | Some location -> Some (`Location [location])
  | None ->
    let@ location =
      Option.map @~ path_definition ~forest ~position:params.position nodes
    in
    `Location [location]