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

open Base

open struct
  module T = Types
end

type 'a _object = {self: string option; methods: (string * 'a) list}
[@@deriving show]

type 'a patch = {
  obj: 'a;
  self: string option;
  super: string option;
  methods: (string * 'a) list;
}
[@@deriving show]

type node =
  | Text of string
  | Verbatim of string
  | Group of delim * t
  | Math of math_mode * t
  | Ident of Trie.path
  | Hash_ident of string
  | Xml_ident of string option * string
  | Subtree of string Range.located option * t
  | Let of Trie.path * string binding list * t
  | Open of Trie.path
  | Scope of t
  | Put of Trie.path Range.located * t
  | Default of Trie.path Range.located * t
  | Get of Trie.path Range.located
  | Fun of string binding list * t
  | Object of t _object
  | Patch of t patch
  | Call of t * string
  | Import of visibility * string Range.located
  | Def of Trie.path * string binding list * t
  | Decl_xmlns of string * string
  | Alloc of Trie.path
  | Namespace of Trie.path * t
  | Dx_sequent of t * t list
  | Dx_query of string * t list * t list
  | Dx_prop of t * t list
  | Dx_var of string
  | Dx_const_content of t
  | Dx_const_uri of t
  | Comment of string
  | Error of string
[@@deriving show]

and t = node Range.located list [@@deriving show]

type tree = {source_path: string option; uri: URI.t option; code: t}
[@@deriving show]

let import_private x = Import (Private, x)
let import_public x = Import (Public, x)
let inline_math e = Math (Inline, e)
let display_math e = Math (Display, e)
let parens e = Group (Parens, e)
let squares e = Group (Squares, e)
let braces e = Group (Braces, e)

let map f node =
  match node with
  | Group (d, t) -> Group (d, f t)
  | Math (m, t) -> Math (m, f t)
  | Namespace (p, t) -> Namespace (p, f t)
  | Dx_const_content t -> Dx_const_content (f t)
  | Dx_const_uri t -> Dx_const_uri (f t)
  | Def (p, b, t) -> Def (p, b, f t)
  | Dx_sequent (t, ts) -> Dx_sequent (f t, List.map f ts)
  | Dx_query (s, ts, rs) -> Dx_query (s, List.map f ts, List.map f @@ rs)
  | Dx_prop (t, ts) -> Dx_prop (t, ts)
  | Subtree (s, t) -> Subtree (s, f t)
  | Let (p, b, t) -> Let (p, b, f t)
  | Default (p, t) -> Default (p, f t)
  | Scope t -> Scope (f t)
  | Put (p, t) -> Put (p, f t)
  | Fun (b, t) -> Fun (b, f t)
  | Call (t, s) -> Call (f t, s)
  | Object {self; methods} ->
    Object {self; methods = List.map (fun (s, t) -> (s, f t)) methods}
  | Patch {obj; self; super; methods} ->
    Patch
      {
        obj = f obj;
        self;
        super;
        methods = List.map (fun (s, t) -> (s, f t)) methods;
      }
  | Text _ | Verbatim _ | Ident _ | Hash_ident _
  | Xml_ident (_, _)
  | Open _ | Get _
  | Import (_, _)
  | Decl_xmlns (_, _)
  | Alloc _ | Dx_var _ | Comment _ | Error _ ->
    node

let children (node : node Range.located) =
  match node.value with
  | Math (_, t)
  | Group (_, t)
  | Let (_, _, t)
  | Scope t
  | Put (_, t)
  | Fun (_, t)
  | Default (_, t)
  | Def (_, _, t)
  | Namespace (_, t)
  | Dx_const_uri t
  | Dx_const_content t
  | Call (t, _)
  | Subtree (_, t) ->
    t
  | Dx_prop (_, t) | Dx_query (_, _, t) | Dx_sequent (_, t) -> List.concat t
  | Object {methods; _} -> methods |> List.map snd |> List.concat
  | Patch {obj; methods; _} ->
    let methods = methods |> List.map snd |> List.concat in
    List.append obj methods
  | Text _ | Verbatim _ | Ident _ | Hash_ident _
  | Xml_ident (_, _)
  | Open _ | Get _
  | Import (_, _)
  | Decl_xmlns (_, _)
  | Alloc _ | Dx_var _ | Comment _ | Error _ ->
    []

let rec equal (a : t) (b : t) = equal_list (equal_located equal_node) a b

and equal_list : 'a. ('a -> 'a -> bool) -> 'a list -> 'a list -> bool =
 fun eq x y ->
  match (x, y) with
  | [], [] -> true
  | a :: x, b :: y -> eq a b && equal_list eq x y
  | _ -> false

and equal_trie_path = List.equal String.equal

and equal_located :
    'a. ('a -> 'a -> bool) -> 'a Range.located -> 'a Range.located -> bool =
 fun eq a b -> eq a.value b.value

and equal_binding eq (v1, a) (v2, b) =
  match (v1, v2) with
  | Base.Strict, Base.Strict | Base.Lazy, Base.Lazy -> eq a b
  | _ -> false

and equal_object eq o1 o2 = false

and equal_patch eq p1 p2 = false

and equal_visibility v1 v2 =
  match (v1, v2) with
  | Base.Public, Base.Public | Base.Private, Base.Private -> true
  | _ -> false

and equal_node (a : node) (b : node) =
  match (a, b) with
  | Text lhs0, Text rhs0 -> String.equal lhs0 rhs0
  | Verbatim lhs0, Verbatim rhs0 -> String.equal lhs0 rhs0
  | Group (lhs0, lhs1), Group (rhs0, rhs1) -> lhs0 = rhs0 && equal lhs1 rhs1
  | Math (lhs0, lhs1), Math (rhs0, rhs1) -> lhs0 = rhs0 && equal lhs1 rhs1
  | Ident lhs0, Ident rhs0 -> equal_trie_path lhs0 rhs0
  | Hash_ident lhs0, Hash_ident rhs0 -> String.equal lhs0 rhs0
  | Xml_ident (lhs0, lhs1), Xml_ident (rhs0, rhs1) ->
    (fun x y ->
      match (x, y) with
      | None, None -> true
      | Some a, Some b -> String.equal a b
      | _ -> false)
      lhs0 rhs0
    && String.equal lhs1 rhs1
  | Subtree (lhs0, lhs1), Subtree (rhs0, rhs1) ->
    (fun x y ->
      match (x, y) with
      | None, None -> true
      | Some a, Some b -> (equal_located String.equal) a b
      | _ -> false)
      lhs0 rhs0
    && equal lhs1 rhs1
  | Let (lhs0, lhs1, lhs2), Let (rhs0, rhs1, rhs2) ->
    equal_trie_path lhs0 rhs0
    && equal_list (equal_binding String.equal) lhs1 rhs1
    && equal lhs2 rhs2
  | Open lhs0, Open rhs0 -> equal_trie_path lhs0 rhs0
  | Scope lhs0, Scope rhs0 -> equal lhs0 rhs0
  | Put (lhs0, lhs1), Put (rhs0, rhs1) ->
    (equal_located equal_trie_path) lhs0 rhs0 && equal lhs1 rhs1
  | Default (lhs0, lhs1), Default (rhs0, rhs1) ->
    (equal_located equal_trie_path) lhs0 rhs0 && equal lhs1 rhs1
  | Get lhs0, Get rhs0 -> (equal_located equal_trie_path) lhs0 rhs0
  | Fun (lhs0, lhs1), Fun (rhs0, rhs1) ->
    equal_list (equal_binding String.equal) lhs0 rhs0 && equal lhs1 rhs1
  | Object lhs0, Object rhs0 -> (equal_object equal) lhs0 rhs0
  | Patch lhs0, Patch rhs0 -> (equal_patch equal) lhs0 rhs0
  | Call (lhs0, lhs1), Call (rhs0, rhs1) ->
    equal lhs0 rhs0 && String.equal lhs1 rhs1
  | Import (lhs0, lhs1), Import (rhs0, rhs1) ->
    equal_visibility lhs0 rhs0 && (equal_located String.equal) lhs1 rhs1
  | Def (lhs0, lhs1, lhs2), Def (rhs0, rhs1, rhs2) ->
    equal_trie_path lhs0 rhs0
    && equal_list (equal_binding String.equal) lhs1 rhs1
    && equal lhs2 rhs2
  | Decl_xmlns (lhs0, lhs1), Decl_xmlns (rhs0, rhs1) ->
    String.equal lhs0 rhs0 && String.equal lhs1 rhs1
  | Alloc lhs0, Alloc rhs0 -> equal_trie_path lhs0 rhs0
  | Namespace (lhs0, lhs1), Namespace (rhs0, rhs1) ->
    equal_trie_path lhs0 rhs0 && equal lhs1 rhs1
  | Dx_sequent (lhs0, lhs1), Dx_sequent (rhs0, rhs1) ->
    equal lhs0 rhs0 && equal_list equal lhs1 rhs1
  | Dx_query (lhs0, lhs1, lhs2), Dx_query (rhs0, rhs1, rhs2) ->
    String.equal lhs0 rhs0 && equal_list equal lhs1 rhs1
    && equal_list equal lhs2 rhs2
  | Dx_prop (lhs0, lhs1), Dx_prop (rhs0, rhs1) ->
    equal lhs0 rhs0 && equal_list equal lhs1 rhs1
  | Dx_var lhs0, Dx_var rhs0 -> String.equal lhs0 rhs0
  | Dx_const_content lhs0, Dx_const_content rhs0 -> equal lhs0 rhs0
  | Dx_const_uri lhs0, Dx_const_uri rhs0 -> equal lhs0 rhs0
  | Comment lhs0, Comment rhs0 -> String.equal lhs0 rhs0
  | Error lhs0, Error rhs0 -> String.equal lhs0 rhs0
  | _ -> false