(*---------------------------------------------------------------------------
  Copyright (c) 2025 Anil Madhavapeddy <anil@recoil.org>. All rights reserved.
  SPDX-License-Identifier: MIT
 ---------------------------------------------------------------------------*)

(* HTML5 DOM node types *)

type doctype_data = {
  name : string option;
  public_id : string option;
  system_id : string option;
}

(** Source location for nodes *)
type location = {
  line : int;
  column : int;
  end_line : int option;
  end_column : int option;
}

type quirks_mode = No_quirks | Quirks | Limited_quirks

type node = {
  mutable name : string;
  mutable namespace : string option;  (* None = html, Some "svg", Some "mathml" *)
  mutable attrs : (string * string) list;
  mutable children : node list;
  mutable parent : node option;
  mutable data : string;  (* For text, comment nodes *)
  mutable template_content : node option;  (* For <template> elements *)
  mutable doctype : doctype_data option;  (* For doctype nodes *)
  mutable location : location option;  (* Source location where node was parsed *)
}

(* Node name constants *)
let document_name = "#document"
let document_fragment_name = "#document-fragment"
let text_name = "#text"
let comment_name = "#comment"
let doctype_name = "!doctype"

(* Base node constructor - all nodes share this structure *)
let make_node ~name ?(namespace=None) ?(attrs=[]) ?(data="") ?template_content ?doctype ?location () = {
  name;
  namespace;
  attrs;
  children = [];
  parent = None;
  data;
  template_content;
  doctype;
  location;
}

(* Constructors *)
let create_element name ?(namespace=None) ?(attrs=[]) ?location () =
  make_node ~name ~namespace ~attrs ?location ()

let create_text ?location data =
  make_node ~name:text_name ~data ?location ()

let create_comment ?location data =
  make_node ~name:comment_name ~data ?location ()

let create_document () =
  make_node ~name:document_name ()

let create_document_fragment () =
  make_node ~name:document_fragment_name ()

let create_doctype ?name ?public_id ?system_id ?location () =
  make_node ~name:doctype_name ~doctype:{ name; public_id; system_id } ?location ()

let create_template ?(namespace=None) ?(attrs=[]) ?location () =
  let node = create_element "template" ~namespace ~attrs ?location () in
  node.template_content <- Some (create_document_fragment ());
  node

(* Predicates *)
let is_element node =
  not (List.mem node.name [text_name; comment_name; document_name; document_fragment_name; doctype_name])

let is_text node = node.name = text_name
let is_comment node = node.name = comment_name
let is_document node = node.name = document_name
let is_document_fragment node = node.name = document_fragment_name
let is_doctype node = node.name = doctype_name
let has_children node = node.children <> []

(* DOM manipulation *)
let append_child parent child =
  child.parent <- Some parent;
  parent.children <- parent.children @ [child]

let insert_before parent new_child ref_child =
  new_child.parent <- Some parent;
  let rec insert acc = function
    | [] -> List.rev acc @ [new_child]
    | x :: xs when x == ref_child -> List.rev acc @ [new_child; x] @ xs
    | x :: xs -> insert (x :: acc) xs
  in
  parent.children <- insert [] parent.children

let remove_child parent child =
  child.parent <- None;
  parent.children <- List.filter (fun c -> c != child) parent.children

(* Find the last text node before a reference point *)
let last_child_text parent =
  match List.rev parent.children with
  | last :: _ when is_text last -> Some last
  | _ -> None

let insert_text_at parent text before_node =
  match before_node with
  | None ->
    (* Append - merge with last child if it's text *)
    (match last_child_text parent with
     | Some txt -> txt.data <- txt.data ^ text
     | None -> append_child parent (create_text text))
  | Some ref ->
    (* Find last text node before ref_child *)
    let rec find_prev_text = function
      | [] | [_] -> None
      | prev :: curr :: _ when curr == ref && is_text prev -> Some prev
      | _ :: rest -> find_prev_text rest
    in
    match find_prev_text parent.children with
    | Some txt -> txt.data <- txt.data ^ text
    | None -> insert_before parent (create_text text) ref

(* Location helpers *)
let make_location ~line ~column ?end_line ?end_column () =
  { line; column; end_line; end_column }

let set_location node ~line ~column ?end_line ?end_column () =
  node.location <- Some { line; column; end_line; end_column }

let get_location node = node.location

(* Attribute helpers *)
let get_attr node name = List.assoc_opt name node.attrs

let set_attr node name value =
  node.attrs <- List.filter (fun (n, _) -> n <> name) node.attrs @ [(name, value)]

let has_attr node name = List.mem_assoc name node.attrs

(* Whitespace splitting for space-separated attribute values per HTML5 spec.
   Handles ASCII whitespace: space, tab, newline, carriage return, form feed *)
let split_on_whitespace s =
  let is_whitespace = function
    | ' ' | '\t' | '\n' | '\r' | '\x0c' -> true
    | _ -> false
  in
  let len = String.length s in
  let rec find_start acc i =
    if i >= len then List.rev acc
    else if is_whitespace s.[i] then find_start acc (i + 1)
    else find_end acc i (i + 1)
  and find_end acc start i =
    if i >= len then List.rev (String.sub s start (i - start) :: acc)
    else if is_whitespace s.[i] then find_start (String.sub s start (i - start) :: acc) (i + 1)
    else find_end acc start (i + 1)
  in
  find_start [] 0

(* Get space-separated attribute as list *)
let get_attr_list node name =
  match get_attr node name with
  | Some s -> split_on_whitespace s
  | None -> []

(* Common space-separated attribute accessors *)
let get_class_list node = get_attr_list node "class"
let get_rel_list node = List.map String.lowercase_ascii (get_attr_list node "rel")
let get_headers_list node = get_attr_list node "headers"
let get_itemref_list node = get_attr_list node "itemref"
let get_itemprop_list node = get_attr_list node "itemprop"
let get_itemtype_list node = get_attr_list node "itemtype"

(* Tree traversal *)
let rec descendants node =
  List.concat_map (fun n -> n :: descendants n) node.children

let ancestors node =
  let rec collect acc n =
    match n.parent with
    | None -> List.rev acc
    | Some p -> collect (p :: acc) p
  in
  collect [] node

let rec get_text_content node =
  if is_text node then node.data
  else String.concat "" (List.map get_text_content node.children)

(* Clone *)
let rec clone ?(deep=false) node =
  let new_node = make_node
    ~name:node.name
    ~namespace:node.namespace
    ~attrs:node.attrs
    ~data:node.data
    ?doctype:node.doctype
    ?location:node.location
    ()
  in
  if deep then begin
    new_node.children <- List.map (clone ~deep:true) node.children;
    List.iter (fun c -> c.parent <- Some new_node) new_node.children;
    Option.iter (fun tc ->
      new_node.template_content <- Some (clone ~deep:true tc)
    ) node.template_content
  end;
  new_node

(* Pretty printers *)
let pp_doctype_data fmt (d : doctype_data) =
  Format.fprintf fmt "<!DOCTYPE %s%s%s>"
    (Option.value ~default:"" d.name)
    (match d.public_id with Some p -> " PUBLIC \"" ^ p ^ "\"" | None -> "")
    (match d.system_id with Some s -> " \"" ^ s ^ "\"" | None -> "")

let pp_quirks_mode fmt = function
  | No_quirks -> Format.pp_print_string fmt "no-quirks"
  | Quirks -> Format.pp_print_string fmt "quirks"
  | Limited_quirks -> Format.pp_print_string fmt "limited-quirks"

let pp fmt node =
  if is_text node then
    Format.fprintf fmt "#text %S" node.data
  else if is_comment node then
    Format.fprintf fmt "<!-- %s -->" node.data
  else if is_document node then
    Format.pp_print_string fmt "#document"
  else if is_document_fragment node then
    Format.pp_print_string fmt "#document-fragment"
  else if is_doctype node then
    (match node.doctype with
     | Some d -> pp_doctype_data fmt d
     | None -> Format.pp_print_string fmt "<!DOCTYPE>")
  else begin
    Format.fprintf fmt "<%s" node.name;
    List.iter (fun (k, v) -> Format.fprintf fmt " %s=%S" k v) node.attrs;
    Format.pp_print_char fmt '>'
  end