(*
 * SPDX-FileCopyrightText: 2026 The Forester Project Contributors
 *
 * SPDX-License-Identifier: GPL-3.0-or-later
 *)

open Forester_core
open Forester_parser

open struct
  module T = Types
  module G = Grace
  module D = G.Diagnostic

  type message = D.Message.t

  let empty_label ~range = D.Label.createf ~range ~priority:Primary ""

  let fold_range ~range msg =
    Option.fold range
      ~some:(fun range -> [D.Label.create ~priority:Primary ~range msg])
      ~none:[]

  let primary_label ~range fmt =
    let@ s = Format.kasprintf @~ fmt in
    [D.Label.create ~priority:Primary ~range (D.Message.create s)]

  let message = D.Message.createf

  let error ?labels ?notes fmt = D.createf Error ?labels ?notes fmt
  let warning ?labels ?notes fmt = D.createf Warning ?labels ?notes fmt
end

open G
open D

type latex_error = {
  range: Grace.Range.t;
  msg: string;
  binary_missing: string option;
}

type unlinked_attribution_warning = T.(content attribution)
type unreferenced_asset = {path: string; uri: URI.t}
type broken_link = {link: T.(content link); suggestion: URI.t option}
type broken_transclusion = {t: T.transclusion; suggestion: URI.t option}
type foreign_blob_error = {path: string; msg: string}

type t =
  [ Eval_error.t
  | Expand_error.t
  | Config_error.t
  | `Parse_error of Parse.error
  | `Io_error of Eio_util.IO_error.t
  | `LaTeX_error of latex_error
  | `Duplicate_tree of URI.t * Range.t list * string list
  | `Unlinked_attribution_warning of unlinked_attribution_warning
  | `Broken_link of broken_link
  | `Broken_transclusion of broken_transclusion
  | `Cant_eval_anonymous_tree
  | `Failed_to_load_foreign_blob of foreign_blob_error
  | `Failed_to_parse_foreign_blob of foreign_blob_error
  | `Failed_to_add_edge of Vertex.t * Vertex.t
  | `Failed_to_add_vertex of Range.t * URI.t
  | `Foreign_import of Range.t * URI.t
  | `Unreferenced_asset of unreferenced_asset ]

type error = t

let tex_range {range; _} = range
let latex_error ?binary_missing ~range msg = {range; msg; binary_missing}
let latex_error_binary_missing {binary_missing; _} = binary_missing
let of_tex_error e = `LaTeX_error e

let unlinked_attribution_warning attribution =
  `Unlinked_attribution_warning attribution
let broken_transclusion ?suggestion t = `Broken_transclusion {t; suggestion}
let broken_link ?suggestion link = `Broken_link {link; suggestion}
let duplicate_tree ~uri ~ranges ~paths = `Duplicate_tree (uri, ranges, paths)
let failed_to_add_edge v w = `Failed_to_add_edge (v, w)
let failed_to_add_vertex ~range v = `Failed_to_add_vertex (range, v)
let foreign_import ~range uri = `Foreign_import (range, uri)
let unreferenced_asset ~path uri = `Unreferenced_asset {path; uri}
let cant_eval_anonymous_tree = `Cant_eval_anonymous_tree
let failed_to_load_foreign_blob ~msg ~path =
  `Failed_to_load_foreign_blob {msg; path}
let failed_to_parse_foreign_blob ~msg ~path =
  `Failed_to_parse_foreign_blob {msg; path}

module Collect_errors = Algaeff.Sequencer.Make (struct
  (* hack to avoid cyclic type error *)
  type error = t
  type t = error
end)

let collect = Collect_errors.run
let yield = Collect_errors.yield

let parse_error err : t = `Parse_error err
let io_error exn : t = `Io_error exn
let eval_error (e : Eval_error.t) : t = (e :> t)
let expand_error (e : Expand_error.t) : t = (e :> t)
let config_error (e : Config_error.t) : t = (e :> t)

let yield_eval_error = eval_error >>> Collect_errors.yield
let yield_expand_error = expand_error >>> Collect_errors.yield

let did_you_mean_note (type a) (pp : Format.formatter -> a -> unit) :
    a option -> message list = function
  | None -> []
  | Some x -> [message "Did you mean %a?" pp x]

let foreign_blob_hint =
  [
    message
      "Make sure the blob was generated with the version of forester that you \
       are using.";
  ]

let render_type_error ({range; got; expected} : Eval_error.type_error) =
  let show_value : Value.t -> string = function
    | Content _ -> "content"
    | Clo (_, _, _) -> "a function"
    | Dx_prop _ -> "a datalog proposition"
    | Dx_sequent _ -> "a datalog sequent"
    | Dx_query _ -> "a datalog query"
    | Dx_var _ -> "a datalog variable"
    | Dx_const _ -> "a datalog constant"
    | Sym _ -> "a symbol"
    | Obj _ -> "an object"
  in
  let labels =
    match got with
    | Some (Clo _) ->
      primary_label ~range
        "This is a function. Did you forget to provide some arguments?"
    | Some (Obj _) ->
      primary_label ~range
        "This is an object. Did you forget to bind this object to an \
         identifier?"
    | Some got -> primary_label ~range "this is %s" (show_value got)
    | None -> [empty_label ~range]
  in
  let notes =
    match expected with
    | [] -> []
    | x :: [] -> [message "expected %s" @@ Eval_error.show_expectation x]
    | xs ->
      [
        message "expected one of @[<v2>[ %a ]@]"
          (Format.pp_print_list Eval_error.pp_expected_value)
          xs;
      ]
  in
  error ~labels ~notes "mismatched type"

let render_eval_error : Eval_error.t -> _ Diagnostic.t =
  let asset_error labels = error ~labels "asset error" in
  function
  | `Type_error err -> render_type_error err
  | `Invalid_URI {range; uri} ->
    let labels = primary_label ~range "%s" uri in
    error ~labels "Invalid URI"
  | `Invalid_date {range; date} ->
    let labels =
      primary_label ~range "`%s` is not a date in the form of YYYY-MM-DD." date
    in
    error ~labels "Failed to parse date"
  | `Unbound_variable {range; name; suggestion} ->
    let notes = did_you_mean_note Format.pp_print_string suggestion in
    let labels = primary_label ~range "no binding named %s is in scope" name in
    error ~labels ~notes "unbound variable"
  | `No_current_uri range -> error ~labels:[empty_label ~range] "no current uri"
  | `Missing_arguments {range; n} ->
    let labels =
      primary_label ~range "this function expected %d more argument%s" n
        (if n = 1 then "" else "s")
    in
    error ~labels "missing arguments"
  | `Unbound_method {range; method_name} ->
    let labels = primary_label ~range "no method named %s" method_name in
    error ~labels "unbound method"
  | `Unbound_fluid_symbol {range; sym} ->
    let labels =
      primary_label ~range "no fluid binding for %s" (Symbol.show sym)
    in
    error ~labels "unbound fluid symbol"
  | `Unresolved_identifier {range; suggestions; path} ->
    let labels = [empty_label ~range] in
    error ~labels ~notes:suggestions "unresolved identifier %a" Trie.pp_path
      path
  | `Not_found {range; path} ->
    asset_error (primary_label ~range "no such file: %s" path)
  | `No_content_address range ->
    asset_error (primary_label ~range "this asset has no content address")
  | `Failed_to_hash range ->
    asset_error (primary_label ~range "failed to hash this asset's content")

let render_parse_error Parse.{range; msg; lexeme} =
  let labels =
    match lexeme with
    | "import" ->
      primary_label ~range "imports may only appear on the top level of a file"
    | _ -> [empty_label ~range]
  in
  error ~labels "%s" msg

let render_io_error (e : Eio_util.IO_error.t) =
  let msg =
    match e with
    | Path_is_not_a_directory d -> d ^ " is not a directory"
    | Path_is_not_a_file f -> f ^ " is not a file"
    | Failed_to_create_directory d -> "failed to create directory " ^ d
    | Failed_to_create_file f -> "failed to create file " ^ f
    | File_already_exists f -> "file " ^ f ^ " already exists"
    | IO_error s -> s
  in
  error "%s" msg

let render_config_error : Config_error.t -> _ Diagnostic.t =
  let pp_path =
    Format.(
      pp_print_list ~pp_sep:(fun out _ -> fprintf out ".") pp_print_string)
  in
  function
  | `Config_parse_error {range; msg} ->
    error ~labels:[empty_label ~range] "%s" msg
  | `Using_default {path; default} ->
    warning "using default %a for option %a" pp_path default pp_path path
  | `Unknown_option {path} -> warning "unknown option %a" pp_path path
  | `Invalid_url {url} -> error "invalid url %s" url
  | `No_path_in_foreign_table -> error "no path in foreign table"
  | `Missing_dir {dir} -> error "directory %s does not exist" dir
  | `Missing_file {file} -> error "file %s does not exist" file
  | `No_forest_table -> error "no top-level [forest] table"

let render_latex_error ({range; msg; _} : latex_error) =
  error ~labels:[empty_label ~range] "%s" msg

let render_broken_transclusion {t = {range; href; _}; suggestion} =
  let notes = did_you_mean_note URI.pp suggestion in
  let labels = fold_range ~range (message "This tree was not found") in
  warning ~labels ~notes "Broken transclusion %a" URI.pp href

let render_broken_link {link = {range; href; _}; suggestion} =
  let notes = did_you_mean_note URI.pp suggestion in
  let labels = fold_range ~range (message "This tree was not found") in
  warning ~labels ~notes "Broken link %a" URI.pp href

let render_unlinked_attribution_warning
    ({role; range; vertex} : unlinked_attribution_warning) =
  let corrected_attribution_code =
    match role with
    | Author -> "\\author/literal"
    | Contributor -> "\\contributor/literal"
  in
  let arg =
    match vertex with
    | T.Uri_vertex uri ->
      Option.value ~default:(Format.asprintf "%a" URI.pp uri) @@ URI.name uri
    | T.Content_vertex _ -> "..."
  in
  let suggestion = Format.asprintf "%s{%s}" corrected_attribution_code arg in
  let notes =
    [
      message
        "If there is a tree trees/%s.tree, you may use the shorthand %s \
         instead of a full URI."
        arg arg;
    ]
  in
  let labels =
    fold_range ~range
    @@ message
         "Expected valid URI in attribution. Use `%s` instead if you intend an \
          unlinked attribution."
         suggestion
  in
  warning ~labels ~notes "Did you mean %s?" suggestion

let render_unreferenced_asset ({path; _} : unreferenced_asset) =
  let notes =
    [
      message
        "This file is not referenced by any tree; consider removing it or \
         referencing it with \\route-asset.";
    ]
  in
  warning ~notes "unreferenced asset: %s" path

let render_foreign_blob_error ~verb ({path; msg} : foreign_blob_error) =
  error ~notes:foreign_blob_hint "Failed to %s foreign blob %s: %s" verb path
    msg

let render ~(config : Config.t) (err : error) : error Diagnostic.t =
  let diag =
    match err with
    | `Cant_eval_anonymous_tree -> failwith "todo: can't eval"
    | #Eval_error.t as e -> render_eval_error e
    | #Expand_error.t as e -> Expand_error.render e
    | #Config_error.t as e -> render_config_error e
    | `Parse_error error -> render_parse_error error
    | `Io_error error -> render_io_error error
    | `LaTeX_error error -> render_latex_error error
    | `Duplicate_tree (uri, ranges, paths) ->
      let labels =
        match ranges with
        | [] -> []
        | first :: rest ->
          primary_label ~range:first "duplicate subtree address"
          @ List.map
              (fun range ->
                Label.createf ~range ~priority:Secondary "also declared here")
              rest
      in
      let notes =
        match paths with
        | [] -> []
        | paths -> [message "also registered at: %s" (String.concat ", " paths)]
      in
      error ~labels ~notes "duplicate tree %a" URI.pp uri
    | `Broken_link link -> render_broken_link link
    | `Broken_transclusion t -> render_broken_transclusion t
    | `Unlinked_attribution_warning w -> render_unlinked_attribution_warning w
    | `Failed_to_load_foreign_blob s -> render_foreign_blob_error ~verb:"load" s
    | `Failed_to_parse_foreign_blob s ->
      render_foreign_blob_error ~verb:"parse" s
    | `Failed_to_add_edge (_, _) -> failwith "failed to add edge"
    | `Failed_to_add_vertex (range, uri) ->
      let name = Option.get @@ URI.name uri in
      let labels = primary_label ~range "No file %s.tree found." name in
      let notes =
        [
          message "checked the following directories: %a"
            Format.(
              pp_print_list
                ~pp_sep:(fun out _ -> fprintf out ", ")
                pp_print_string)
            config.trees;
        ]
      in
      error ~labels ~notes ""
    | `Foreign_import (range, uri) ->
      let host = Option.value ~default:"?" (URI.host uri) in
      let labels =
        primary_label ~range "%a belongs to %s, a foreign forest" URI.pp uri
          host
      in
      let notes = [message "foreign trees cannot be imported"] in
      error ~labels ~notes ""
    | `Unreferenced_asset asset -> render_unreferenced_asset asset
  in
  {diag with code = Some err}

let range ~(config : Config.t) error : Range.t option =
  match (render ~config error).labels with
  | label :: _ -> Some label.range
  | [] -> None

let source_name (source : Grace.Source.t) : string option =
  match source with
  | `File filename -> Some filename
  | `String {name; _} -> name
  | `Reader _ -> None

let compare ~(config : Config.t) e1 e2 =
  match (range ~config e1, range ~config e2) with
  | None, None -> 0
  | None, Some _ -> -1
  | Some _, None -> 1
  | Some r1, Some r2 ->
    compare (source_name (Range.source r1)) (source_name (Range.source r2))

let code_to_string : error -> string = function
  | `Type_error _ -> "type_mismatch"
  | `Invalid_URI _ -> "invalid_uri"
  | `Invalid_date _ -> "invalid_date"
  | `Unbound_variable _ -> "unbound_variable"
  | `No_current_uri _ -> "no_current_uri"
  | `Missing_arguments _ -> "missing_arguments"
  | `Unbound_method _ -> "unbound_method"
  | `Unbound_fluid_symbol _ -> "unbound_fluid_symbol"
  | `Unresolved_identifier _ | `Expand_unresolved_identifier _ ->
    "unresolved_identifier"
  | `Not_found _ | `No_content_address _ | `Failed_to_hash _ -> "invalid_asset"
  | `Import_not_found _ | `Failed_to_add_vertex _ -> "import_not_found"
  | `Unresolved_xmlns _ -> "unresolved_xmlns"
  | `Config_parse_error _ | `No_path_in_foreign_table | `No_forest_table ->
    "invalid_configuration"
  | `Using_default _ -> "default_configuration_used"
  | `Unknown_option _ -> "unknown_configuration_option"
  | `Invalid_url _ -> "invalid_url"
  | `Missing_dir _ -> "missing_directory"
  | `Missing_file _ -> "missing_file"
  | `Parse_error _ -> "invalid_syntax"
  | `Io_error (Path_is_not_a_directory _) -> "not_a_directory"
  | `Io_error (Path_is_not_a_file _) -> "not_a_file"
  | `Io_error (Failed_to_create_directory _) -> "failed_to_create_directory"
  | `Io_error (Failed_to_create_file _) -> "failed_to_create_file"
  | `Io_error (File_already_exists _) -> "file_already_exists"
  | `Io_error (IO_error _) -> "io_failure"
  | `LaTeX_error _ -> "latex_failure"
  | `Duplicate_tree _ -> "duplicate_tree"
  | `Unlinked_attribution_warning _ -> "unlinked_attribution"
  | `Broken_link _ -> "broken_link"
  | `Broken_transclusion _ -> "broken_transclusion"
  | `Cant_eval_anonymous_tree | `Failed_to_add_edge _ -> "internal_failure"
  | `Failed_to_load_foreign_blob _ | `Failed_to_parse_foreign_blob _ ->
    "foreign_forest_failure"
  | `Foreign_import _ -> "foreign_import"
  | `Unreferenced_asset _ -> "unreferenced_asset"

let render_lsp_related_info (label : Diagnostic.Label.t) :
    Lsp.Types.DiagnosticRelatedInformation.t option =
  let@ uri =
    Option.map @~ Lsp_shims.lsp_uri_of_source (Range.source label.range)
  in
  let range = Lsp_shims.lsp_range_of_range label.range in
  let location = Lsp.Types.Location.create ~uri ~range in
  let text = Format.asprintf "%a" Diagnostic.Message.pp label.message in
  Lsp.Types.DiagnosticRelatedInformation.create ~location ~message:text

let render_lsp_diagnostic uri diag =
  let belongs_to_this_doc (label : Diagnostic.Label.t) =
    match Lsp_shims.lsp_uri_of_source (Range.source label.range) with
    | Some label_uri -> Lsp.Uri.equal label_uri uri
    | None -> false
  in
  let range =
    match (List.find_opt belongs_to_this_doc diag.labels, diag.labels) with
    | Some label, _ | None, label :: _ ->
      Lsp_shims.lsp_range_of_range label.range
    | None, [] ->
      let start = Lsp.Types.Position.create ~line:0 ~character:0 in
      Lsp.Types.Range.create ~start ~end_:start
  in
  let severity = Lsp_shims.Diagnostic.lsp_severity_of_severity diag.severity in
  let code =
    let@ code = Option.map @~ diag.code in
    `String (code_to_string code)
  in
  let text = Format.asprintf "%a" Diagnostic.Message.pp diag.message in
  let relatedInformation =
    List.filter_map render_lsp_related_info diag.labels
  in
  Lsp.Types.Diagnostic.create ~range ~severity ?code ~source:"forester"
    ~message:(`String text) ~relatedInformation ()

let print_diagnostic =
  let config =
    Grace_ansi_renderer.Config.{default with num_contextual_lines = 3}
  in
  Format.eprintf "%a@."
  @@ Grace_ansi_renderer.pp_diagnostic ~config ~code_to_string

let print ~config error = print_diagnostic (render ~config error)

let is_fatal : error -> bool = function
  | `Broken_link _ | `Broken_transclusion _ | `Unreferenced_asset _ -> false
  | _ -> true

let any_fatal = List.fold_left (fun acc x -> acc || is_fatal x) false

let print_config_error error =
  print_diagnostic
    {(render_config_error error) with code = Some (config_error error)}